maxframe 2.0.0b2__cp37-cp37m-win32.whl → 2.3.0rc1__cp37-cp37m-win32.whl

maxframe/lib/dtypes_extension/blob.py

@@ -0,0 +1,304 @@
+# Copyright 1999-2025 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import abc
+import contextlib
+import io
+from typing import Any, ContextManager, List, Optional, Sequence, Union
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+from pandas.api.extensions import (
+    ExtensionArray,
+    ExtensionDtype,
+    register_extension_dtype,
+)
+
+from ...utils import tokenize
+from .dtypes import ArrowDtype
+
+
+class ArrowBlobType(pa.ExtensionType):
+    def __init__(self):
+        super().__init__(pa.binary(), "maxframe.blob")
+
+    def __arrow_ext_serialize__(self):
+        return b""
+
+    @classmethod
+    def __arrow_ext_deserialize__(cls, storage_type, serialized):
+        return ArrowBlobType()
+
+    def __eq__(self, other):
+        return isinstance(other, ArrowBlobType)
+
+    def __hash__(self):
+        return hash(str(self))
+
+    def to_pandas_dtype(self):
+        return ExternalBlobDtype()
+
+
+class AbstractExternalBlob(metaclass=abc.ABCMeta):
+    _blob_types = {}
+
+    @property
+    def reference(self) -> Any:
+        raise NotImplementedError
+
+    def __eq__(self, other) -> bool:
+        if isinstance(other, AbstractExternalBlob):
+            return self.reference == other.reference
+        else:
+            return False
+
+    def __hash__(self):
+        return hash((type(self), self.reference))
+
+    def __repr__(self):
+        return f"<{type(self).__name__} reference={self.reference}>"
+
+    @abc.abstractmethod
+    def open(self, mode: str = "r") -> ContextManager[io.IOBase]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def copy(self):
+        raise NotImplementedError
+
+    def __copy__(self):
+        return self.copy()
+
+    @classmethod
+    def register(cls):
+        cls._blob_types[cls.__name__.lower()] = cls
+
+    @classmethod
+    def get_cls_by_name(cls, cls_name: str):
+        return cls._blob_types[cls_name.lower()]
+
+
+class SolidBlob(AbstractExternalBlob):
+    def __init__(self, content: Optional[bytes] = None):
+        self._reference = content
+
+    @property
+    def reference(self) -> bytes:
+        return self._reference
+
+    def __reduce__(self):
+        return type(self), (self._reference,)
+
+    @contextlib.contextmanager
+    def open(self, mode: str = "r"):
+        if "w" in mode:
+            sio = io.BytesIO()
+        else:
+            sio = io.BytesIO(self._reference)
+        try:
+            yield sio
+        except Exception:
+            raise
+        else:
+            if "w" in mode:
+                self._reference = sio.getvalue()
+
+    def copy(self) -> "SolidBlob":
+        return SolidBlob(self._reference)
+
+    def __maxframe_tokenize__(self):
+        return type(self), self._reference
+
+
+SolidBlob.register()
+
+
+def new_blob(content: Optional[bytes] = None):
+    # todo this function currently maps directly to SolidBlob. It should
+    #  provide appropriate instance base on running envs.
+    return SolidBlob(content)
+
+
+def blob_from_reference(reference: Optional[bytes] = None):
+    # todo this function currently maps directly to SolidBlob. It should
+    #  provide appropriate instance base on running envs.
+    return SolidBlob(reference)
+
+
+@register_extension_dtype
+class ExternalBlobDtype(ExtensionDtype):
+    """
+    Extension dtype for ExternalBlob data.
+    """
+
+    name = "blob"
+
+    @classmethod
+    def construct_array_type(cls):
+        """
+        Return the array type associated with this dtype.
+        """
+        return ExternalBlobExtensionArray
+
+    @classmethod
+    def construct_from_string(cls, string):
+        if string == cls.name:
+            return cls()
+        raise TypeError(f"Cannot construct a '{cls.__name__}' from '{string}'")
+
+    @property
+    def type(self):
+        return AbstractExternalBlob
+
+    def __repr__(self) -> str:
+        return self.name
+
+    def __from_arrow__(self, array):
+        return ExternalBlobExtensionArray(
+            [blob_from_reference(a) for a in array.to_pylist()]
+        )
+
+
+class ExternalBlobExtensionArray(ExtensionArray):
+    """
+    Extension array for storing AbstractExternalBlob objects.
+    """
+
+    _data: List[AbstractExternalBlob]
+
+    def __init__(self, values: Union[List[AbstractExternalBlob], np.ndarray]):
+        if not isinstance(values, (list, np.ndarray)):
+            raise TypeError("values must be a list or numpy array")
+
+        if isinstance(values, np.ndarray):
+            if values.ndim != 1:
+                raise ValueError("values must be a 1-dimensional array")
+            values = values.tolist()
+
+        new_values = [None] * len(values)
+        for idx, val in enumerate(values):
+            if isinstance(val, str):
+                val = val.encode()
+            if isinstance(val, (bytes, bytearray)):
+                val = new_blob(val)
+
+            if pd.isna(val):
+                val = None
+            if val is not None and not isinstance(val, AbstractExternalBlob):
+                raise TypeError(
+                    "All values must be instances of AbstractExternalBlob or None"
+                )
+            else:
+                new_values[idx] = val
+
+        # Store with python list instead of numpy array
+        self._data = new_values
+        self._dtype = ExternalBlobDtype()
+
+    @classmethod
+    def _from_sequence(cls, scalars, dtype=None, copy: bool = False):
+        """
+        Construct a new ExtensionArray from a sequence of scalars.
+        """
+        return cls(scalars)
+
+    @classmethod
+    def _from_factorized(cls, values, original):
+        """
+        Reconstruct an ExtensionArray after factorization.
+        """
+        return cls(values)
+
+    @property
+    def dtype(self) -> ExtensionDtype:
+        """Return the dtype of the array."""
+        return self._dtype
+
+    def __len__(self) -> int:
+        """Return length of the array."""
+        return len(self._data)
+
+    def __getitem__(self, item):
+        """Select a subset of self."""
+        if isinstance(item, int):
+            return self._data[item]
+        else:
+            # For slices and other index types, return a new array with the selected items
+            return type(self)([self._data[i] for i in np.arange(len(self._data))[item]])
+
+    def __maxframe_tokenize__(self):
+        return type(self), tokenize(*self._data)
+
+    def isna(self):
+        """A 1-D array indicating if each value is missing."""
+        return np.array([x is None for x in self._data])
+
+    def copy(self):
+        """Return a copy of the array."""
+        return type(self)(self._data.copy())
+
+    @property
+    def nbytes(self) -> int:
+        """The number of bytes needed to store this object in memory."""
+        # Since we're storing references, we'll return a minimal size estimate
+        return len(self._data) * 8
+
+    def take(self, indices, allow_fill=False, fill_value=None):
+        """
+        Take elements from an array.
+
+        Parameters
+        ----------
+        indices : sequence of int
+            Indices to be taken.
+        allow_fill : bool, default False
+            How to handle negative values in `indices`.
+        fill_value : any, optional
+            Fill value to use for NA-indices when `allow_fill` is True.
+        """
+        if allow_fill:
+            # Handle negative indices
+            values = []
+            for i in indices:
+                if i < 0:
+                    if not allow_fill:
+                        raise IndexError("negative index not allowed")
+                    values.append(fill_value)
+                else:
+                    values.append(self._data[i])
+            return type(self)(values)
+        else:
+            return type(self)([self._data[i] for i in indices])
+
+    @classmethod
+    def _concat_same_type(
+        cls, to_concat: Sequence["ExternalBlobExtensionArray"]
+    ) -> "ExternalBlobExtensionArray":
+        return cls([x for sublist in to_concat for x in sublist._data])
+
+    def __arrow_array__(self, type=None):
+        return pa.ExtensionArray.from_storage(
+            ArrowBlobType(),
+            pa.array([x.reference for x in self._data], type=pa.binary()),
+        )
+
+    def astype(self, dtype, copy=True):
+        if isinstance(dtype, ArrowDtype) and dtype.pyarrow_dtype == pa.binary():
+            blobs = [None] * len(self)
+            for i, blob in enumerate(self._data):
+                with blob.open("rb") as reader:
+                    blobs[i] = reader.read()
+            return pd.array(blobs, dtype=dtype)
+        return super().astype(dtype, copy=copy)

maxframe/lib/dtypes_extension/dtypes.py

@@ -46,6 +46,25 @@ def list_(value_type: Union[pa.DataType, pa.Field]):
     return pd.ArrowDtype(pa.list_(value_type))
 
 
+def struct_(fields: Union[pa.Schema, list]):
+    """
+    Create ``pd.ArrowDtype(pa.StructType)`` instance from fields.
+
+    Parameters
+    ----------
+    fields : pyarrow.Schema or list of pyarrow.Field
+        The fields to create the struct type from.
+
+    Returns
+    -------
+    ArrowDtype
+        An ArrowDtype instance with a StructType.
+    """
+    if ArrowDtype is None:
+        raise ImportError("ArrowDtype is not supported in current environment")
+    return pd.ArrowDtype(pa.struct(fields))
+
+
 def is_map_dtype(dtype: ArrowDtype) -> bool:
     """
     Check whether the dtype is a map type.
@@ -64,3 +83,24 @@ def is_list_dtype(dtype: ArrowDtype) -> bool:
     return isinstance(dtype, ArrowDtype) and isinstance(
         dtype.pyarrow_dtype, pa.ListType
     )
+
+
+def is_struct_dtype(dtype: ArrowDtype) -> bool:
+    """
+    Check whether the dtype is a struct dtype.
+
+    Parameters
+    ----------
+    dtype : ArrowDtype
+        The dtype to check.
+
+    Returns
+    -------
+    bool
+        True if the dtype is a struct type, False otherwise.
+    """
+    if ArrowDtype is None:
+        raise ImportError("ArrowDtype is not supported in current environment")
+    return isinstance(dtype, ArrowDtype) and isinstance(
+        dtype.pyarrow_dtype, pa.StructType
+    )

maxframe/lib/dtypes_extension/tests/test_blob.py

@@ -0,0 +1,88 @@
+import copy
+import pickle
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+from ....lib.version import parse as parse_version
+from ....utils import deserialize_serializable, serialize_serializable, tokenize
+from ...wrapped_pickle import switch_unpickle
+from .. import ArrowDtype
+from ..blob import ArrowBlobType, ExternalBlobDtype, SolidBlob
+
+
+@switch_unpickle
+def test_blob_object():
+    blob = SolidBlob(b"text_content")
+    assert blob == SolidBlob(b"text_content")
+    assert blob == copy.copy(blob)
+    assert blob == deserialize_serializable(serialize_serializable(blob))
+    assert hash(blob) == hash(blob)
+    assert tokenize(blob) == tokenize(blob)
+    assert blob != b"other_content"
+
+    with blob.open("rb") as reader:
+        assert reader.read() == b"text_content"
+
+    blob = SolidBlob()
+    with pytest.raises(ValueError), blob.open("wb"):
+        raise ValueError
+    with blob.open("wb") as writer:
+        writer.write(b"text_content")
+    assert blob.reference == b"text_content"
+    with switch_unpickle(forbidden=False):
+        assert pickle.loads(pickle.dumps(blob)) == blob
+
+
+@switch_unpickle
+def test_blob_series():
+    with pytest.raises(ValueError):
+        pd.Series(np.array([["a", "b"], ["c", "d"]]), dtype="blob")
+
+    pd_ser = pd.Series(["abcd", "efgh", "ijkl"], dtype="blob")
+    assert tokenize(pd_ser) == tokenize(pd_ser)
+    assert pd_ser[0].reference == b"abcd"
+    pd.testing.assert_series_equal(pd_ser, pd_ser.copy(deep=True))
+    pd.testing.assert_series_equal(
+        pd_ser, deserialize_serializable(serialize_serializable(pd_ser))
+    )
+
+    part_ser = pd_ser.iloc[np.array([0, 1])]
+    pd.testing.assert_series_equal(
+        part_ser, pd.Series(np.array([b"abcd", b"efgh"], dtype="O"), dtype="blob")
+    )
+    ix = pd.Index([0, 2, 3])
+    part_ser = pd_ser.reindex(ix)
+    pd.testing.assert_series_equal(
+        part_ser, pd.Series([b"abcd", b"ijkl", None], index=ix, dtype="blob")
+    )
+
+    pd_ser2 = pd.Series(["abcd"], dtype="blob")
+    cat_ser = pd.concat([pd_ser, pd_ser2], ignore_index=True)
+    pd.testing.assert_series_equal(
+        cat_ser, pd.Series([b"abcd", b"efgh", b"ijkl", b"abcd"], dtype="blob")
+    )
+
+    cat_ser_bin = cat_ser.astype(ArrowDtype(pa.binary()))
+    pd.testing.assert_series_equal(
+        cat_ser_bin,
+        pd.Series([b"abcd", b"efgh", b"ijkl", b"abcd"], dtype=ArrowDtype(pa.binary())),
+    )
+
+
+def test_blob_arrow_conversion():
+    pd_ser = pd.Series([SolidBlob(b"abcd"), SolidBlob(b"efgh")], dtype="blob")
+    pa_arr = pa.Array.from_pandas(pd_ser)
+    assert pa_arr.type == ArrowBlobType()
+    try:
+        assert pa_arr.tolist() == [b"abcd", b"efgh"]
+    except NotImplementedError:
+        # compatibility test for arrow 1.0
+        assert pa_arr.storage.tolist() == [b"abcd", b"efgh"]
+
+    if parse_version(pa.__version__) >= parse_version("3.0"):
+        to_pd_ser = pa_arr.to_pandas()
+        assert to_pd_ser.dtype == ExternalBlobDtype()
+        pd.testing.assert_series_equal(pd_ser, to_pd_ser)

maxframe/lib/dtypes_extension/tests/test_dtypes.py

@@ -17,7 +17,7 @@ import pyarrow as pa
 import pytest
 
 from ....utils import ARROW_DTYPE_NOT_SUPPORTED
-from ..dtypes import dict_, is_list_dtype, is_map_dtype, list_
+from ..dtypes import dict_, is_list_dtype, is_map_dtype, is_struct_dtype, list_, struct_
 
 try:
     from pandas import ArrowDtype
@@ -46,3 +46,18 @@ def test_list_dtype():
     dt = pd.ArrowDtype(pa.map_(pa.int64(), pa.string()))
     assert not is_list_dtype(dt)
     assert not is_list_dtype(pd.Int64Dtype)
+
+
+def test_struct_dtype():
+    fields = [pa.field("a", pa.int64()), pa.field("b", pa.string())]
+    dt = struct_(fields)
+    assert is_struct_dtype(dt)
+
+    # Test with schema
+    schema = pa.schema(fields)
+    dt2 = struct_(schema)
+    assert is_struct_dtype(dt2)
+
+    dt = pd.ArrowDtype(pa.map_(pa.int64(), pa.string()))
+    assert not is_struct_dtype(dt)
+    assert not is_struct_dtype(pd.Int64Dtype)

maxframe/lib/dtypes_extension/tests/test_fake_arrow_dtype.py

@@ -0,0 +1,75 @@
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+from pandas.api.types import pandas_dtype
+
+from ....lib.version import parse as parse_version
+from ....utils import deserialize_serializable, serialize_serializable, tokenize
+from ...wrapped_pickle import switch_unpickle
+
+try:
+    from pandas import ArrowDtype  # noqa: F401
+
+    pytestmark = pytest.mark.skip("Only test when ArrowDtype not available in pandas")
+except ImportError:
+    from .._fake_arrow_dtype import FakeArrowDtype, to_pyarrow_type
+
+
+def test_fake_arrow_dtype():
+    assert to_pyarrow_type(np.dtype("int64")) == pa.int64()
+    assert to_pyarrow_type(pa.string()) == pa.string()
+    assert to_pyarrow_type(FakeArrowDtype(pa.string())) == pa.string()
+    assert to_pyarrow_type(FakeArrowDtype(pa.bool_())) == pa.bool_()
+    assert to_pyarrow_type(FakeArrowDtype(pa.int8())) == pa.int8()
+
+    pd_type = pandas_dtype("binary[pyarrow]")
+    assert isinstance(pd_type, FakeArrowDtype)
+    assert pd_type.pyarrow_dtype == pa.binary()
+
+
+@switch_unpickle
+def test_arrow_series():
+    if parse_version(pa.__version__).major < 2:
+        pytest.skip("pyarrow need to be >= 2.0 to run this case")
+
+    empty_pd_ser = pd.Series(np.array([]), dtype=FakeArrowDtype(pa.binary()))
+    assert len(empty_pd_ser) == 0
+
+    pd_ser = pd.Series([b"abcd", b"efgh", b"ijkl"], dtype=FakeArrowDtype(pa.binary()))
+    assert tokenize(pd_ser) == tokenize(pd_ser)
+    assert pd_ser[0] == b"abcd"
+    pd.testing.assert_series_equal(pd_ser, pd_ser.copy(deep=True))
+    pd.testing.assert_series_equal(
+        pd_ser, deserialize_serializable(serialize_serializable(pd_ser))
+    )
+    part_ser = pd_ser.iloc[np.array([0, 1])]
+    pd.testing.assert_series_equal(
+        part_ser,
+        pd.Series(
+            np.array([b"abcd", b"efgh"], dtype="O"), dtype=FakeArrowDtype(pa.binary())
+        ),
+    )
+    ix = pd.Index([0, 2, 3])
+    part_ser = pd_ser.reindex(ix)
+    pd.testing.assert_series_equal(
+        part_ser,
+        pd.Series(
+            [b"abcd", b"ijkl", None], index=ix, dtype=FakeArrowDtype(pa.binary())
+        ),
+    )
+    pd_ser2 = pd.Series([b"abcd"], dtype=FakeArrowDtype(pa.binary()))
+    cat_ser = pd.concat([pd_ser, pd_ser2], ignore_index=True)
+    pd.testing.assert_series_equal(
+        cat_ser,
+        pd.Series(
+            [b"abcd", b"efgh", b"ijkl", b"abcd"], dtype=FakeArrowDtype(pa.binary())
+        ),
+    )
+    cat_ser_bin = cat_ser.astype(FakeArrowDtype(pa.binary()))
+    pd.testing.assert_series_equal(
+        cat_ser_bin,
+        pd.Series(
+            [b"abcd", b"efgh", b"ijkl", b"abcd"], dtype=FakeArrowDtype(pa.binary())
+        ),
+    )