feat: allow options on ingest queries (#74)

## What's Changed

- Fix null values for structs
- Allow ingest cases to set statement options

Author: lidavidm

adbc_drivers_validation/arrowjson.py

@@ -111,29 +111,41 @@ def array_from_values(values: list, field: pyarrow.Field) -> pyarrow.Array:
             mask=pyarrow.array(mask, type=pyarrow.bool_()) if field.nullable else None,
         )
     elif isinstance(field.type, pyarrow.StructType):
-        children = structlike_from_rows(values, field.type.fields)
-        return pyarrow.StructArray.from_arrays(children, type=field.type)
+        children, valid = structlike_from_rows(values, field.type.fields)
+        if not field.nullable:
+            if not all(valid):
+                raise ValueError(
+                    f"Field '{field.name}' is not nullable but contains None values"
+                )
+        return pyarrow.StructArray.from_arrays(
+            children,
+            type=field.type,
+            mask=pyarrow.array([not v for v in valid], type=pyarrow.bool_()),
+        )
     return pyarrow.array(values, type=field.type)
 
 
 def structlike_from_rows(
     rows: list[dict], fields: list[pyarrow.Field]
-) -> list[pyarrow.Array]:
+) -> tuple[list[pyarrow.Array], list[bool]]:
     cols = []
     for field in fields:
-        values = [row.pop(field.name, None) for row in rows]
+        values = [(row.pop(field.name, None) if row else None) for row in rows]
         array = array_from_values(values, field)
         cols.append(array)
+    valid = [row is not None for row in rows]
 
     for row in rows:
         if row:
             raise ValueError(f"Extra fields in row: {row}")
 
-    return cols
+    return cols, valid
 
 
 def table_from_rows(rows: list[dict], schema: pyarrow.Schema) -> pyarrow.Table:
-    cols = structlike_from_rows(rows, list(schema))
+    cols, valid = structlike_from_rows(rows, list(schema))
+    if not all(valid):
+        raise ValueError("Top-level table cannot have null rows")
     table = pyarrow.Table.from_arrays(cols, schema=schema)
     return table
 

adbc_drivers_validation/tests/ingest.py

@@ -28,6 +28,7 @@
 
 from adbc_drivers_validation import compare, model
 from adbc_drivers_validation.model import Query
+from adbc_drivers_validation.utils import setup_statement
 
 
 def generate_tests(all_quirks: list[model.DriverQuirks], metafunc) -> None:
@@ -107,7 +108,8 @@ def test_create(
 
         with conn.cursor() as cursor:
             cursor.execute(driver.drop_table(table_name=table_name))
-            cursor.adbc_ingest(table_name, data, mode="create")
+            with setup_statement(query, cursor):
+                cursor.adbc_ingest(table_name, data, mode="create")
 
         idx = driver.quote_identifier("idx")
         value = driver.quote_identifier("value")

adbc_drivers_validation/tests/query.py

@@ -28,7 +28,7 @@
 
 from adbc_drivers_validation import compare, model
 from adbc_drivers_validation.model import Query
-from adbc_drivers_validation.utils import scoped_trace
+from adbc_drivers_validation.utils import scoped_trace, setup_statement
 
 
 def generate_tests(all_quirks: list[model.DriverQuirks], metafunc) -> None:
@@ -112,7 +112,7 @@ def test_query(
                 cursor.adbc_statement.execute_update()
 
         with conn.cursor() as cursor:
-            with _setup_statement(query, cursor):
+            with setup_statement(query, cursor):
                 cursor.adbc_statement.set_sql_query(sql)
                 handle, _ = cursor.adbc_statement.execute_query()
                 with pyarrow.RecordBatchReader._import_from_c(handle.address) as reader:
@@ -135,7 +135,7 @@ def test_execute_schema(
         _setup_query(driver, conn, query)
 
         with conn.cursor() as cursor:
-            with _setup_statement(query, cursor):
+            with setup_statement(query, cursor):
                 schema = cursor.adbc_execute_schema(sql)
 
         compare.compare_schemas(expected_schema, schema)
@@ -227,31 +227,3 @@ def _setup_query(
                 with scoped_trace(f"setup statement: {statement}"):
                     cursor.adbc_statement.set_sql_query(statement)
                     cursor.adbc_statement.execute_update()
-
-
-@contextlib.contextmanager
-def _setup_statement(query: Query, cursor: adbc_driver_manager.dbapi.Cursor) -> None:
-    md = query.metadata()
-    if "statement" not in md:
-        yield
-        return
-
-    statement_md = md["statement"]
-    if "options" not in statement_md:
-        yield
-        return
-
-    options = {}
-    options_revert = {}
-    for key, value in statement_md["options"].items():
-        if isinstance(value, dict):
-            assert "apply" in value
-            assert "revert" in value
-            options[key] = value["apply"]
-            options_revert[key] = value["revert"]
-        else:
-            options[key] = value
-
-    cursor.adbc_statement.set_options(**options)
-    yield
-    cursor.adbc_statement.set_options(**options_revert)

adbc_drivers_validation/utils.py

@@ -15,6 +15,11 @@
 import contextlib
 import typing
 
+import adbc_driver_manager.dbapi
+
+if typing.TYPE_CHECKING:
+    from adbc_drivers_validation.model import Query
+
 
 def merge_into(target: dict[str, typing.Any], values: dict[str, typing.Any]) -> None:
     for key, value in values.items():
@@ -36,3 +41,31 @@ def scoped_trace(msg: str) -> None:
     except Exception as e:
         e.add_note(msg)
         raise
+
+
+@contextlib.contextmanager
+def setup_statement(query: "Query", cursor: adbc_driver_manager.dbapi.Cursor) -> None:
+    md = query.metadata()
+    if "statement" not in md:
+        yield
+        return
+
+    statement_md = md["statement"]
+    if "options" not in statement_md:
+        yield
+        return
+
+    options = {}
+    options_revert = {}
+    for key, value in statement_md["options"].items():
+        if isinstance(value, dict):
+            assert "apply" in value
+            assert "revert" in value
+            options[key] = value["apply"]
+            options_revert[key] = value["revert"]
+        else:
+            options[key] = value
+
+    cursor.adbc_statement.set_options(**options)
+    yield
+    cursor.adbc_statement.set_options(**options_revert)

tests/test_arrowjson.py

@@ -220,7 +220,6 @@ def test_array_from_values_list() -> None:
 
 
 def test_array_from_values_struct() -> None:
-    values = [{"a": 1, "b": base64.b64encode(b"hello").decode("utf-8")}]
     field = pyarrow.field(
         "struct_field",
         pyarrow.struct(
@@ -230,9 +229,15 @@ def test_array_from_values_struct() -> None:
             ]
         ),
     )
+
+    values = [
+        None,
+        {"a": 1, "b": base64.b64encode(b"hello").decode("utf-8")},
+        {"a": None, "b": base64.b64encode(b"hello").decode("utf-8")},
+    ]
     actual = arrowjson.array_from_values(values, field)
     expected = pyarrow.array(
-        [{"a": 1, "b": b"hello"}],
+        [None, {"a": 1, "b": b"hello"}, {"a": None, "b": b"hello"}],
         type=field.type,
     )
     assert actual == expected