mortie.arrow

Arrow interop: the morton_index pyarrow ExtensionType and the library-agnostic Arrow C Data Interface surface. MortonIndexType and MortonIndexExtArray are built lazily behind a module __getattr__ (pyarrow is optional), so they are documented narratively in Arrow interchange rather than here.

The morton_index Arrow skin: a pyarrow ExtensionType over the words.

A pyarrow :class:pyarrow.ExtensionType over uint64 storage carrying the morton_index tag (issue #35, phase 4; issue #58 flipped the storage to uint64).

This is the Arrow-interop sibling of the pandas ExtensionArray in :mod:mortie.morton_index. The packed 64-bit decimal-Morton words live in Rust (src_rust/src/decimal_morton.rs); this module only wraps them so the same words can travel through an Arrow array and survive a parquet round-trip with their morton_index identity attached as extension metadata. Storage is the raw uint64 words verbatim (over the kernel's bit layout), so the raw word order is the Z-order, the same convention as the pandas skin.

pyarrow is an optional dependency exactly like pandas: importing mortie succeeds with neither installed. The extension type is built lazily on first use and a clear ImportError is raised if it is touched without pyarrow.

morton_index_type()

Return the (registered) morton_index pyarrow extension type.

Returns:

Type	Description
ExtensionType	The singleton type instance, registered with pyarrow on first call.

Raises:

Type	Description
ImportError	If pyarrow is not installed.

Source code in mortie/arrow.py

def morton_index_type():
    """Return the (registered) ``morton_index`` pyarrow extension type.

    Returns
    -------
    pyarrow.ExtensionType
        The singleton type instance, registered with pyarrow on first call.

    Raises
    ------
    ImportError
        If pyarrow is not installed.
    """
    return _build_type()

from_morton_index(array)

Wrap a :class:~mortie.morton_index.MortonIndexArray as an Arrow array.

Builds a pyarrow ExtensionArray of the morton_index type over the same uint64 words. Missing elements -- a MortonIndexArray for which :meth:isna is True, i.e. the all-zero empty sentinel word -- emit Arrow nulls, so a null survives the round-trip back through :func:to_morton_index. (The missing mask is read off the uint64 words, so a sentinel word in a raw array is treated as a null too; an already-built Arrow array goes back through :func:to_morton_index, not here.)

Parameters:

Name	Type	Description	Default
array	MortonIndexArray or array_like	The words to wrap; may also be a raw uint64 array-like of words.	required

Returns:

Type	Description
ExtensionArray	A morton_index-typed Arrow array over the same words.

Raises:

Type	Description
ImportError	If pyarrow is not installed.

Source code in mortie/arrow.py

def from_morton_index(array):
    """Wrap a :class:`~mortie.morton_index.MortonIndexArray` as an Arrow array.

    Builds a pyarrow ``ExtensionArray`` of the ``morton_index`` type over the
    same ``uint64`` words. Missing elements -- a ``MortonIndexArray`` for which
    :meth:`isna` is True, i.e. the all-zero empty sentinel word -- emit Arrow
    nulls, so a null survives the round-trip back through
    :func:`to_morton_index`. (The missing mask is read off the ``uint64``
    words, so a sentinel word in a raw array is treated as a null too; an
    already-built Arrow array goes back through :func:`to_morton_index`, not
    here.)

    Parameters
    ----------
    array : MortonIndexArray or array_like
        The words to wrap; may also be a raw ``uint64`` array-like of words.

    Returns
    -------
    pyarrow.ExtensionArray
        A ``morton_index``-typed Arrow array over the same words.

    Raises
    ------
    ImportError
        If pyarrow is not installed.
    """
    pa = _require_pyarrow()
    ext_type = _build_type()
    data = np.asarray(getattr(array, "_data", array), dtype=np.uint64)
    # The empty sentinel (all-zero word, prefix 0) is the missing value on the
    # pandas side; mirror it as an Arrow null so isna() round-trips both ways.
    from .morton_index import MortonIndexArray

    mask = data == MortonIndexArray._SENTINEL
    storage = pa.array(data, type=pa.uint64(), mask=mask)
    return pa.ExtensionArray.from_storage(ext_type, storage)

to_morton_index(array)

Convert an Arrow morton_index array back to a MortonIndexArray.

Arrow nulls come back as the all-zero empty sentinel word, so the pandas :meth:isna reports them as missing.

Parameters:

Name	Type	Description	Default
array	ExtensionArray or Array	The extension array, or its plain uint64 storage.	required

Returns:

Type	Description
MortonIndexArray	The pandas-side :class:~mortie.morton_index.MortonIndexArray over the same words.

Raises:

Type	Description
ImportError	If pyarrow is not installed.

Source code in mortie/arrow.py

def to_morton_index(array):
    """Convert an Arrow ``morton_index`` array back to a ``MortonIndexArray``.

    Arrow nulls come back as the all-zero empty sentinel word, so the pandas
    :meth:`isna` reports them as missing.

    Parameters
    ----------
    array : pyarrow.ExtensionArray or pyarrow.Array
        The extension array, or its plain ``uint64`` storage.

    Returns
    -------
    MortonIndexArray
        The pandas-side :class:`~mortie.morton_index.MortonIndexArray` over
        the same words.

    Raises
    ------
    ImportError
        If pyarrow is not installed.
    """
    _require_pyarrow()
    from .morton_index import MortonIndexArray

    storage = getattr(array, "storage", array)
    # Fill nulls with the empty sentinel before materializing: a uint64 array
    # with a null buffer cannot go straight to numpy.
    if storage.null_count:
        storage = storage.fill_null(int(MortonIndexArray._SENTINEL))
    words = storage.to_numpy(zero_copy_only=False).astype(np.uint64, copy=False)
    return MortonIndexArray(words)

export_c_array(words)

Export packed uint64 words as an Arrow C Data Interface capsule pair.

Consumable by any Arrow lib without pandas or pyarrow.

Parameters:

Name	Type	Description	Default
words	array_like	Any uint64 array-like (e.g. a raw numpy array or a MortonIndexArray).	required

Returns:

Type	Description
tuple of PyCapsule	The (schema_capsule, array_capsule) pair, carrying the words as a morton_index extension column (ARROW:extension:name on the schema), with the all-zero empty sentinel mapped to an Arrow null via a real validity bitmap.

Source code in mortie/arrow.py

def export_c_array(words):
    """Export packed ``uint64`` words as an Arrow C Data Interface capsule pair.

    Consumable by any Arrow lib without pandas or pyarrow.

    Parameters
    ----------
    words : array_like
        Any ``uint64`` array-like (e.g. a raw numpy array or a
        ``MortonIndexArray``).

    Returns
    -------
    tuple of PyCapsule
        The ``(schema_capsule, array_capsule)`` pair, carrying the words as a
        ``morton_index`` extension column (``ARROW:extension:name`` on the
        schema), with the all-zero empty sentinel mapped to an Arrow null via
        a real validity bitmap.
    """
    from . import _rustie

    data = np.ascontiguousarray(
        np.asarray(getattr(words, "_data", words), dtype=np.uint64)
    )
    return _rustie.rust_mi_export_c_array(data)

export_c_schema()

Return the morton_index Arrow schema capsule.

The __arrow_c_schema__ half of the C Data Interface surface.

Returns:

Type	Description
PyCapsule	An ArrowSchema capsule carrying the morton_index extension type.

Source code in mortie/arrow.py

def export_c_schema():
    """Return the ``morton_index`` Arrow schema capsule.

    The ``__arrow_c_schema__`` half of the C Data Interface surface.

    Returns
    -------
    PyCapsule
        An ``ArrowSchema`` capsule carrying the ``morton_index`` extension
        type.
    """
    from . import _rustie

    return _rustie.rust_mi_export_c_schema()

import_c_array(source)

Import an Arrow C Data Interface array/stream as packed uint64 words.

Arrow nulls come back as the all-zero empty sentinel, so the null<->sentinel convention round-trips byte-for-byte. No pyarrow dependency on any path.

Parameters:

Name	Type	Description	Default
source	object or tuple	One of: an object exposing __arrow_c_array__ (a contiguous arro3-core / pyarrow / polars array), an object exposing __arrow_c_stream__ (a chunked column / multi-batch source -- every chunk is concatenated), or a (schema_capsule, array_capsule) tuple.	required

Returns:

Type	Description
ndarray	The packed words as a uint64 array.

Source code in mortie/arrow.py

def import_c_array(source):
    """Import an Arrow C Data Interface array/stream as packed ``uint64`` words.

    Arrow nulls come back as the all-zero empty sentinel, so the null<->sentinel
    convention round-trips byte-for-byte. No pyarrow dependency on any path.

    Parameters
    ----------
    source : object or tuple
        One of:

        * an object exposing ``__arrow_c_array__`` (a contiguous arro3-core /
          pyarrow / polars array),
        * an object exposing ``__arrow_c_stream__`` (a **chunked** column /
          multi-batch source -- every chunk is concatenated),
        * or a ``(schema_capsule, array_capsule)`` tuple.

    Returns
    -------
    numpy.ndarray
        The packed words as a ``uint64`` array.
    """
    from . import _rustie

    # A single contiguous array is preferred when both are present; only a
    # chunked source (no __arrow_c_array__) goes through the stream path.
    if hasattr(source, "__arrow_c_array__"):
        schema_capsule, array_capsule = source.__arrow_c_array__()
        return _rustie.rust_mi_import_c_array(schema_capsule, array_capsule)
    if hasattr(source, "__arrow_c_stream__"):
        return _rustie.rust_mi_import_c_stream(source.__arrow_c_stream__())
    schema_capsule, array_capsule = source
    return _rustie.rust_mi_import_c_array(schema_capsule, array_capsule)