Python decode function (#2563)

Author: gatesn

docs/api/python/arrays.rst

@@ -85,7 +85,7 @@ functionality.
 
 .. autoclass:: vortex.Float64TypeArray
      :members:
-     
+
 .. autoclass:: vortex.Utf8TypeArray
      :members:
 
@@ -201,5 +201,23 @@ Compressed Encodings
 Pluggable Encodings
 -------------------
 
-.. autoclass:: vortex.PyEncoding
+Subclasses of :class:`~vortex.PyArray` can be used to implement custom Vortex encodings in Python. These encodings
+can be registered with the :attr:`~vortex.registry` so they are available to use when reading Vortex files.
+
+.. autoclass:: vortex.PyArray
+    :members:
+
+
+Registry and Serde
+------------------
+
+.. autodata:: vortex.registry
+
+.. autoclass:: vortex.Registry
+    :members:
+
+.. autoclass:: vortex.ArrayContext
+    :members:
+
+.. autoclass:: vortex.ArrayParts
     :members:

pyvortex/pyproject.toml

@@ -40,5 +40,6 @@ dev = [
     "mypy>=1.14.1",
     "numpy>=2.2.2",
     "pandas-stubs>=2.2.3.241126",
+    "pcodec>=0.3.3",
     "pyarrow-stubs>=17.16",
 ]

pyvortex/python/vortex/__init__.py

@@ -1,5 +1,4 @@
 from . import _lib
-from ._lib import register
 from ._lib.arrays import (
     AlpArray,
     AlpRdArray,
@@ -32,7 +31,6 @@
     NullTypeArray,
     PrimitiveArray,
     PrimitiveTypeArray,
-    PyEncoding,
     RunEndArray,
     SparseArray,
     StructArray,
@@ -68,6 +66,7 @@
     uint,
     utf8,
 )
+from ._lib.registry import Registry
 from ._lib.scalar import (
     BinaryScalar,
     BoolScalar,
@@ -80,16 +79,17 @@
     Utf8Scalar,
     scalar,
 )
-from .arrays import Array, array
+from ._lib.serde import ArrayContext, ArrayParts
+from .arrays import Array, PyArray, array
 
 assert _lib, "Ensure we eagerly import the Vortex native library"
 
 __all__ = [
     "array",
     "compress",
-    "register",
     # Arrays
     "Array",
+    "PyArray",
     "NullTypeArray",
     "BoolTypeArray",
     "PrimitiveTypeArray",
@@ -143,7 +143,6 @@
     "StructArray",
     "ListArray",
     "ExtensionArray",
-    "PyEncoding",
     "AlpArray",
     "AlpRdArray",
     "DateTimePartsArray",
@@ -166,4 +165,11 @@
     "StructScalar",
     "ListScalar",
     "ExtensionScalar",
+    # Registry + Serde
+    "Registry",
+    "ArrayContext",
+    "ArrayParts",
 ]
+
+#: The default registry for Vortex
+registry = Registry()

pyvortex/python/vortex/_lib/__init__.pyi

@@ -1,3 +0,0 @@
-import vortex as vx
-
-def register(cls: type[vx.PyEncoding]): ...

pyvortex/python/vortex/_lib/arrays.pyi

@@ -17,7 +17,7 @@ class Array:
     def from_arrow(obj: pa.Array | pa.ChunkedArray | pa.Table) -> Array: ...
     def to_arrow_array(self) -> pa.Array: ...
     @property
-    def encoding(self) -> str: ...
+    def id(self) -> str: ...
     @property
     def nbytes(self) -> int: ...
     @property
@@ -34,6 +34,9 @@ class Array:
     def to_polars_dataframe(self) -> pl.DataFrame: ...
     def to_polars_series(self) -> pl.Series: ...
     def to_pylist(self) -> list: ...
+    def serialize(self, ctx: vx.ArrayContext) -> bytes: ...
+
+class NativeArray(vx.Array): ...
 
 @final
 class NullArray(vx.Array): ...
@@ -60,8 +63,6 @@ class ListArray(vx.Array): ...
 @final
 class ExtensionArray(vx.Array): ...
 
-class PyEncoding(vx.Array): ...
-
 @final
 class ConstantArray(vx.Array):
     def scalar(self) -> vx.Scalar: ...

pyvortex/python/vortex/_lib/mypy.allowlist

@@ -8,4 +8,6 @@ vortex\._lib\.dataset
 vortex\._lib\.dtype
 vortex\._lib\.expr
 vortex\._lib\.io
+vortex\._lib\.registry
 vortex\._lib\.scalar
+vortex\._lib\.serde

pyvortex/python/vortex/_lib/registry.pyi

@@ -0,0 +1,8 @@
+from typing import final
+
+import vortex as vx
+
+@final
+class Registry:
+    def register(self, cls: type[vx.Array]): ...
+    def array_ctx(self, encodings: list[type[vx.Array] | str]) -> vx.ArrayContext: ...

pyvortex/python/vortex/_lib/serde.pyi

@@ -0,0 +1,25 @@
+from typing import final
+
+import pyarrow as pa
+
+import vortex as vx
+
+@final
+class ArrayParts:
+    @staticmethod
+    def parse(data: bytes) -> ArrayParts: ...
+    @property
+    def metadata(self) -> bytes | None: ...
+    @property
+    def nbuffers(self) -> int: ...
+    @property
+    def buffers(self) -> list[pa.Buffer]: ...
+    @property
+    def nchildren(self) -> int: ...
+    @property
+    def children(self) -> list[ArrayParts]: ...
+    def decode(self, ctx: ArrayContext, dtype: vx.DType, len: int) -> pa.Array: ...
+
+@final
+class ArrayContext:
+    def __len__(self): ...

pyvortex/python/vortex/arrays.py

@@ -1,7 +1,9 @@
+import abc
 from typing import TYPE_CHECKING, Any
 
 import pyarrow
 
+import vortex as vx
 from vortex._lib import arrays as _arrays
 
 try:
@@ -380,3 +382,23 @@ def array(obj: pyarrow.Array | list | Any) -> Array:
     except ImportError:
         pass
     return Array.from_arrow(obj)
+
+
+class PyArray(Array, metaclass=abc.ABCMeta):
+    """Abstract base class for Python-based Vortex arrays."""
+
+    id: str
+
+    @abc.abstractmethod
+    def __len__(self) -> int:
+        """Return the logical length of the array."""
+
+    @classmethod
+    @abc.abstractmethod
+    def decode(cls, parts: vx.ArrayParts, ctx: vx.ArrayContext, dtype: vx.DType, len: int) -> Array:
+        """Decode an array from its component parts.
+
+        :class:`ArrayParts` contains the metadata, buffers and child :class:`ArrayParts` that represent the
+        current array. Implementations of this function should validate this information, and then construct
+        a new array.
+        """

pyvortex/python/vortex/serde.py

@@ -0,0 +1,3 @@
+from vortex._lib.serde import ArrayContext, ArrayParts
+
+__all__ = ["ArrayParts", "ArrayContext"]