bugfix #4 - validate_records drops col if first row cell null with list

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:

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