diff --git a/src/amplify_db_utils/schema.py b/src/amplify_db_utils/schema.py index b56bbbc..095c166 100644 --- a/src/amplify_db_utils/schema.py +++ b/src/amplify_db_utils/schema.py @@ -128,8 +128,10 @@ def validate_records( table = records elif isinstance(records, list): processed = _preprocess_list_records(records, schema) + present_names = {key for record in records for key in record} + build_schema = pa.schema([f for f in schema if f.name in present_names]) try: - table = pa.Table.from_pylist(processed) + table = pa.Table.from_pylist(processed, schema=build_schema) except Exception as e: raise ValueError(f"Failed to convert records to Arrow table: {e}") from e else: diff --git a/tests/test_schema.py b/tests/test_schema.py index 6e59fac..6f7cfc7 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -163,6 +163,36 @@ def test_validate_pyarrow_table_input(): assert len(result) == 2 +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), + ] + ) + records = [ + {"id": "a"}, + {"id": "b", "vec": [1.0, 2.0]}, + ] + table = validate_records(records, schema) + assert table.column("vec").to_pylist() == [None, [1.0, 2.0]] + + +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), + ] + ) + records = [ + {"id": "b", "vec": [1.0, 2.0]}, + {"id": "a"}, + ] + table = validate_records(records, schema) + assert table.column("vec").to_pylist() == [[1.0, 2.0], None] + + def test_validate_dict_serialized_to_json(): class WithData(BaseModel): image_id: str