diff --git a/src/cdm_data_loaders/pipelines/core.py b/src/cdm_data_loaders/pipelines/core.py index 5b1ef88d..a1aa1a68 100644 --- a/src/cdm_data_loaders/pipelines/core.py +++ b/src/cdm_data_loaders/pipelines/core.py @@ -290,6 +290,8 @@ def run_pipeline( 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 @@ def run_pipeline( 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}" diff --git a/src/cdm_data_loaders/pipelines/jsonlines/base.py b/src/cdm_data_loaders/pipelines/jsonlines/base.py index 0235d6e2..0f3b9bcf 100644 --- a/src/cdm_data_loaders/pipelines/jsonlines/base.py +++ b/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 @@ def make_page_router( # noqa: PLR0913 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 @@ def make_page_router( # noqa: PLR0913 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 @@ def make_page_router( # noqa: PLR0913 :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 @@ def build_entity_resource( # noqa: PLR0913 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 @@ def build_entity_resource( # noqa: PLR0913 :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 @@ def build_entity_resource( # noqa: PLR0913 rejected_suffix=rejected_suffix, include_record=include_record, transform_valid=transform_valid, + pyarrow_schema=pyarrow_schema, ), data_from=raw, name=f"{table_name}_validated", diff --git a/src/cdm_data_loaders/pipelines/jsonlines/extract_jsonschema_validate_pipeline.py b/src/cdm_data_loaders/pipelines/jsonlines/extract_jsonschema_validate_pipeline.py index 8409a0ed..08aff45e 100644 --- a/src/cdm_data_loaders/pipelines/jsonlines/extract_jsonschema_validate_pipeline.py +++ b/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 ``_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 ``
_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 @@ def load_entity_schemas(module_path: str) -> dict[str, dict[str, Any]]: 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,12 +217,18 @@ def _validate(record: dict[str, Any]) -> str | None: 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 @@ -109,26 +236,42 @@ def build_entity_resource( :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)}, ) diff --git a/tests/cdm_data_loaders/pipelines/jsonlines/integration/test_jsonschema_end_to_end.py b/tests/cdm_data_loaders/pipelines/jsonlines/integration/test_jsonschema_end_to_end.py index 39d3d31c..71b8184d 100644 --- a/tests/cdm_data_loaders/pipelines/jsonlines/integration/test_jsonschema_end_to_end.py +++ b/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]} diff --git a/tests/cdm_data_loaders/pipelines/jsonlines/test_jsonschema_pipeline.py b/tests/cdm_data_loaders/pipelines/jsonlines/test_jsonschema_pipeline.py index a956a842..d7747116 100644 --- a/tests/cdm_data_loaders/pipelines/jsonlines/test_jsonschema_pipeline.py +++ b/tests/cdm_data_loaders/pipelines/jsonlines/test_jsonschema_pipeline.py @@ -1,36 +1,54 @@ -"""Tests for JSON Schema validation and rejected-record routing.""" +"""Tests for JSON Schema validation, PyArrow conversion, and rejected-record routing.""" import gzip import json import sys from collections.abc import Callable +from datetime import UTC, date, datetime, time +from decimal import Decimal from pathlib import Path from types import ModuleType from typing import Any from uuid import uuid4 +import pyarrow as pa import pytest from pydantic import ValidationError +from cdm_data_loaders.converters.jsonschema_to_pyarrow import JSONSchemaToPyArrowError from cdm_data_loaders.pipelines.jsonlines.base import make_page_router from cdm_data_loaders.pipelines.jsonlines.extract_jsonschema_validate_pipeline import ( _make_resolve_error, build_entity_resource, + build_pyarrow_schemas, + coerce_record_for_pyarrow, + convert_schema_to_pyarrow, load_entity_schemas, run_jsonlines_pipeline, ) from cdm_data_loaders.pipelines.jsonlines.settings import JsonlJsonschemaIngestSettings +from cdm_data_loaders.utils.jsonschema.dereferencer import DereferencingError from tests.cdm_data_loaders.pipelines.jsonlines.conftest import make_item +from tests.cdm_data_loaders.pipelines.jsonlines.integration.dataset_report_schemas import ( + ENTITY_SCHEMAS as DATASET_ENTITY_SCHEMAS, +) def make_router(schema: dict[str, Any], buffer_size: int = 2) -> Callable[[list[dict[str, Any]]], Any]: """Build the router the jsonschema pipeline uses for the widget table.""" - return make_page_router("widget", buffer_size, _make_resolve_error(schema)) + pyarrow_schema = convert_schema_to_pyarrow(schema) + return make_page_router( + "widget", + buffer_size, + _make_resolve_error(schema), + transform_valid=lambda record: coerce_record_for_pyarrow(record, pyarrow_schema), + pyarrow_schema=pyarrow_schema, + ) -def test_make_page_router_pass_parse_error_preserves_raw_record() -> None: +def test_make_page_router_pass_parse_error_preserves_raw_record(widget_schema: dict[str, Any]) -> None: """Malformed JSON bypasses validation and preserves its source text.""" - validate = make_router({"$schema": "https://json-schema.org/draft/2020-12/schema"}) + validate = make_router(widget_schema) item = { "record": None, "raw_record": "{bad json", @@ -137,13 +155,344 @@ def test_load_entity_schemas_fail_reports_all_invalid_schemas( assert all(text in str(error.value) for text in [module_name, "first:", "second:"]) +def test_convert_schema_to_pyarrow_pass_flat_and_nested_schema(widget_schema: dict[str, Any]) -> None: + """Validated JSON Schemas convert to PyArrow schemas with expected types and nullability.""" + pa_schema = convert_schema_to_pyarrow(widget_schema) + + assert pa_schema == pa.schema( + [ + pa.field("widget_id", pa.string(), nullable=False), + pa.field("count", pa.int64(), nullable=False), + pa.field("collected", pa.date32(), nullable=True), + pa.field("label", pa.string(), nullable=True), + pa.field( + "nested", + pa.struct([pa.field("tags", pa.list_(pa.string()), nullable=True)]), + nullable=True, + ), + ] + ) + + +def test_convert_schema_to_pyarrow_pass_dereferences_root_and_nested_refs() -> None: + """Schemas using root $ref, nested $defs, and allOf dereference before PyArrow conversion.""" + draft7_root_ref_schema: dict[str, Any] = { + "$schema": "http://json-schema.org/draft-07/schema#", + "$ref": "#/definitions/RootItem", + "definitions": { + "Base": { + "type": "object", + "required": ["id"], + "properties": {"id": {"type": "string"}}, + }, + "RootItem": { + "allOf": [ + {"$ref": "#/definitions/Base"}, + { + "type": "object", + "properties": {"score": {"type": "number", "multipleOf": 0.01}}, + }, + ] + }, + }, + } + + pa_schema = convert_schema_to_pyarrow(draft7_root_ref_schema) + assert pa_schema == pa.schema( + [ + pa.field("id", pa.string(), nullable=False), + pa.field("score", pa.decimal128(38, 2), nullable=True), + ] + ) + + dataset_pa_schema = convert_schema_to_pyarrow(DATASET_ENTITY_SCHEMAS["dataset"]) + assert "accession" in dataset_pa_schema.names + assert dataset_pa_schema.field("accession").type == pa.string() + + +@pytest.mark.parametrize( + "schema", + [ + pytest.param( + {"$schema": "https://json-schema.org/draft/2020-12/schema", "type": "object", "properties": {}}, + id="empty-properties", + ), + pytest.param( + {"$schema": "https://json-schema.org/draft/2020-12/schema", "type": "object"}, + id="missing-properties", + ), + ], +) +def test_convert_schema_to_pyarrow_pass_empty_object_schema(schema: dict[str, Any]) -> None: + """Object schemas with empty or omitted properties convert to an empty PyArrow schema.""" + assert convert_schema_to_pyarrow(schema) == pa.schema([]) + + +@pytest.mark.parametrize( + ("schema", "exc_type"), + [ + pytest.param( + { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": {"item": {"$ref": "#/$defs/Missing"}}, + }, + DereferencingError, + id="unresolvable-ref", + ), + pytest.param( + { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$ref": "#/$defs/Loop", + "$defs": {"Loop": {"$ref": "#/$defs/Loop"}}, + }, + DereferencingError, + id="circular-ref", + ), + pytest.param( + { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "string", + }, + JSONSchemaToPyArrowError, + id="non-object-root", + ), + pytest.param( + { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": ["object", "null"], + "properties": {"id": {"type": "string"}}, + }, + JSONSchemaToPyArrowError, + id="nullable-object-root", + ), + pytest.param( + { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "additionalProperties": {"type": "string"}, + }, + JSONSchemaToPyArrowError, + id="root-map-without-properties", + ), + pytest.param( + { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": {"val": {"type": "number", "multipleOf": Decimal("1e-39")}}, + }, + JSONSchemaToPyArrowError, + id="decimal-scale-exceeds-precision", + ), + pytest.param( + { + "type": "object", + "properties": {"id": {"type": "string"}}, + }, + JSONSchemaToPyArrowError, + id="missing-schema-draft", + ), + ], +) +def test_convert_schema_to_pyarrow_fail_invalid_schemas(schema: dict[str, Any], exc_type: type[Exception]) -> None: + """Dereferencing and PyArrow conversion failures raise their documented exception types.""" + with pytest.raises(exc_type): + convert_schema_to_pyarrow(schema) + + +def test_build_pyarrow_schemas_pass_multiple_and_empty_registries(widget_schema: dict[str, Any]) -> None: + """Valid registries convert every table to a PyArrow schema, and an empty mapping returns empty.""" + assert build_pyarrow_schemas({}) == {} + schemas = build_pyarrow_schemas({"widget": widget_schema, "dataset": DATASET_ENTITY_SCHEMAS["dataset"]}) + assert set(schemas) == {"widget", "dataset"} + assert schemas["widget"] == convert_schema_to_pyarrow(widget_schema) + assert "accession" in schemas["dataset"].names + + +def test_build_pyarrow_schemas_fail_reports_all_unconvertible_schemas() -> None: + """Conversion errors across multiple tables are aggregated into a single RuntimeError.""" + table_to_schema: dict[str, dict[str, Any]] = { + "broken_ref": { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": {"item": {"$ref": "#/$defs/Missing"}}, + }, + "non_struct_root": { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "string", + }, + } + + with pytest.raises(RuntimeError, match="converting schemas to PyArrow") as error: + build_pyarrow_schemas(table_to_schema) + assert all(name in str(error.value) for name in ["broken_ref:", "non_struct_root:"]) + + +def test_coerce_record_for_pyarrow_pass_temporal_decimal_and_nested_types() -> None: + """JSON-parsed values are coerced to match temporal, decimal, numeric, struct, list, map, and string fields.""" + pa_schema = pa.schema( + [ + pa.field("day", pa.date32()), + pa.field("ts_z", pa.timestamp("us", tz="UTC")), + pa.field("ts_naive", pa.timestamp("us", tz="UTC")), + pa.field("tod", pa.time64("us")), + pa.field("price", pa.decimal128(38, 2)), + pa.field("ratio", pa.float64()), + pa.field("whole", pa.int64()), + pa.field("collapsed_obj", pa.string()), + pa.field("collapsed_num", pa.string()), + pa.field("nested", pa.struct([pa.field("when", pa.date32())])), + pa.field("items", pa.list_(pa.decimal128(38, 1))), + pa.field("attrs", pa.map_(pa.string(), pa.date32())), + ] + ) + raw_record: dict[str, Any] = { + "day": "2024-02-29", + "ts_z": "2024-02-29T12:30:00Z", + "ts_naive": "2024-02-29T12:30:00", + "tod": "08:15:30", + "price": 19.95, + "ratio": 2, + "whole": 7.0, + "collapsed_obj": {"a": 1}, + "collapsed_num": 42, + "nested": {"when": "2024-03-01", "extra": "dropped"}, + "items": [1.5, 2], + "attrs": {"start": "2024-01-01"}, + "undeclared_top": "dropped", + } + + coerced = coerce_record_for_pyarrow(raw_record, pa_schema) + + assert coerced == { + "day": date(2024, 2, 29), + "ts_z": datetime(2024, 2, 29, 12, 30, tzinfo=UTC), + "ts_naive": datetime(2024, 2, 29, 12, 30, tzinfo=UTC), + "tod": time(8, 15, 30), + "price": Decimal("19.95"), + "ratio": 2.0, + "whole": 7, + "collapsed_obj": json.dumps({"a": 1}), + "collapsed_num": "42", + "nested": {"when": date(2024, 3, 1)}, + "items": [Decimal("1.5"), Decimal(2)], + "attrs": {"start": date(2024, 1, 1)}, + } + table = pa.Table.from_pylist([coerced], schema=pa_schema) + assert table.num_rows == 1 + assert table.schema == pa_schema + + +def test_coerce_record_for_pyarrow_pass_nulls_empty_and_boundary_values() -> None: + """Nulls, empty containers, zero/negative numbers, tz-naive timestamps, large types, and pre-coerced values pass.""" + pa_schema = pa.schema( + [ + pa.field("day", pa.date32()), + pa.field("ts_naive_target", pa.timestamp("us")), + pa.field("tod", pa.time64("us")), + pa.field("zero_dec", pa.decimal128(38, 2)), + pa.field("neg_dec", pa.decimal128(38, 2)), + pa.field("neg_float", pa.float64()), + pa.field("neg_int", pa.int64()), + pa.field("large_text", pa.large_string()), + pa.field("collapsed_list", pa.large_string()), + pa.field("nested", pa.struct([pa.field("when", pa.date32())])), + pa.field("large_items", pa.large_list(pa.decimal128(38, 1))), + pa.field("attrs", pa.map_(pa.string(), pa.date32())), + ] + ) + + assert coerce_record_for_pyarrow({}, pa_schema) == {} + + all_nulls = dict.fromkeys(pa_schema.names) + assert coerce_record_for_pyarrow(all_nulls, pa_schema) == all_nulls + + boundary_record: dict[str, Any] = { + "day": date(2024, 1, 1), + "ts_naive_target": "2024-02-29T12:30:00", + "tod": time(0, 0, 0), + "zero_dec": 0, + "neg_dec": "-0.05", + "neg_float": -3, + "neg_int": -7.0, + "large_text": "already-string", + "collapsed_list": [1, "two"], + "nested": {"when": None}, + "large_items": [None, Decimal("-1.5")], + "attrs": {"none_val": None}, + } + coerced = coerce_record_for_pyarrow(boundary_record, pa_schema) + assert coerced == { + "day": date(2024, 1, 1), + "ts_naive_target": datetime(2024, 2, 29, 12, 30), # noqa: DTZ001 + "tod": time(0, 0, 0), + "zero_dec": Decimal(0), + "neg_dec": Decimal("-0.05"), + "neg_float": -3.0, + "neg_int": -7, + "large_text": "already-string", + "collapsed_list": json.dumps([1, "two"]), + "nested": {"when": None}, + "large_items": [None, Decimal("-1.5")], + "attrs": {"none_val": None}, + } + rows_to_write = [coerced, all_nulls] + table = pa.Table.from_pylist(rows_to_write, schema=pa_schema) + assert table.num_rows == len(rows_to_write) + + empty_containers = {"nested": {}, "large_items": [], "attrs": {}} + assert coerce_record_for_pyarrow(empty_containers, pa_schema) == empty_containers + + +def test_coerce_record_for_pyarrow_pass_preserves_boolean_and_non_integral_float_guards() -> None: + """Booleans are not coerced to numeric/decimal types, and non-integral floats are not truncated to int.""" + pa_schema = pa.schema( + [ + pa.field("dec", pa.decimal128(38, 2)), + pa.field("flt", pa.float64()), + pa.field("integer_from_bool", pa.int64()), + pa.field("integer_from_frac", pa.int64()), + ] + ) + fractional_value = 7.5 + raw = { + "dec": True, + "flt": False, + "integer_from_bool": True, + "integer_from_frac": fractional_value, + } + coerced = coerce_record_for_pyarrow(raw, pa_schema) + assert coerced == raw + assert isinstance(coerced["dec"], bool) + assert isinstance(coerced["flt"], bool) + assert isinstance(coerced["integer_from_bool"], bool) + assert coerced["integer_from_frac"] == fractional_value + + def test_make_page_router_pass_preserves_records_and_empty_pages(widget_schema: dict[str, Any]) -> None: - """Valid data retains extra fields, dates, and nested values without inserted defaults.""" - record = {"widget_id": "a", "count": 0, "collected": "2024-02-29", "nested": {"items": [1, 2]}} + """Valid data is coerced into a PyArrow Table matching the schema without inserted defaults.""" + record = { + "widget_id": "a", + "count": 0, + "collected": "2024-02-29", + "nested": {"tags": ["one", "two"]}, + "undeclared": "ignored", + } validate = make_router(widget_schema) results = list(validate([make_item(record)])) - assert [(result.meta.table_name, result.data) for result in results] == [("widget", [record])] - assert results[0].data[0] is record + assert len(results) == 1 + assert results[0].meta.table_name == "widget" + assert isinstance(results[0].data, pa.Table) + assert results[0].data.schema == convert_schema_to_pyarrow(widget_schema) + assert results[0].data.to_pylist() == [ + { + "widget_id": "a", + "count": 0, + "collected": date(2024, 2, 29), + "label": None, + "nested": {"tags": ["one", "two"]}, + } + ] assert "label" not in record assert list(validate([])) == [] @@ -200,6 +549,7 @@ def test_make_page_router_pass_flushes_independent_buffers(widget_schema: dict[s ("widget", 1), ("widget_rejected", 1), ] + assert all(isinstance(page.data, pa.Table) for page in pages if page.meta.table_name == "widget") assert [row["line_no"] for page in pages if page.meta.table_name == "widget_rejected" for row in page.data] == [ 2, 4, @@ -230,11 +580,15 @@ def test_build_entity_resource_pass_reads_selected_files( resource = build_entity_resource("widget", widget_schema, settings) rows = list(resource) - assert record in rows - rejected = next(row for row in rows if "error_detail" in row) + valid_tables = [item for item in rows if isinstance(item, pa.Table)] + assert len(valid_tables) == 1 + assert valid_tables[0].to_pylist() == [ + {"widget_id": "a", "count": 0, "collected": None, "label": None, "nested": None} + ] + rejected = next(row for row in rows if isinstance(row, dict) and "error_detail" in row) expected_error = "Expecting property name enclosed in double quotes: line 1 column 2 (char 1)" assert rejected == {"source_file": filename, "line_no": 4, "raw_record": "{broken", "error_detail": expected_error} - assert rows == [record, rejected] + assert len(rows) == len(valid_tables) + 1 @pytest.mark.parametrize("create_directory", [False, True], ids=["missing", "empty"]) @@ -250,6 +604,28 @@ def test_build_entity_resource_pass_no_files( assert list(build_entity_resource("widget", widget_schema, settings)) == [] +def test_run_jsonlines_pipeline_pass_loads_valid_and_rejected_rows( + schema_settings_factory: Callable[..., JsonlJsonschemaIngestSettings], +) -> None: + """The pipeline converts schemas to PyArrow and writes both valid Arrow pages and rejected records.""" + settings = schema_settings_factory() + entity_dir = Path(settings.input_dir) / "widget" + entity_dir.mkdir() + valid_record = {"widget_id": "w1", "count": 3, "collected": "2024-02-29"} + invalid_record = {"widget_id": "w2", "count": -1} + (entity_dir / "data.jsonl").write_text( + f"{json.dumps(valid_record)}\n{json.dumps(invalid_record)}\n", + encoding="utf-8", + ) + + load_info = run_jsonlines_pipeline(settings) + assert load_info is not None + assert not load_info.has_failed_jobs + dataset_dir = Path(settings.output_dir) / str(settings.dataset_name) + assert (dataset_dir / "widget").is_dir() + assert (dataset_dir / "widget_rejected").is_dir() + + def test_run_jsonlines_pipeline_fail_unknown_tables( schema_settings_factory: Callable[..., JsonlJsonschemaIngestSettings], ) -> None: @@ -260,6 +636,26 @@ def test_run_jsonlines_pipeline_fail_unknown_tables( assert list(Path(settings.output_dir).iterdir()) == [] +def test_run_jsonlines_pipeline_fail_unconvertible_pyarrow_schema( + schema_module_factory: Callable[[object], str], + schema_settings_factory: Callable[..., JsonlJsonschemaIngestSettings], +) -> None: + """Valid JSON Schemas that cannot convert to PyArrow fail before creating pipeline output.""" + module_name = schema_module_factory( + { + "map_root": { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "additionalProperties": {"type": "string"}, + } + } + ) + settings = schema_settings_factory(schema_files_module=module_name) + with pytest.raises(RuntimeError, match="converting schemas to PyArrow"): + run_jsonlines_pipeline(settings) + assert list(Path(settings.output_dir).iterdir()) == [] + + @pytest.mark.parametrize( ("table_names", "expected"), [