Repository: Point72/ccflow

Issue:
Allow Validation of ArrowSchema External Type
**Is your feature request related to a problem? Please describe.**
Pydantic cannot validate the schema of the `ArrowSchema` ext type, requiring that any model containing them to allow arbitrary types.

**Describe the solution you'd like**
Allow validation of the `ArrowSchema` type

**Describe alternatives you've considered**
While models can serialize the representation of a pa.Schema in other ways, such as a dict of field names/types, it seems reasonable to allow validation of the `ArrowSchema` for this purpose to provide a general solution.


Discussion:

Patch:
diff --git a/ccflow/exttypes/arrow.py b/ccflow/exttypes/arrow.py
index dbed26c..ebdf79b 100644
--- a/ccflow/exttypes/arrow.py
+++ b/ccflow/exttypes/arrow.py
@@ -4,7 +4,7 @@ import narwhals.stable.v1 as nw
 import pyarrow as pa
 from pydantic import TypeAdapter
 from pydantic_core import core_schema
-from typing_extensions import Literal, Self, get_args
+from typing_extensions import Any, Literal, Self, get_args
 
 __all__ = ("ArrowSchema", "ArrowTable", "PyArrowDatatype")
 
@@ -39,6 +39,25 @@ class ArrowSchema(type):
 
         return newclass
 
+    @classmethod
+    def __get_pydantic_core_schema__(cls, source_type, handler):
+        return core_schema.no_info_plain_validator_function(cls._validate)
+
+    @classmethod
+    def _validate(cls, value: Any) -> Self:
+        if isinstance(value, ArrowSchema):
+            return value
+
+        if isinstance(value, pa.Schema):
+            return ArrowSchema.make(value)
+
+        raise ValueError(f"Cannot convert value to ArrowSchema, expects ArrowSchema or pa.Schema: {value}")
+
+    @classmethod
+    def validate(cls, value: Any) -> Self:
+        """Try to convert/validate an arbitrary value to an ArrowSchema."""
+        return _TYPE_ADAPTER_ARROW_SCHEMA.validate_python(value)
+
 
 S = TypeVar("S", bound=ArrowSchema)
 
@@ -120,7 +139,8 @@ class PyArrowDatatype(str):
     @classmethod
     def validate(cls, value) -> Self:
         """Try to convert/validate an arbitrary value to a PyArrowDatatype."""
-        return _TYPE_ADAPTER.validate_python(value)
+        return _TYPE_ADAPTER_PYARROW_DATA_TYPE.validate_python(value)
 
 
-_TYPE_ADAPTER = TypeAdapter(PyArrowDatatype)
+_TYPE_ADAPTER_PYARROW_DATA_TYPE = TypeAdapter(PyArrowDatatype)
+_TYPE_ADAPTER_ARROW_SCHEMA = TypeAdapter(ArrowSchema)


Test patch:
diff --git a/ccflow/tests/exttypes/test_arrow.py b/ccflow/tests/exttypes/test_arrow.py
index ffafba8..2e82d7b 100644
--- a/ccflow/tests/exttypes/test_arrow.py
+++ b/ccflow/tests/exttypes/test_arrow.py
@@ -40,6 +40,24 @@ class TestArrowSchema(TestCase):
         # You can't actually construct the schema type (because the schema is on the type, not the instance)
         self.assertRaises(TypeError, SCHEMA_WEAK)
 
+    def test_validate(self):
+        # Validation failures should raise
+        self.assertRaises(ValueError, ArrowSchema.validate, "foo")
+        self.assertRaises(ValueError, ArrowSchema.validate, 20.0)
+        self.assertRaises(ValueError, ArrowSchema.validate, False)
+
+        # Validation of a pyarrow schema should provide an ArrowSchema
+        pyarrow_schema = pa.schema({"A": pa.float64(), "B": pa.int64(), "C": pa.string()})
+        result = ArrowSchema.validate(pyarrow_schema)
+        self.assertIsInstance(result, ArrowSchema)
+        self.assertTrue(pyarrow_schema.equals(result.schema))
+
+        # Validation of an ArrowSchema should provide an ArrowSchema
+        schema = ArrowSchema.make(pyarrow_schema, strict=True)
+        result = ArrowSchema.validate(schema)
+        self.assertIsInstance(result, ArrowSchema)
+        self.assertEqual(schema, result)
+
 
 class TestArrowTable(TestCase):
     def setUp(self) -> None: