Adding in pyarrow schema for jsonschema ingestion pw

src/cdm_data_loaders/pipelines/core.py

@@ -290,6 +290,8 @@
 
     pipeline = dlt.pipeline(**pipeline_kwargs)
 
+    logger.info("Prepping pipeline...")
+
     try:
         pipeline.destination = dlt.destination(settings.use_destination, **(destination_kwargs or {}))
     except Exception as e:
@@ -305,7 +307,9 @@
     load_info: LoadInfo | None = None
     err_msg: str | None = None
     try:
+        logger.info("Running pipeline!")
         load_info = pipeline.run(resource, **(pipeline_run_kwargs or {}))
+        logger.info("Pipeline complete!")
     except PipelineStepFailed as e:
         # errors raise PipelineStepFailed
         err_msg = f"Pipeline failed: {e!s}"

src/cdm_data_loaders/pipelines/jsonlines/base.py

@@ -13,6 +13,7 @@
 from typing import TYPE_CHECKING, Any, Final
 
 import dlt
+import pyarrow as pa
 from dlt.common.pipeline import LoadInfo
 from dlt.common.schema.typing import TSchemaContractDict
 from dlt.sources.filesystem import filesystem
@@ -23,7 +24,7 @@
     run_pipeline,
 )
 from cdm_data_loaders.readers.jsonlines import stream_jsonl_lines
-from cdm_data_loaders.utils.buffer import ListBuffer
+from cdm_data_loaders.utils.buffer import ArrowListBuffer, ListBuffer
 
 if TYPE_CHECKING:
     from dlt.extract.resource import DltResource
@@ -67,6 +68,7 @@
     rejected_suffix: str = DEFAULT_REJECTED_SUFFIX,
     include_record: bool = False,
     transform_valid: Callable[[Any], Any] | None = None,
+    pyarrow_schema: pa.Schema | None = None,
 ) -> Callable[[list[dict[str, Any]]], Generator[Any, Any, Any]]:
     """Build a function that routes pages of line records to valid and rejected tables.
 
@@ -76,6 +78,8 @@
     through `transform_valid` (identity when None) into `table_name`; failed
     records yield a provenance row into ``{table_name}{rejected_suffix}``. With
     `include_record`, the rejected row also carries the parsed `record` value.
+    When `pyarrow_schema` is supplied, valid records are accumulated in an
+    `ArrowListBuffer` and yielded as `pyarrow.Table` pages.
 
     :param table_name: table name for valid records
     :type  table_name: str
@@ -89,13 +93,19 @@
     :type  include_record: bool
     :param transform_valid: function applied to each valid record before yielding
     :type  transform_valid: Callable[[Any], Any] | None
+    :param pyarrow_schema: optional PyArrow schema for buffering valid records as Arrow tables
+    :type  pyarrow_schema: pa.Schema | None
     :return: function taking one page of line records and yielding routed pages
     :rtype : Callable[[list[dict[str, Any]]], Generator[Any, Any, Any]]
     """
     rejected_table = f"{table_name}{rejected_suffix}"
 
     def _route(page: list[dict[str, Any]]) -> Generator[Any, Any, Any]:
-        valid_buffer = ListBuffer(table_name=table_name, max_items=buffer_size)
+        valid_buffer: ListBuffer | ArrowListBuffer = (
+            ArrowListBuffer(table_to_schema_map={table_name: pyarrow_schema}, max_items=buffer_size)
+            if pyarrow_schema is not None
+            else ListBuffer(table_name=table_name, max_items=buffer_size)
+        )
         rejected_buffer = ListBuffer(table_name=rejected_table, max_items=buffer_size)
         for item in page:
             error_detail = item["parse_error"]
@@ -126,6 +136,7 @@
     rejected_suffix: str = DEFAULT_REJECTED_SUFFIX,
     include_record: bool = False,
     transform_valid: Callable[[Any], Any] | None = None,
+    pyarrow_schema: pa.Schema | None = None,
 ) -> "DltResource":
     """Read plain or gzip JSONL files for one entity and route validated pages.
 
@@ -144,6 +155,8 @@
     :type  include_record: bool
     :param transform_valid: function applied to each valid record before yielding
     :type  transform_valid: Callable[[Any], Any] | None
+    :param pyarrow_schema: optional PyArrow schema for buffering valid records as Arrow tables
+    :type  pyarrow_schema: pa.Schema | None
     :return: resource yielding validated and rejected records for this entity
     :rtype: DltResource
     """
@@ -170,6 +183,7 @@
             rejected_suffix=rejected_suffix,
             include_record=include_record,
             transform_valid=transform_valid,
+            pyarrow_schema=pyarrow_schema,
         ),
         data_from=raw,
         name=f"{table_name}_validated",

src/cdm_data_loaders/pipelines/jsonlines/extract_jsonschema_validate_pipeline.py

@@ -2,21 +2,29 @@
 
 The module selected by ``schema_files_module`` must define ``ENTITY_SCHEMAS``:
 a mapping of table names to JSON Schema dictionaries, each declaring ``$schema``.
-Files are read from one input subdirectory per table. Valid records are loaded
-unchanged; parsing and validation failures go to ``<table>_rejected`` with
-source_file, line_no, raw_record, and error_detail. Format checks are enabled.
+Each validated JSON Schema is dereferenced and converted to a ``pyarrow.Schema``
+via ``JSONSchemaToPyArrow``. Files are read from one input subdirectory per
+table. Valid records are coerced to the target ``pyarrow.Schema`` and buffered
+into ``pyarrow.Table`` pages via ``ArrowListBuffer``; parsing and validation
+failures go to ``<table>_rejected`` with source_file, line_no, raw_record, and
+error_detail. Format checks are enabled.
 """
 
 import json
+from collections.abc import Mapping
+from datetime import UTC, date, datetime, time
+from decimal import Decimal
 from logging import Logger, getLogger
 from typing import Any
 
+import pyarrow as pa
 from dlt.common.pipeline import LoadInfo
 from dlt.extract.resource import DltResource
 from jsonschema import FormatChecker
 from jsonschema.exceptions import SchemaError
 from jsonschema.validators import validator_for
 
+from cdm_data_loaders.converters.jsonschema_to_pyarrow import JSONSchemaToPyArrow, JSONSchemaToPyArrowError
 from cdm_data_loaders.pipelines.core import PipelineCli
 from cdm_data_loaders.pipelines.jsonlines.base import (
     ParseErrorResolver,
@@ -27,6 +35,7 @@
     build_entity_resource as build_shared_entity_resource,
 )
 from cdm_data_loaders.pipelines.jsonlines.settings import JSONSCHEMA_PIPELINE_NAME, JsonlJsonschemaIngestSettings
+from cdm_data_loaders.utils.jsonschema.dereferencer import DereferencingError, dereference_schema
 
 logger: Logger = getLogger(__name__)
 
@@ -79,6 +88,118 @@
     return registry
 
 
+def convert_schema_to_pyarrow(schema: dict[str, Any]) -> pa.Schema:
+    """Dereference a validated JSON Schema and convert it to a PyArrow schema.
+
+    :param schema: validated JSON Schema dictionary declaring ``$schema``
+    :type schema: dict[str, Any]
+    :return: target PyArrow schema for the entity table
+    :rtype: pa.Schema
+    :raises DereferencingError: if ``$ref`` or ``allOf`` cannot be resolved
+    :raises JSONSchemaToPyArrowError: if the dereferenced schema cannot be converted to a PyArrow schema
+    """
+    dereferenced = dereference_schema(schema)
+    if "$schema" not in dereferenced and "$schema" in schema:
+        dereferenced = {"$schema": schema["$schema"], **dereferenced}
+    return JSONSchemaToPyArrow().convert(dereferenced)
+
+
+def build_pyarrow_schemas(table_to_schema: dict[str, dict[str, Any]]) -> dict[str, pa.Schema]:
+    """Convert all entity JSON Schemas to PyArrow schemas up front.
+
+    Collects any dereferencing or conversion errors across all tables and raises
+    a single ``RuntimeError`` listing every failing table.
+
+    :param table_to_schema: mapping of table name to validated JSON Schema dictionary
+    :type table_to_schema: dict[str, dict[str, Any]]
+    :return: mapping of table name to PyArrow schema
+    :rtype: dict[str, pa.Schema]
+    :raises RuntimeError: if one or more schemas cannot be converted to PyArrow schemas
+    """
+    pyarrow_schemas: dict[str, pa.Schema] = {}
+    errors: dict[str, str] = {}
+    for table_name, schema in table_to_schema.items():
+        try:
+            pyarrow_schemas[table_name] = convert_schema_to_pyarrow(schema)
+        except (DereferencingError, JSONSchemaToPyArrowError) as error:
+            logger.exception("Error converting JSON Schema to PyArrow for %s", table_name)
+            errors[table_name] = str(error)
+
+    if errors:
+        details = "".join(f"{table_name}: {error}\n" for table_name, error in errors.items())
+        err_msg = f"The following errors were found when converting schemas to PyArrow:\n{details}"
+        raise RuntimeError(err_msg)
+    return pyarrow_schemas
+
+
+def _parse_iso_timestamp(value: str, data_type: pa.TimestampType) -> datetime:
+    """Parse an ISO-8601 timestamp string, attaching UTC when target type is tz-aware."""
+    parsed_dt = datetime.fromisoformat(value)
+    if parsed_dt.tzinfo is None and data_type.tz is not None:
+        return parsed_dt.replace(tzinfo=UTC)
+    return parsed_dt
+
+
+def _coerce_scalar_for_pyarrow(value: Any, data_type: pa.DataType) -> Any:  # noqa: ANN401
+    """Coerce a JSON-parsed value to match a PyArrow scalar type."""
+    coerced = value
+    if pa.types.is_date(data_type) and isinstance(value, str):
+        coerced = date.fromisoformat(value)
+    elif pa.types.is_timestamp(data_type) and isinstance(value, str):
+        coerced = _parse_iso_timestamp(value, data_type)
+    elif pa.types.is_time(data_type) and isinstance(value, str):
+        coerced = time.fromisoformat(value)
+    elif pa.types.is_decimal(data_type) and isinstance(value, int | float | str) and not isinstance(value, bool):
+        coerced = Decimal(str(value))
+    elif pa.types.is_floating(data_type) and isinstance(value, int) and not isinstance(value, bool):
+        coerced = float(value)
+    elif pa.types.is_integer(data_type) and isinstance(value, float) and value.is_integer():
+        coerced = int(value)
+    elif (pa.types.is_string(data_type) or pa.types.is_large_string(data_type)) and not isinstance(value, str):
+        coerced = json.dumps(value) if isinstance(value, Mapping | list) else str(value)
+    return coerced
+
+
+def _coerce_value_for_pyarrow(value: Any, data_type: pa.DataType) -> Any:  # noqa: ANN401
+    """Recursively coerce a JSON-parsed value to match the target PyArrow DataType."""
+    if value is None:
+        return None
+    if pa.types.is_struct(data_type) and isinstance(value, Mapping):
+        return {
+            field.name: _coerce_value_for_pyarrow(value[field.name], field.type)
+            for field in data_type
+            if field.name in value
+        }
+    if (pa.types.is_list(data_type) or pa.types.is_large_list(data_type)) and isinstance(value, list):
+        return [_coerce_value_for_pyarrow(item, data_type.value_type) for item in value]
+    if pa.types.is_map(data_type) and isinstance(value, Mapping):
+        return {
+            _coerce_value_for_pyarrow(k, data_type.key_type): _coerce_value_for_pyarrow(v, data_type.item_type)
+            for k, v in value.items()
+        }
+    return _coerce_scalar_for_pyarrow(value, data_type)
+
+
+def coerce_record_for_pyarrow(record: dict[str, Any], schema: pa.Schema) -> dict[str, Any]:
+    """Coerce a JSON-validated record dictionary to match a PyArrow schema.
+
+    Only fields declared in ``schema`` are retained and coerced recursively so
+    ``pyarrow.Table.from_pylist`` can construct a table matching ``schema``.
+
+    :param record: JSON-parsed and schema-validated record dictionary
+    :type record: dict[str, Any]
+    :param schema: target PyArrow schema for the record's table
+    :type schema: pa.Schema
+    :return: coerced record dictionary ready for ``ArrowListBuffer``
+    :rtype: dict[str, Any]
+    """
+    return {
+        field.name: _coerce_value_for_pyarrow(record[field.name], field.type)
+        for field in schema
+        if field.name in record
+    }
+
+
 def _make_resolve_error(schema: dict[str, Any]) -> ParseErrorResolver:
     """Build the per-record JSON Schema validation closure.
 
@@ -96,39 +217,61 @@
 
 
 def build_entity_resource(
-    table_name: str, schema: dict[str, Any], settings: JsonlJsonschemaIngestSettings
+    table_name: str,
+    schema: dict[str, Any],
+    settings: JsonlJsonschemaIngestSettings,
+    *,
+    pyarrow_schema: pa.Schema | None = None,
 ) -> DltResource:
     """Read plain or gzip JSONL files for one entity and route validated pages.
 
-    A missing entity directory produces no rows. Nested values are preserved
-    rather than normalized into child tables.
+    A missing entity directory produces no rows. Validated records are coerced
+    to the target PyArrow schema and buffered into ``pyarrow.Table`` pages via
+    ``ArrowListBuffer``. Nested values are preserved rather than normalized into
+    child tables.
 
     :param table_name: table name
     :type  table_name: str
     :param schema: JSONSchema for the table
     :type  schema: dict[str, Any]
     :param settings: pipeline configuration
     :type  settings: JsonlJsonschemaIngestSettings
+    :param pyarrow_schema: precomputed PyArrow schema, or None to convert ``schema``
+    :type  pyarrow_schema: pa.Schema | None
     :return: resource yielding validated and rejected records for this entity
     :rtype: DltResource
     """
-    return build_shared_entity_resource(table_name, settings, resolve_error=_make_resolve_error(schema))
+    resolved_pa_schema = pyarrow_schema if pyarrow_schema is not None else convert_schema_to_pyarrow(schema)
+    return build_shared_entity_resource(
+        table_name,
+        settings,
+        resolve_error=_make_resolve_error(schema),
+        transform_valid=lambda record: coerce_record_for_pyarrow(record, resolved_pa_schema),
+        pyarrow_schema=resolved_pa_schema,
+    )
 
 
 def run_jsonlines_pipeline(settings: JsonlJsonschemaIngestSettings) -> LoadInfo | None:
-    """Validate the registry and load requested tables using the shared runner.
+    """Validate the registry, build PyArrow schemas, and load requested tables.
 
     :param settings: pipeline configuration
     :type settings:  JsonlJsonschemaIngestSettings
     :raises ValueError: if settings.table_names includes a name not in ENTITY_SCHEMAS
+    :raises RuntimeError: if any schema fails JSON Schema validation or PyArrow conversion
     :return: load information for the pipeline
     :rtype:  LoadInfo | None
     """
     table_to_schema = load_entity_schemas(settings.schema_files_module)
+    pyarrow_schemas = build_pyarrow_schemas(table_to_schema)
     return run_registry_ingest(
         settings,
         table_to_schema,
-        build_resource=build_entity_resource,
+        build_resource=lambda table_name, schema, cfg: build_entity_resource(
+            table_name,
+            schema,
+            cfg,
+            pyarrow_schema=pyarrow_schemas[table_name],
+        ),
         pipeline_name=JSONSCHEMA_PIPELINE_NAME,
         pipeline_run_kwargs={"loader_file_format": str(settings.loader_file_format)},
     )

tests/cdm_data_loaders/pipelines/jsonlines/integration/test_jsonschema_end_to_end.py

@@ -6,8 +6,11 @@
 from pathlib import Path
 from typing import Any
 
+import pyarrow as pa
+import pyarrow.parquet as pq
 import pytest
 
+from cdm_data_loaders.core.fields import PARQUET
 from cdm_data_loaders.pipelines.jsonlines.settings import JSONSCHEMA_PIPELINE_NAME, JsonlJsonschemaIngestSettings
 from tests.cdm_data_loaders.pipelines.framework import PipelineTestRunner
 from tests.cdm_data_loaders.pipelines.helpers import (
@@ -23,6 +26,20 @@
 
 WIDGET_RECORD = {"widget_id": "a", "count": 0, "nested": {"tags": ["one", "two"]}}
 GADGET_RECORD = {"widget_id": "g", "count": 3}
+EXPECTED_WIDGET_ROW = {
+    "widget_id": "a",
+    "count": 0,
+    "collected": None,
+    "label": None,
+    "nested": {"tags": ["one", "two"]},
+}
+EXPECTED_GADGET_ROW = {
+    "widget_id": "g",
+    "count": 3,
+    "collected": None,
+    "label": None,
+    "nested": None,
+}
 TYPE_ERROR_RAW = '{"widget_id":"bad", "count":"2"}'
 FORMAT_ERROR_RAW = '{"widget_id":"bad","count":1,"collected":"2023-02-29"}'
 MALFORMED_RAW = "{broken"
@@ -73,9 +90,17 @@ def test_run_jsonlines_pipeline_pass_persists_valid_and_rejected_rows(
     tables = read_data_tables(settings, load_info.dataset_name)
     assert set(tables) == {"widget", "widget_rejected", "gadget"}
     assert [{key: parse_json_string(value) for key, value in row.items()} for row in tables["widget"]] == [
-        WIDGET_RECORD
+        EXPECTED_WIDGET_ROW
     ] * 3
-    assert tables["gadget"] == [GADGET_RECORD]
+    assert tables["gadget"] == [EXPECTED_GADGET_ROW]
+    if loader_file_format == PARQUET:
+        widget_parquet = next((Path(settings.output_dir) / load_info.dataset_name / "widget").glob("*.parquet"))
+        parquet_schema = pq.read_schema(widget_parquet)
+        assert parquet_schema.field("widget_id").type == pa.string()
+        assert parquet_schema.field("count").type == pa.int64()
+        assert parquet_schema.field("collected").type == pa.date32()
+        assert parquet_schema.field("label").type == pa.string()
+        assert parquet_schema.field("nested").type == pa.struct([pa.field("tags", pa.list_(pa.string()))])
     assert sorted(tables["widget_rejected"], key=lambda row: (row["source_file"], row["line_no"])) == [
         {
             "source_file": "compressed.jsonl.gz",
@@ -151,4 +176,4 @@ def test_cli_pass_loads_selected_entity(
     )
     assert load_info.dataset_name == "schema_cli"
     assert_pipeline_report_written(Path(settings.output_dir), load_info, metadata_moved=True)
-    assert read_data_tables(settings, load_info.dataset_name) == {"gadget": [GADGET_RECORD]}
+    assert read_data_tables(settings, load_info.dataset_name) == {"gadget": [EXPECTED_GADGET_ROW]}