feat: add bidirectional arrow conversion

src/schematools/arrow.py

@@ -0,0 +1,229 @@
+import typing as t
+from functools import singledispatchmethod
+
+import pyarrow as pa
+
+from schematools import jsonschema
+from schematools.jsonschema import (
+    ArrayType,
+    BaseJSONSchemaType,
+    BooleanType,
+    IntegerType,
+    JSONSchemaParser,
+    NullType,
+    NumberType,
+    ObjectType,
+    StringType,
+    UnionType,
+)
+
+from .convert import SchemaConverterBase
+
+
+class JSONSchemaToArrowTypeMap:
+
+    @singledispatchmethod
+    def convert(self, jsontype: BaseJSONSchemaType) -> pa.DataType:
+        """Convert JSON type to Apache Arrow type."""
+        raise NotImplementedError(f"Conversion of {jsontype} is not supported")
+
+    @convert.register
+    def convert_union(self, jsontype: UnionType) -> pa.DataType:
+        """Convert UnionType to Apache Arrow type."""
+        return pa.union(
+            [pa.field(i, self.convert(t)) for i, t in enumerate(jsontype.types)],
+            mode="dense",
+        )
+
+    @convert.register
+    def convert_integer(self, jsontype: IntegerType) -> pa.DataType:
+        """Convert IntegerType to Apache Arrow type."""
+        return pa.int64()
+
+    @convert.register
+    def convert_number(self, jsontype: NumberType) -> pa.DataType:
+        """Convert NumberType to Apache Arrow type."""
+        return pa.float64()
+
+    @convert.register
+    def convert_string(self, jsontype: StringType) -> pa.DataType:
+        """Convert StringType to Apache Arrow type."""
+        return pa.string()
+
+    @convert.register
+    def convert_boolean(self, jsontype: BooleanType) -> pa.DataType:
+        """Convert BooleanType to Apache Arrow type."""
+        return pa.bool_()
+
+    @convert.register
+    def convert_object(self, jsontype: ObjectType) -> pa.DataType:
+        """Convert ObjectType to Apache Arrow type."""
+        if jsontype.properties is not None:
+
+            fields = []
+            for k, v in jsontype.properties.items():
+                type_ = self.convert(v)
+                if isinstance(type_, UnionType):
+                    if type_.is_simple_nullable():
+                        fields.append(
+                            pa.field(
+                                name=k,
+                                type=type_.simple_nullable_type,
+                                nullable=True,
+                            )
+                        )
+
+                pa.field(
+                    name=k,
+                    type=type_,
+                    nullable=type_.is_nullable(),
+                )
+
+            return pa.struct(fields), False
+        return pa.struct([]), False
+
+    @convert.register
+    def convert_array(self, jsontype: ArrayType) -> pa.DataType:
+        """Convert ArrayType to Apache Arrow type."""
+        return pa.list_(self.convert(jsontype.items))
+
+    @convert.register
+    def convert_null(self, jsontype: NullType) -> pa.DataType:
+        """Convert NullType to Apache Arrow type."""
+        return pa.null()
+
+
+class ArrowToJSONSchemaTypeMap:
+
+    @singledispatchmethod
+    def convert(
+        self, arrowtype: pa.DataType, nullable: bool = False
+    ) -> BaseJSONSchemaType:
+        """Convert Apache Arrow type to JSON type.
+
+        PyArrow does not provide importable Type classes for all types, so we have to use
+        the `is_*` methods from pa.types for simple types.
+        """
+        if pa.types.is_string(arrowtype):
+            if nullable:
+                return UnionType.from_types([StringType, NullType])
+            return StringType()
+
+        if pa.types.is_integer(arrowtype):
+            if nullable:
+                return UnionType.from_types([IntegerType, NullType])
+            return IntegerType()
+
+        if pa.types.is_floating(arrowtype):
+            if nullable:
+                return UnionType.from_types([NumberType, NullType])
+            return NumberType()
+
+        if pa.types.is_boolean(arrowtype):
+            if nullable:
+                return UnionType.from_types([BooleanType, NullType])
+            return BooleanType()
+
+        if pa.types.is_null(arrowtype):
+            return NullType()
+
+        raise NotImplementedError(f"Conversion of {arrowtype} is not supported")
+
+    @convert.register
+    def convert_union(
+        self, arrowtype: pa.UnionType, nullable: bool = False
+    ) -> BaseJSONSchemaType:
+        """Convert Apache Arrow type to JSON type."""
+        fields = [arrowtype.field(i) for i in range(arrowtype.num_fields)]
+        field_types = [field.type for field in fields]
+        jsonschema_types = [
+            type(self.convert(field_type)) for field_type in field_types
+        ]
+        if nullable:
+            jsonschema_types.append(NullType)
+        return UnionType.from_types(jsonschema_types)
+
+    @convert.register
+    def convert_array(
+        self, arrowtype: pa.ListType | pa.LargeListType, nullable: bool = False
+    ) -> ArrayType:
+        """Convert Apache Arrow type to ArrayType."""
+        if nullable:
+            nullable_array = UnionType.from_types([ArrayType, NullType])
+            nullable_array.items = self.convert(arrowtype.value_type)
+            return nullable_array
+        return ArrayType(items=self.convert(arrowtype.value_type))
+
+    @convert.register
+    def convert_struct(
+        self, arrowtype: pa.StructType, nullable: bool = False
+    ) -> ObjectType:
+        """Convert Apache Arrow type to ObjectType."""
+        fields = [field for field in arrowtype]
+        properties = {
+            field.name: self.convert(field.type, nullable=field.nullable)
+            for field in fields
+        }
+        if nullable:
+            return UnionType.from_types([ObjectType(properties=properties), NullType])
+        return ObjectType(properties=properties)
+
+
+class ArrowToJSONSchemaConverter(SchemaConverterBase):
+    """Apache Arrow schema representation."""
+
+    @staticmethod
+    def from_jsonschema(jsonschema: BaseJSONSchemaType) -> pa.Schema:
+        """Convert JSON schema to Apache Arrow schema."""
+
+        if isinstance(jsonschema, ObjectType) and jsonschema.has_properties():
+            return pa.schema(
+                [
+                    pa.field(
+                        name=k,
+                        type=JSONSchemaToArrowTypeMap().convert(v),
+                        nullable=jsonschema.is_nullable(),
+                    )
+                    for k, v in jsonschema.properties.items()
+                ]
+            )
+
+        if isinstance(jsonschema, UnionType):
+            if jsonschema.is_simple_nullable():
+                return pa.schema(
+                    [
+                        pa.field(
+                            name="root",
+                            type=JSONSchemaToArrowTypeMap().convert(jsonschema),
+                            nullable=True,
+                        )
+                    ]
+                )
+
+        return pa.schema(
+            [
+                pa.field(
+                    name="root",
+                    type=JSONSchemaToArrowTypeMap().convert(jsonschema),
+                    nullable=jsonschema.is_nullable(),
+                )
+            ]
+        )
+
+    @staticmethod
+    def to_jsonschema(arrow_schema: pa.Schema) -> BaseJSONSchemaType:
+        """Convert Apache Arrow schema to JSON schema."""
+        if len(arrow_schema) == 1:
+            if isinstance(arrow_schema[0], pa.DataType):
+                return ArrowToJSONSchemaTypeMap().convert(arrow_schema[0].type)
+            if isinstance(arrow_schema[0], pa.Field):
+                return ArrowToJSONSchemaTypeMap().convert(
+                    arrow_schema[0].type, nullable=arrow_schema[0].nullable
+                )
+        properties = {
+            field.name: ArrowToJSONSchemaTypeMap().convert(
+                field.type, nullable=field.nullable
+            )
+            for field in arrow_schema
+        }
+        return ObjectType(properties=properties)

src/schematools/convert.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+import typing as t
+
+from schematools.jsonschema import BaseJSONSchemaType
+
+
+class SchemaConverterBase:
+
+    @staticmethod
+    def to_jsonschema(schema: t.Any) -> BaseJSONSchemaType:
+        raise NotImplementedError
+
+    @staticmethod
+    def from_jsonschema(jsontype: BaseJSONSchemaType) -> t.Any:
+        raise NotImplementedError

src/schematools/jsonschema/__init__.py

@@ -4,7 +4,7 @@
 from .parse import JSONSchemaParser
 from .types import (
     ArrayType,
-    BaseJSONType,
+    BaseJSONSchemaType,
     BooleanType,
     DateTimeType,
     DateType,
@@ -22,6 +22,7 @@
     RelativeJSONPointerType,
     StringType,
     TimeType,
+    UnionType,
     URIReferenceType,
     URITemplateType,
     URIType,
@@ -30,7 +31,7 @@
 
 __all__ = [
     "ArrayType",
-    "BaseJSONType",
+    "BaseJSONSchemaType",
     "BooleanType",
     "DateTimeType",
     "DateType",
@@ -43,16 +44,17 @@
     "IPv6Type",
     "JSONPointerType",
     "JSONSchema",
+    "JSONSchemaParser",
     "NullType",
     "NumberType",
     "ObjectType",
     "RegexType",
     "RelativeJSONPointerType",
     "StringType",
     "TimeType",
+    "UnionType",
     "URIReferenceType",
     "URITemplateType",
     "URIType",
     "UUIDType",
-    "JSONSchemaParser",
 ]