cdm_data_loaders.converters.emitters.pyarrow ¶

class PyArrowEmitterError(ConversionError):
    """A typed schema cannot be represented by the configured Arrow policy."""

class PyArrowEmitter(_TabularEmitterBase[pa.DataType]):
    """Convert typed schema nodes into PyArrow schemas and empty tables."""

    _error_cls: ClassVar[type[ConversionError]] = PyArrowEmitterError
    _target_name: ClassVar[str] = "Arrow"
    _decimal_target_label: ClassVar[str] = "Arrow"
    _string_type_label: ClassVar[str] = "string"
    _struct_type_label: ClassVar[str] = "an Arrow struct"
    _empty_struct_label: ClassVar[str] = "empty struct"
    _map_type_label: ClassVar[str] = "an Arrow map"
    _map_value_label: ClassVar[str] = "map value type"
    _union_approx_reason: ClassVar[str] = "Arrow has no union type compatible with table storage."
    _array_label: ClassVar[str] = "Arrow lists"
    _slot_label: ClassVar[str] = "list's"

    format_map: frozendict[str, pa.DataType] = Field(default_factory=lambda: DEFAULT_FORMAT_MAP)

    @field_validator("format_map", mode="before")
    @classmethod
    def _merge_with_builtin_format_map(cls, value: object) -> object:
        """Merge nonempty format overrides into the built-in mapping."""
        return merge_format_map(value) if value else value

    def emit(self, document: SchemaDocument) -> pa.Schema:
        """Emit a table schema, requiring the root to be a struct.

        :param document: the schema document to convert
        :type document: SchemaDocument
        :return: the Arrow schema
        :rtype: pa.Schema
        :raises PyArrowEmitterError: if the root does not resolve to a struct
        """
        result = self._emit_root_node(document)
        if pa.types.is_struct(result):
            return pa.schema([result.field(index) for index in range(result.num_fields)])
        msg = (
            "Root JSON Schema did not resolve to a struct.\nThis can happen if it has no "
            "'properties' but does have 'patternProperties' or 'additionalProperties', which map to "
            "a map type instead. An Arrow table requires a struct at the top level of its schema."
        )
        raise PyArrowEmitterError(msg)

    def emit_table(self, document: SchemaDocument) -> pa.Table:
        """Emit an empty table with the document's schema.

        :param document: the schema document to convert
        :type document: SchemaDocument
        :return: a zero-row table
        :rtype: pa.Table
        :raises PyArrowEmitterError: if the root does not resolve to a struct
        """
        return self.emit(document).empty_table()

    def emit_python(
        self,
        document: SchemaDocument | Mapping[str, SchemaDocument],
        *,
        docstring: str | None = None,
    ) -> str:
        """Render a static Python module defining PyArrow struct components and schemas.

        :param document: A single ``SchemaDocument`` or a mapping of table names to ``SchemaDocument``s.
        :type document: SchemaDocument | Mapping[str, SchemaDocument]
        :param docstring: Optional module docstring header.
        :type docstring: str | None
        :return: Formatted Python source code defining ``*_STRUCT``, ``*_SCHEMA``, ``SCHEMA``, and ``SCHEMAS``.
        :rtype: str
        :raises PyArrowEmitterError: If any document root does not resolve to a struct.
        """
        return render_python_schema_module(self, document, docstring=docstring)

    def emit_python_file(
        self,
        document: SchemaDocument | Mapping[str, SchemaDocument],
        output: str | Path,
        *,
        docstring: str | None = None,
    ) -> Path:
        """Render and write a static Python schema module to ``output``.

        :param document: A single ``SchemaDocument`` or a mapping of table names to ``SchemaDocument``s.
        :type document: SchemaDocument | Mapping[str, SchemaDocument]
        :param output: Destination file path for the generated Python module.
        :type output: str | Path
        :param docstring: Optional module docstring header.
        :type docstring: str | None
        :return: Resolved output ``Path``.
        :rtype: Path
        :raises PyArrowEmitterError: If any document root does not resolve to a struct.
        """
        return write_python_schema_file(self, document, output, docstring=docstring)

    def build_metadata(
        self, node: TypedNode, ctx: ConversionContext, field: SchemaField | None = None
    ) -> dict[str, str]:
        """Select field metadata and encode it as Arrow's string key/value pairs.

        :param node: the node whose constraints, annotations and extensions are searched
        :type node: TypedNode
        :param ctx: draft-specific metadata context
        :type ctx: ConversionContext
        :param field: the owning field, whose annotations and extensions take precedence
        :type field: SchemaField | None
        :return: ``jsonschema`` (JSON object), ``comment`` and ``original_union`` (JSON array) entries
        :rtype: dict[str, str]
        """
        values, jsonschema, original_union = self._select_metadata_parts(node, ctx, field)
        metadata: dict[str, str] = {}
        if jsonschema:
            metadata["jsonschema"] = _dump_json(jsonschema)
        description = values.get("description")
        if isinstance(description, str):
            metadata["comment"] = description
        for prov in (field.provenance if field is not None else None, node.provenance):
            if prov is not None:
                for key in ("jsonschema_type", "pydantic_type", "xsd_type"):
                    val = prov.metadata.get(key)
                    if isinstance(val, str):
                        metadata.setdefault(key, val)
        if original_union is not None:
            metadata["original_union"] = _dump_json(original_union)
        return metadata

    def _string_type(self) -> pa.DataType:
        return pa.string()

    def _int32_type(self) -> pa.DataType:
        return pa.int32()

    def _int64_type(self) -> pa.DataType:
        return pa.int64()

    def _float64_type(self) -> pa.DataType:
        return pa.float64()

    def _decimal128_type(self, precision: int, scale: int) -> pa.DataType:
        return pa.decimal128(precision, scale)

    def _scalar_types(self) -> Mapping[str, pa.DataType]:
        return SCALAR_TYPE_MAP

    def _infer_enum_type(self, values: list[object]) -> pa.DataType:
        return infer_type_from_enum(values)

    def _struct_type(self, properties: tuple[SchemaField, ...], ctx: ConversionContext, depth: int) -> pa.DataType:
        return pa.struct(
            [
                pa.field(
                    prop.name,
                    self.emit_node(prop.node, ctx, depth + 1),
                    nullable=not prop.required,
                    metadata=self.build_metadata(prop.node, ctx, prop) or None,
                )
                for prop in properties
            ]
        )

    def _map_type(self, key_type: pa.DataType, value_type: pa.DataType, *, value_nullable: bool) -> pa.MapType:
        return pa.map_(key_type, pa.field("value", value_type, nullable=value_nullable))

    def _array_type(self, element_type: pa.DataType, *, element_nullable: bool) -> pa.ListType:
        return pa.list_(pa.field("item", element_type, nullable=element_nullable))

    def _emit_map(self, node: TypedNode, ctx: ConversionContext, depth: int = 0) -> pa.MapType:
        """Emit a map with independently typed keys, values, and value field metadata."""
        if node.key_type is None or node.value_type is None:
            msg = "Map nodes require both key_type and value_type"
            raise self._error_cls(msg)
        return pa.map_(
            self.emit_node(node.key_type, ctx, depth + 1),
            pa.field(
                "value",
                self.emit_node(node.value_type, ctx, depth + 1),
                nullable=element_nullable(node.value_type),
                metadata=self.build_metadata(node.value_type, ctx) or None,
            ),
        )

    def _emit_array(self, node: TypedNode, ctx: ConversionContext, depth: int = 0) -> pa.ListType:
        """Emit homogeneous arrays, preserving element field nullability and metadata."""
        element = self._resolve_array_element(node)
        if element is None:
            return self._array_type(self._string_type(), element_nullable=True)
        return pa.list_(
            pa.field(
                "item",
                self.emit_node(element, ctx, depth + 1),
                nullable=element_nullable(element),
                metadata=self.build_metadata(element, ctx) or None,
            )
        )

    @staticmethod
    def _emit_logical(node: TypedNode) -> pa.DataType | None:
        """Preserve physical widths and logical source types without a JSON encoding."""
        if node.provenance is not None and node.provenance.source == "json-schema":
            logical_key = node.hints.logical_type
            return (
                JSON_SCHEMA_LOGICAL_TYPE_MAP.get(logical_key)
                if logical_key is not None and "jsonschema_type" in node.provenance.metadata
                else None
            )
        hints = node.hints
        logical = hints.logical_type
        if logical in {"decimal", "wei"}:
            return _decimal_type(hints, wei=logical == "wei")
        if logical in {"timestamp", "timestamp-with-tz", "timestamp-without-tz"}:
            unit = _timestamp_unit(hints.precision)
            naive = hints.timezone is False or logical == "timestamp-without-tz"
            return pa.timestamp(unit) if naive else pa.timestamp(unit, tz="UTC")
        if node.type == "integer":
            return _emit_logical_integer(hints)
        if node.type == "number" and (logical == "float" or hints.bit_width == INT32_WIDTH):
            return pa.float32()
        fixed = pa.binary(hints.length) if hints.length is not None else pa.binary()
        return {"date": pa.date32(), "time": pa.time64("us"), "binary": pa.binary(), "fixed": fixed}.get(logical)