Skip to content

validate_records drops list values when the first row omits a nullable list column #4

Description

@sbatchelder

validate_records drops list values when the first row omits a nullable list column

Summary

schema.validate_records(records: list[dict], schema) can silently drop values from nullable list columns when the first input row omits that column.

The behavior is order-dependent. If the first row has no value for a nullable list field, validate_records first constructs a PyArrow table without using the declared schema. Because pa.Table.from_pylist(...) infers the table shape from the input rows, the nullable list column can be omitted or materialized incorrectly before the later cast to the target schema. When the missing nullable column is then added back, it is filled entirely with nulls, so list values from later rows are lost.

This makes mixed record batches unsafe when some row types leave a nullable vector/list column empty and other rows in the same batch populate it.

Reproduce

import pyarrow as pa

from amplify_db_utils.schema import validate_records

schema = pa.schema([
    pa.field("id", pa.string(), nullable=False),
    pa.field("vec", pa.list_(pa.float32()), nullable=True),
])

rows = [
    {"id": "a"},                      # first row: "vec" absent
    {"id": "b", "vec": [1.0, 2.0]},   # later row: "vec" populated
]

table = validate_records(rows, schema)

print(table.column("vec").to_pylist())

# expected:
# [None, [1.0, 2.0]]
#
# actual:
# [None, None]

The same records in the opposite order preserve the list value:

rows = [
    {"id": "b", "vec": [1.0, 2.0]},
    {"id": "a"},
]

table = validate_records(rows, schema)

print(table.column("vec").to_pylist())

# expected:
# [[1.0, 2.0], None]
#
# actual:
# [[1.0, 2.0], None]

Root cause

For list[dict] input, validate_records currently does this:

processed = _preprocess_list_records(records, schema)
table = pa.Table.from_pylist(processed)

The declared target schema is not passed into from_pylist. As a result, PyArrow infers the initial table from the records themselves before validate_records has enforced the registered schema.

Later, validate_records fills missing nullable columns like this:

table = table.append_column(
    field,
    pa.array([None] * len(table), type=field.type),
)

At that point, any values that were present only in later rows have already been discarded from the inferred table. The final table.cast(schema) cannot recover those values.

Suggested fix

Build the Arrow table using the declared schema from the start, instead of inferring the table first and casting afterward.

One direct fix is to pass the target schema to pa.Table.from_pylist:

processed = _preprocess_list_records(records, schema)

try:
    table = pa.Table.from_pylist(processed, schema=schema)
except Exception as e:
    raise ValueError(f"Failed to convert records to Arrow table: {e}") from e

Because from_pylist(..., schema=schema) produces a table with exactly the requested schema, this should also preserve the existing behavior of dropping unknown columns.

If missing nullable fields need to be normalized before conversion, another safe option is to null-fill only the schema fields first and then build arrays using each declared field type:

columns = {}

for field in schema:
    values = [record.get(field.name) for record in processed]
    columns[field.name] = pa.array(values, type=field.type)

table = pa.table(columns, schema=schema)

Either approach makes the registered schema determine the column types and names, not the first row of the batch.

Test plan

Add coverage for nullable list fields where the first row omits the list column and a later row provides it.

def test_validate_records_preserves_list_values_when_first_row_missing_list_column():
    schema = pa.schema([
        pa.field("id", pa.string(), nullable=False),
        pa.field("vec", pa.list_(pa.float32()), nullable=True),
    ])

    rows = [
        {"id": "a"},
        {"id": "b", "vec": [1.0, 2.0]},
    ]

    table = validate_records(rows, schema)

    assert table.column("vec").to_pylist() == [None, [1.0, 2.0]]

Also add the reverse ordering to guard against order-dependent behavior:

def test_validate_records_preserves_list_values_when_first_row_has_list_column():
    schema = pa.schema([
        pa.field("id", pa.string(), nullable=False),
        pa.field("vec", pa.list_(pa.float32()), nullable=True),
    ])

    rows = [
        {"id": "b", "vec": [1.0, 2.0]},
        {"id": "a"},
    ]

    table = validate_records(rows, schema)

    assert table.column("vec").to_pylist() == [[1.0, 2.0], None]

Regression expectations:

  • Nullable list values are preserved regardless of row order.
  • Unknown columns are still dropped.
  • Missing nullable columns are still filled with nulls.
  • Missing non-nullable columns still raise ValueError.
  • Existing scalar-column behavior is unchanged.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions