New xsd/pyarrow XML data loader

.github/workflows/tests.yml

@@ -85,8 +85,8 @@ jobs:
         suite:
           - name: pipelines
             args: tests/cdm_data_loaders/pipelines/
-          - name: docker
-            args: --run-docker -m "docker and not requires_spark and not requires_ceph" tests/cdm_data_loaders/pipelines/
+          # - name: docker
+          #   args: --run-docker -m "docker and not requires_spark and not requires_ceph" tests/cdm_data_loaders/pipelines/
           - name: integration
             args: tests/integration/
           - name: everything else

pyproject.toml

@@ -49,6 +49,7 @@ pdb_manifest = "cdm_data_loaders.pipelines.pdb_manifest:cli"
 uniprot_cdm = "cdm_data_loaders.pipelines.uniprot_kb_cdm:cli"
 uniref_cdm = "cdm_data_loaders.pipelines.uniref_cdm:cli"
 xmltodict = "cdm_data_loaders.pipelines.xmltodict.pipeline:cli"
+xmltopyarrow = "cdm_data_loaders.pipelines.xmltopyarrow.pipeline:cli"
 xml2db = "cdm_data_loaders.pipelines.xml2db.pipeline:cli"
 xsv = "cdm_data_loaders.pipelines.xsv.pipeline:cli"
 hello_world = "cdm_data_loaders.utils.hello_world:cli"

src/cdm_data_loaders/pipelines/ncbi_ftp_promote.py

@@ -4,7 +4,6 @@
 final location in the Lakehouse.
 """
 
-import logging
 from pathlib import Path, PurePosixPath
 from typing import Final
 
@@ -22,8 +21,6 @@
 DEFAULT_TRANSFER_MANIFEST_FILE: PurePosixPath = PurePosixPath("transfer_manifest.txt")
 ESTIMATED_FILES_PER_ASSEMBLY: int = 21
 
-logger = logging.getLogger("dlt")
-
 
 class PromoteSettings(LoggerSettings):
     """Configuration for the NCBI FTP file promotion pipeline."""

src/cdm_data_loaders/pipelines/xmltodict/pipeline.py

@@ -1,7 +1,6 @@
 """XML to dictionary pipeline for the KBase CTS."""
 
 from collections.abc import Generator, Iterator
-from logging import Logger, getLogger
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
@@ -18,15 +17,15 @@
 if TYPE_CHECKING:
     from dlt.extract import DltResource
 
-logger: Logger = getLogger(__name__)
 
-
-def _read_items(items: Iterator[FileItemDict], settings: XmlToDictSettings) -> Generator[TDataItems, Any, Any]:
+def _read_xmltodict_items(
+    items: Iterator[FileItemDict], settings: XmlToDictSettings
+) -> Generator[TDataItems, Any, Any]:
     """Read each file in items. Yields one dict per `xml_tag` element.
 
     :param items: file items to read
     :type  items: Iterator[FileItemDict]
-    :param settings: pipeline config, including xml_tag and xsd-derived paths
+    :param settings: pipeline config, including xml_tag
     :type  settings: XmlToDictSettings
     :yield: one dict per matched element
     :rtype: Generator[TDataItems, Any, Any]
@@ -41,8 +40,7 @@
     :param settings: pipeline configuration
     :type  settings: XmlToDictSettings
     """
-    reader = dlt.transformer(_read_items, name="xmltodict_reader", parallelized=True)
-
+    reader = dlt.transformer(_read_xmltodict_items, name="xmltodict_reader", parallelized=True)
     files = filesystem(bucket_url=str(settings.input_dir), file_glob=settings.file_glob)
 
     xmltodict_resource: DltResource = files | reader(settings)

src/cdm_data_loaders/pipelines/xmltodict/settings.py

@@ -1,9 +1,8 @@
 """Settings class for the XML to dictionary pipeline for the KBase CTS."""
 
-from functools import cached_property
-from typing import Annotated, Any, Final, Self
+from typing import Annotated, Final
 
-from pydantic import Field, FilePath, PrivateAttr, computed_field, model_validator
+from pydantic import Field
 from pydantic_settings import SettingsConfigDict
 
 from cdm_data_loaders.core.fields import (
@@ -17,7 +16,6 @@
     FileFormatIngestSettings,
     default_settings_with_shortcuts,
 )
-from cdm_data_loaders.readers.xsd import find_list_and_single_child_paths, load_schema
 
 PIPELINE_NAME: Final[str] = "xmltodict"
 
@@ -33,9 +31,6 @@ class XmlToDictSettings(FileFormatIngestSettings):
         },
     )
 
-    _paths: Annotated[dict[str, Any], Field(description="parent-child paths in the schema")] = PrivateAttr(
-        default_factory=dict
-    )
     file_glob: Annotated[
         FileGlob,
         Field(
@@ -50,45 +45,3 @@ class XmlToDictSettings(FileFormatIngestSettings):
             description="XML tag to capture the contents of",
         ),
     ]
-    xsd_file: Annotated[
-        FilePath | None,
-        Field(
-            None,
-            description="Path to the XSD file describing the XML documents to be imported. Including the schema will ensure that the appropriate relationships (1-to-1 vs 1-to-many) are created between parent and child elements.",
-        ),
-    ]
-
-    @model_validator(mode="after")
-    def check_xsd_file(self) -> Self:
-        """Check whether the schema file is present and correct; if so, init self._paths."""
-        if self.xsd_file:
-            try:
-                schema = load_schema(self.xsd_file)
-                list_paths, single_paths = find_list_and_single_child_paths(schema)
-                self._paths = {"list": list_paths, "single": single_paths}
-            except Exception as e:
-                err_msg = f"Could not generate parent-child relationships for schema: {e!s}"
-                raise ValueError(err_msg) from e
-
-        return self
-
-    @computed_field
-    @cached_property
-    def xmltodict_args(self) -> dict[str, Any]:
-        """Extra parser configuration arguments for running xmltodict over a dataset.
-
-        :return: kwargs for configuring xmltodict.parse()
-        :rtype: dict[str, Any]
-        """
-        if not self._paths:
-            return {}
-
-        single_paths = self._paths["single"]
-
-        def force_list(path: list[tuple[str, str | None]], key: str, _: Any) -> bool:
-            if not path:
-                return False
-
-            return (path[-1][0], key) not in single_paths
-
-        return {"force_list": force_list}

src/cdm_data_loaders/pipelines/xmltopyarrow/pipeline.py

@@ -0,0 +1,63 @@
+"""XML to pyarrow pipeline for the KBase CTS."""
+
+from collections.abc import Generator, Iterator
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import dlt
+from dlt.common.pipeline import LoadInfo
+from dlt.common.storages.fsspec_filesystem import FileItemDict
+from dlt.common.typing import TDataItems
+from dlt.sources.filesystem import filesystem
+
+from cdm_data_loaders.pipelines.core import PipelineCli, run_pipeline
+from cdm_data_loaders.pipelines.xmltopyarrow.settings import PIPELINE_NAME, XmlToPyArrowSettings
+from cdm_data_loaders.readers.xml import process_xml_file_to_pyarrow
+
+if TYPE_CHECKING:
+    from dlt.extract import DltResource
+
+
+def _read_items(items: Iterator[FileItemDict], settings: XmlToPyArrowSettings) -> Generator[TDataItems, Any, Any]:
+    """Read each file in items. Yields PyArrow tables for matched `xml_tag` elements.
+
+    :param items: file items to read
+    :type  items: Iterator[FileItemDict]
+    :param settings: pipeline config, including xml_tag and XSD schema
+    :type  settings: XmlToPyArrowSettings
+    :yield: PyArrow tables of matched elements
+    :rtype: Generator[TDataItems, Any, Any]
+    """
+    for file_item in items:
+        yield from process_xml_file_to_pyarrow(settings, file_path=Path(file_item.local_file_path))
+
+
+def run_xmltopyarrow_pipeline(settings: XmlToPyArrowSettings) -> LoadInfo | None:
+    """Run the XML to PyArrow pipeline on matching files in settings.input_dir.
+
+    :param settings: pipeline configuration
+    :type  settings: XmlToPyArrowSettings
+    """
+    reader = dlt.transformer(_read_items, name="xmltopyarrow_reader", parallelized=True)
+    files = filesystem(bucket_url=str(settings.input_dir), file_glob=settings.file_glob)
+
+    xmltopyarrow_resource: DltResource = files | reader(settings)
+
+    return run_pipeline(
+        settings=settings,
+        resource=xmltopyarrow_resource,
+        pipeline_kwargs={
+            "pipeline_name": PIPELINE_NAME,
+            "dataset_name": settings.dataset_name,
+        },
+        pipeline_run_kwargs={
+            "loader_file_format": str(settings.loader_file_format),
+        },
+    )
+
+
+cli = PipelineCli(XmlToPyArrowSettings, run_xmltopyarrow_pipeline)
+
+
+if __name__ == "__main__":
+    cli()

src/cdm_data_loaders/pipelines/xmltopyarrow/settings.py

@@ -0,0 +1,95 @@
+"""Settings class for the XML to PyArrow pipeline for the KBase CTS."""
+
+from typing import Annotated, Final, Self
+
+import pyarrow as pa
+import xmlschema
+from pydantic import Field, FilePath, PositiveInt, PrivateAttr, model_validator
+from pydantic_settings import SettingsConfigDict
+
+from cdm_data_loaders.core.fields import (
+    DEFAULT_XML_FILE_GLOB,
+    FILE_GLOB,
+    FileGlob,
+    NonEmptyStr,
+    TableName,
+)
+from cdm_data_loaders.core.settings import (
+    FileFormatIngestSettings,
+    default_settings_with_shortcuts,
+)
+from cdm_data_loaders.readers.xsd import (
+    DEFAULT_MAX_RECURSION_DEPTH,
+    generate_pyarrow_schema,
+    load_schema,
+)
+
+PIPELINE_NAME: Final[str] = "xmltopyarrow"
+
+
+class XmlToPyArrowSettings(FileFormatIngestSettings):
+    """Settings for the XML to PyArrow ingestion pipeline."""
+
+    model_config: SettingsConfigDict = default_settings_with_shortcuts(
+        cli_prog_name=PIPELINE_NAME,
+        cli_shortcuts={
+            **FileFormatIngestSettings.model_config["cli_shortcuts"],
+            FILE_GLOB: "g",
+        },
+    )
+
+    _xsd_schema: xmlschema.XMLSchema = PrivateAttr()
+    _pyarrow_schema: pa.Schema = PrivateAttr()
+    file_glob: Annotated[
+        FileGlob,
+        Field(
+            default=DEFAULT_XML_FILE_GLOB,
+            description="Glob pattern for XML files inside each entity's input subdirectory.",
+        ),
+    ]
+    table_name: TableName
+    xml_tag: Annotated[
+        NonEmptyStr,
+        Field(
+            description="XML tag to capture the contents of",
+        ),
+    ]
+    xsd_file: Annotated[
+        FilePath,
+        Field(
+            description="Path to the XSD file describing the XML documents to be imported. Including the schema will ensure that the appropriate relationships (1-to-1 vs 1-to-many) are created between parent and child elements, and that element text is stored consistently.",
+        ),
+    ]
+    max_recursion_depth: Annotated[
+        PositiveInt,
+        Field(
+            default=DEFAULT_MAX_RECURSION_DEPTH,
+            description="Maximum recursion depth for self-referential XSD types before serializing deeper content as JSON.",
+        ),
+    ]
+
+    @model_validator(mode="after")
+    def check_xsd_file(self) -> Self:
+        """Check whether the schema file is present and valid; load the XSD and generate the PyArrow schema."""
+        try:
+            self._xsd_schema = load_schema(self.xsd_file)
+            self._pyarrow_schema = generate_pyarrow_schema(
+                self._xsd_schema,
+                self.xml_tag,
+                max_recursion_depth=self.max_recursion_depth,
+            )
+        except Exception as e:
+            err_msg = f"Could not generate PyArrow schema from XSD: {e!s}"
+            raise ValueError(err_msg) from e
+
+        return self
+
+    @property
+    def xsd_schema(self) -> xmlschema.XMLSchema:
+        """Return the loaded XMLSchema instance."""
+        return self._xsd_schema
+
+    @property
+    def pyarrow_schema(self) -> pa.Schema:
+        """Return the generated PyArrow schema for xml_tag."""
+        return self._pyarrow_schema

src/cdm_data_loaders/readers/jsonschema_xsv/xsv_validator/helpers.py

@@ -4,7 +4,6 @@
 import subprocess
 from dataclasses import dataclass
 from functools import cached_property
-from logging import Logger, getLogger
 from pathlib import Path
 from shutil import copy, move
 from typing import Annotated, Final, Self
@@ -35,9 +34,6 @@
 VALID_SUFFIX: Final[str] = "-valid"
 
 
-logger: Logger = getLogger(__name__)
-
-
 class ErrorRecord(BaseModel):
     """A single error encountered while cleaning/validating a file.