vortex-data 0.86.1__cp311-abi3-win_amd64.whl

vortex/__init__.py

@@ -0,0 +1,231 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+import importlib
+import importlib.metadata
+import importlib.util
+from types import ModuleType
+
+from . import _lib, arrays, dataset, expr, file, io, ray, registry, scan
+from ._lib.arrays import (
+    AlpArray,
+    AlpRdArray,
+    BoolArray,
+    ByteBoolArray,
+    ChunkedArray,
+    ConstantArray,
+    DateTimePartsArray,
+    # DecimalArray # TODO(connor): Is this missing a `DecimalArray`?
+    DictArray,
+    ExtensionArray,
+    FastLanesBitPackedArray,
+    FastLanesDeltaArray,
+    FastLanesFoRArray,
+    FixedSizeListArray,
+    FsstArray,
+    ListArray,
+    NullArray,
+    PrimitiveArray,
+    RunEndArray,
+    SequenceArray,
+    SparseArray,
+    StructArray,
+    VarBinArray,
+    VarBinViewArray,
+    ZigZagArray,
+)
+from ._lib.compress import compress
+from ._lib.dtype import (
+    BinaryDType,
+    BoolDType,
+    DecimalDType,
+    DType,
+    ExtensionDType,
+    FixedSizeListDType,
+    ListDType,
+    NullDType,
+    PrimitiveDType,
+    PType,
+    StructDType,
+    Utf8DType,
+    binary,
+    bool_,
+    date,
+    decimal,
+    fixed_size_list,
+    float_,
+    int_,
+    list_,
+    null,
+    struct,
+    time,
+    timestamp,
+    uint,
+    utf8,
+)
+from ._lib.iter import ArrayIterator
+from ._lib.runtime import (
+    set_worker_threads,
+    worker_threads,
+)
+from ._lib.scalar import (
+    BinaryScalar,
+    BoolScalar,
+    # TODO(connor): Is this missing a `DecimalScalar`?
+    ExtensionScalar,
+    ListScalar,
+    NullScalar,
+    PrimitiveScalar,
+    Scalar,
+    StructScalar,
+    Utf8Scalar,
+    scalar,
+)
+from ._lib.serde import ArrayContext, SerializedArray
+from .arrays import (
+    Array,
+    PyArray,
+    _unpickle_array,
+    array,
+)
+from .file import VortexFile, open
+from .scan import RepeatedScan
+
+_ = _lib  # Ensure we eagerly import the Vortex native library.
+
+# Resolve the installed distribution version so it is available as vortex.__version__.
+
+__version__ = "unknown"
+try:
+    # Try to read the installed distribution version for the Python package name.
+    __version__ = importlib.metadata.version("vortex-data")
+except importlib.metadata.PackageNotFoundError:
+    # If the distribution is not installed, keep the unknown fallback.
+    pass
+
+
+def cuda_extension_installed() -> bool:
+    """Return whether the optional Vortex CUDA extension package is importable.
+
+    The base ``vortex-data`` wheel is CPU-only. Optional CUDA functionality is
+    provided by the separate ``vortex-data-cuda`` extension package. This returns
+    ``True`` when the ``vortex_cuda`` import package can be found in the current
+    environment, which is what ``vortex-data[cuda]`` installs.
+
+    This does not probe the CUDA driver or attached devices, and it does not
+    imply that any particular GPU interop API is available. After installing the
+    extension package, use ``vortex_cuda.cuda_available()`` to check whether CUDA
+    is usable at runtime.
+    """
+    return importlib.util.find_spec("vortex_cuda") is not None
+
+
+def __getattr__(name: str) -> ModuleType:
+    # `datasets` is exposed lazily and deliberately kept out of __all__: importing it pulls in the
+    # optional `vortex-data[hf]` dependencies, so it must not be imported by `from vortex import *`
+    # or by merely importing `vortex`.
+    if name == "datasets":
+        module = importlib.import_module(".datasets", __name__)
+        globals()[name] = module
+        return module
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    # --- Modules ---
+    "arrays",
+    "dataset",
+    "expr",
+    "file",
+    "scan",
+    "io",
+    "registry",
+    "ray",
+    # --- Objects and Functions ---
+    "array",
+    "compress",
+    "cuda_extension_installed",
+    # Arrays
+    "Array",
+    "PyArray",
+    # DTypes
+    "DType",
+    "PType",
+    "NullDType",
+    "BoolDType",
+    "DecimalDType",
+    "PrimitiveDType",
+    "Utf8DType",
+    "BinaryDType",
+    "StructDType",
+    "ListDType",
+    "FixedSizeListDType",
+    "ExtensionDType",
+    "null",
+    "bool_",
+    "decimal",
+    "int_",
+    "uint",
+    "float_",
+    "utf8",
+    "binary",
+    "struct",
+    "list_",
+    "fixed_size_list",
+    "date",
+    "time",
+    "timestamp",
+    # Encodings
+    "ConstantArray",
+    "ChunkedArray",
+    "NullArray",
+    "BoolArray",
+    "ByteBoolArray",
+    "PrimitiveArray",
+    "VarBinArray",
+    "VarBinViewArray",
+    "StructArray",
+    "ListArray",
+    "FixedSizeListArray",
+    "ExtensionArray",
+    "AlpArray",
+    "AlpRdArray",
+    "DateTimePartsArray",
+    "DictArray",
+    "FsstArray",
+    "RunEndArray",
+    "SequenceArray",
+    "SparseArray",
+    "ZigZagArray",
+    "FastLanesBitPackedArray",
+    "FastLanesDeltaArray",
+    "FastLanesFoRArray",
+    # Scalars
+    "scalar",
+    "Scalar",
+    "NullScalar",
+    "BoolScalar",
+    "PrimitiveScalar",
+    "Utf8Scalar",
+    "BinaryScalar",
+    "StructScalar",
+    "ListScalar",
+    "ExtensionScalar",
+    # Serde
+    "ArrayContext",
+    "SerializedArray",
+    # Pickle support
+    "_unpickle_array",
+    # File
+    "VortexFile",
+    "open",
+    # Iterator
+    "ArrayIterator",
+    # Scan
+    "RepeatedScan",
+    # Runtime
+    "set_worker_threads",
+    "worker_threads",
+    # Version
+    "__version__",
+]

vortex/_lib/__init__.pyi

@@ -0,0 +1,61 @@
+#  SPDX-License-Identifier: Apache-2.0
+#  SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+class CosStore:
+    """A Tencent Cloud COS object store, backed by OpenDAL.
+
+    Construct it with explicit configuration and pass it to
+    ``vortex.io.read_url(url, store=cos_store)`` /
+    ``vortex.io.write(arrays, path, store=cos_store)``.
+    """
+
+    def __init__(
+        self,
+        bucket: str,
+        endpoint: str,
+        *,
+        secret_id: str | None = None,
+        secret_key: str | None = None,
+        root: str | None = None,
+        disable_config_load: bool = False,
+    ) -> None: ...
+
+class GoosefsStore:
+    """A Tencent Cloud GooseFS object store, backed by OpenDAL.
+
+    Construct it with explicit configuration and pass it to
+    ``vortex.io.read_url(url, store=goosefs_store)`` /
+    ``vortex.io.write(arrays, path, store=goosefs_store)``.
+    """
+
+    def __init__(
+        self,
+        master_addr: str,
+        *,
+        root: str | None = None,
+        block_size: int | None = None,
+        chunk_size: int | None = None,
+        write_type: str | None = None,
+        auth_type: str | None = None,
+        auth_username: str | None = None,
+    ) -> None: ...
+
+class HfStore:
+    """A Hugging Face Hub object store, rooted at one repository and revision.
+
+    Reading an ``hf://`` URL needs no store; this covers a token held in a variable
+    rather than the environment, and reads that must stay anonymous regardless of it.
+
+    Construct it and pass it to ``vortex.io.read_url(path, store=hf_store)``, where
+    ``path`` is a path within the repository.
+    """
+
+    def __init__(
+        self,
+        repo_id: str,
+        *,
+        repo_type: str = "dataset",
+        revision: str | None = None,
+        token: bool | str | None = None,
+        endpoint: str | None = None,
+    ) -> None: ...

vortex/_lib/arrays.pyi

@@ -0,0 +1,142 @@
+#  SPDX-License-Identifier: Apache-2.0
+#  SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+from typing import final
+
+import numpy as np
+import pandas as pd
+import polars as pl
+import pyarrow as pa
+
+from .dtype import DType, PType
+from .expr import Expr
+from .scalar import Scalar, ScalarPyType
+from .serde import ArrayContext
+
+class Array:
+    def __lt__(self, other: Array) -> Array: ...
+    def __le__(self, other: Array) -> Array: ...
+    def __gt__(self, other: Array) -> Array: ...
+    def __ge__(self, other: Array) -> Array: ...
+    def __len__(self) -> int: ...
+    @staticmethod
+    def from_arrow(
+        obj: pa.Array[pa.Scalar[pa.DataType]] | pa.ChunkedArray[pa.Scalar[pa.DataType]] | pa.Table,
+    ) -> Array: ...
+    @staticmethod
+    def from_range(obj: range, *, dtype: DType | None = None) -> Array: ...
+    def to_arrow_array(self, *, arrow_type: pa.DataType | None = None) -> pa.Array[pa.Scalar[pa.DataType]]: ...
+    def __vortex_array_metadata__(self) -> tuple[str, bytes, int, bytes, list[object], list[object]]: ...
+    @property
+    def id(self) -> str: ...
+    @property
+    def nbytes(self) -> int: ...
+    @property
+    def dtype(self) -> DType: ...
+    def filter(self, mask: Array) -> Array: ...
+    def scalar_at(self, index: int) -> Scalar: ...
+    def take(self, indices: Array) -> Array: ...
+    def slice(self, start: int, end: int) -> Array: ...
+    def display_tree(self) -> str: ...
+    def to_arrow_table(self) -> pa.Table: ...
+    def to_numpy(self, *, zero_copy_only: bool = True) -> np.ndarray: ...
+    def to_pandas(self) -> pd.DataFrame: ...
+    def to_polars_dataframe(self) -> pl.DataFrame: ...
+    def to_polars_series(self) -> pl.Series: ...
+    def to_pylist(self) -> list[ScalarPyType]: ...
+    def serialize(self, ctx: ArrayContext) -> bytes: ...
+    def apply(self, expr: Expr) -> Array: ...
+
+class NativeArray(Array): ...
+
+class PythonArray(Array):
+    def __init__(self, len: int, dtype: DType) -> None: ...
+
+@final
+class NullArray(Array): ...
+
+@final
+class BoolArray(Array): ...
+
+@final
+class PrimitiveArray(Array):
+    @property
+    def ptype(self) -> PType: ...
+
+@final
+class DecimalArray(Array):
+    @property
+    def precision(self) -> int: ...
+    @property
+    def scale(self) -> int: ...
+
+@final
+class VarBinArray(Array): ...
+
+@final
+class VarBinViewArray(Array): ...
+
+@final
+class StructArray(Array):
+    def field(self, name: str) -> Array: ...
+    def names(self) -> list[str]: ...
+
+@final
+class ListArray(Array): ...
+
+@final
+class FixedSizeListArray(Array): ...
+
+@final
+class ExtensionArray(Array): ...
+
+@final
+class ConstantArray(Array):
+    def scalar(self) -> Scalar: ...
+
+@final
+class ChunkedArray(Array):
+    def chunks(self) -> list[Array]: ...
+
+@final
+class ByteBoolArray(Array): ...
+
+@final
+class SequenceArray(Array): ...
+
+@final
+class SparseArray(Array): ...
+
+@final
+class AlpArray(Array): ...
+
+@final
+class AlpRdArray(Array): ...
+
+@final
+class DateTimePartsArray(Array): ...
+
+@final
+class DictArray(Array): ...
+
+@final
+class FsstArray(Array): ...
+
+@final
+class RunEndArray(Array): ...
+
+@final
+class ZigZagArray(Array):
+    @staticmethod
+    def encode(array: Array) -> Array: ...
+
+@final
+class FastLanesBitPackedArray(Array):
+    @property
+    def bit_width(self) -> int: ...
+
+@final
+class FastLanesDeltaArray(Array): ...
+
+@final
+class FastLanesFoRArray(Array): ...

vortex/_lib/cli.pyi

@@ -0,0 +1,4 @@
+#  SPDX-License-Identifier: Apache-2.0
+#  SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+def launch(args: list[str]) -> None: ...

vortex/_lib/compress.pyi

@@ -0,0 +1,6 @@
+#  SPDX-License-Identifier: Apache-2.0
+#  SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+from .arrays import Array
+
+def compress(array: Array) -> Array: ...

vortex/_lib/dataset.pyi

@@ -0,0 +1,34 @@
+#  SPDX-License-Identifier: Apache-2.0
+#  SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+import pyarrow
+
+from .arrays import Array
+from .expr import Expr
+
+class VortexDataset:
+    def to_array(
+        self,
+        columns: list[str] | list[int] | None = None,
+        row_filter: Expr | None = None,
+        indices: Array | None = None,
+        row_range: tuple[int, int] | None = None,
+    ) -> Array: ...
+    def to_record_batch_reader(
+        self,
+        columns: list[str] | list[int] | None = None,
+        row_filter: Expr | None = None,
+        indices: Array | None = None,
+        split_by: int | None = None,
+        row_range: tuple[int, int] | None = None,
+    ) -> pyarrow.RecordBatchReader: ...
+    def count_rows(
+        self,
+        row_filter: Expr | None = None,
+        split_by: int | None = None,
+        row_range: tuple[int, int] | None = None,
+    ) -> int: ...
+    def schema(self) -> pyarrow.Schema: ...
+    def splits(self) -> list[tuple[int, int]]: ...
+
+def dataset_from_url(url: str) -> VortexDataset: ...

vortex/_lib/dtype.pyi

@@ -0,0 +1,80 @@
+#  SPDX-License-Identifier: Apache-2.0
+#  SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+from typing import Literal, final
+
+import pyarrow as pa
+
+class DType:
+    def is_nullable(self) -> bool: ...
+    def to_arrow_schema(self) -> pa.Schema: ...
+    def to_arrow_type(self) -> pa.DataType: ...
+    @classmethod
+    def from_arrow(cls, arrow_dtype: pa.DataType | pa.Schema, *, non_nullable: bool = False) -> DType: ...
+
+@final
+class NullDType(DType): ...
+
+@final
+class BoolDType(DType): ...
+
+@final
+class PrimitiveDType(DType):
+    @property
+    def ptype(self) -> PType: ...
+
+@final
+class DecimalDType(DType):
+    @property
+    def precision(self) -> int: ...
+    @property
+    def scale(self) -> int: ...
+
+@final
+class Utf8DType(DType): ...
+
+@final
+class BinaryDType(DType): ...
+
+@final
+class StructDType(DType):
+    def names(self) -> list[str]: ...
+    def fields(self) -> list[DType]: ...
+
+@final
+class ListDType(DType): ...
+
+@final
+class FixedSizeListDType(DType): ...
+
+@final
+class ExtensionDType(DType): ...
+
+@final
+class PType:
+    U8: PType
+    U16: PType
+    U32: PType
+    U64: PType
+    I8: PType
+    I16: PType
+    I32: PType
+    I64: PType
+    F16: PType
+    F32: PType
+    F64: PType
+
+def null() -> DType: ...
+def bool_(*, nullable: bool = False) -> DType: ...
+def int_(width: Literal[8, 16, 32, 64] = 64, *, nullable: bool = False) -> DType: ...
+def uint(width: Literal[8, 16, 32, 64] = 64, *, nullable: bool = False) -> DType: ...
+def float_(width: Literal[16, 32, 64] = 64, *, nullable: bool = False) -> DType: ...
+def decimal(*, precision: int = 10, scale: int = 0, nullable: bool = False) -> DType: ...
+def utf8(*, nullable: bool = False) -> DType: ...
+def binary(*, nullable: bool = False) -> DType: ...
+def struct(fields: dict[str, DType] | None = None, *, nullable: bool = False) -> DType: ...
+def list_(element: DType, *, nullable: bool = False) -> DType: ...
+def fixed_size_list(element: DType, size: int, *, nullable: bool = False) -> DType: ...
+def date(unit: Literal["ms", "days"], *, nullable: bool = False) -> DType: ...
+def time(unit: Literal["s", "ms", "us", "ns"], *, nullable: bool = False) -> DType: ...
+def timestamp(unit: Literal["s", "ms", "us", "ns"], *, tz: str | None = None, nullable: bool = False) -> DType: ...

vortex/_lib/expr.pyi

@@ -0,0 +1,122 @@
+#  SPDX-License-Identifier: Apache-2.0
+#  SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+from collections.abc import Iterable, Mapping, Sequence
+from datetime import date, datetime
+from typing import Literal, TypeAlias, final
+
+from typing_extensions import override
+
+from .dtype import DType
+from .scalar import ScalarPyType
+
+IntoExpr: TypeAlias = Expr | bool | int | float | str | bytes | date | datetime | None
+"""A value accepted anywhere an expression is expected. Non-``Expr`` values become literals."""
+
+VariantPath: TypeAlias = str | int | Sequence[str | int]
+
+@final
+class Expr:
+    @override
+    def __eq__(self, other: IntoExpr) -> Expr: ...  # ty: ignore[invalid-method-override]
+    @override
+    def __ne__(self, other: IntoExpr) -> Expr: ...  # ty: ignore[invalid-method-override]
+    def __lt__(self, other: IntoExpr) -> Expr: ...
+    def __le__(self, other: IntoExpr) -> Expr: ...
+    def __gt__(self, other: IntoExpr) -> Expr: ...
+    def __ge__(self, other: IntoExpr) -> Expr: ...
+    def __and__(self, other: IntoExpr) -> Expr: ...
+    def __rand__(self, other: IntoExpr) -> Expr: ...
+    def __or__(self, other: IntoExpr) -> Expr: ...
+    def __ror__(self, other: IntoExpr) -> Expr: ...
+    def __invert__(self) -> Expr: ...
+    def __add__(self, other: IntoExpr) -> Expr: ...
+    def __radd__(self, other: IntoExpr) -> Expr: ...
+    def __sub__(self, other: IntoExpr) -> Expr: ...
+    def __rsub__(self, other: IntoExpr) -> Expr: ...
+    def __mul__(self, other: IntoExpr) -> Expr: ...
+    def __rmul__(self, other: IntoExpr) -> Expr: ...
+    def __truediv__(self, other: IntoExpr) -> Expr: ...
+    def __rtruediv__(self, other: IntoExpr) -> Expr: ...
+    def __getitem__(self, field: str) -> Expr: ...
+    def serialize(self) -> bytes: ...
+    @override
+    def __reduce__(self) -> tuple[object, tuple[bytes]]: ...
+
+# Leaves and scope
+def root() -> Expr: ...
+def column(name: str) -> Expr: ...
+def literal(dtype: DType, value: ScalarPyType) -> Expr: ...
+def get_item(field: str, child: IntoExpr | None = None) -> Expr: ...
+
+# Boolean logic
+def not_(child: IntoExpr) -> Expr: ...
+def and_(left: IntoExpr, right: IntoExpr) -> Expr: ...
+def or_(left: IntoExpr, right: IntoExpr) -> Expr: ...
+def and_collect(exprs: Iterable[IntoExpr]) -> Expr | None: ...
+def or_collect(exprs: Iterable[IntoExpr]) -> Expr | None: ...
+
+# Comparisons and arithmetic
+def eq(left: IntoExpr, right: IntoExpr) -> Expr: ...
+def not_eq(left: IntoExpr, right: IntoExpr) -> Expr: ...
+def gt(left: IntoExpr, right: IntoExpr) -> Expr: ...
+def gt_eq(left: IntoExpr, right: IntoExpr) -> Expr: ...
+def lt(left: IntoExpr, right: IntoExpr) -> Expr: ...
+def lt_eq(left: IntoExpr, right: IntoExpr) -> Expr: ...
+def add(left: IntoExpr, right: IntoExpr) -> Expr: ...
+def sub(left: IntoExpr, right: IntoExpr) -> Expr: ...
+def mul(left: IntoExpr, right: IntoExpr) -> Expr: ...
+def div(left: IntoExpr, right: IntoExpr) -> Expr: ...
+def between(
+    child: IntoExpr,
+    lower: IntoExpr,
+    upper: IntoExpr,
+    *,
+    lower_strict: bool = False,
+    upper_strict: bool = False,
+) -> Expr: ...
+
+# Nullability
+def is_null(child: IntoExpr) -> Expr: ...
+def is_not_null(child: IntoExpr) -> Expr: ...
+def fill_null(child: IntoExpr, fill_value: IntoExpr) -> Expr: ...
+
+# Strings
+def like(child: IntoExpr, pattern: IntoExpr) -> Expr: ...
+def ilike(child: IntoExpr, pattern: IntoExpr) -> Expr: ...
+def not_like(child: IntoExpr, pattern: IntoExpr) -> Expr: ...
+def not_ilike(child: IntoExpr, pattern: IntoExpr) -> Expr: ...
+def byte_length(child: IntoExpr) -> Expr: ...
+
+# Structs
+def select(fields: str | Iterable[str], child: IntoExpr | None = None) -> Expr: ...
+def select_exclude(fields: str | Iterable[str], child: IntoExpr | None = None) -> Expr: ...
+def pack(
+    fields: Mapping[str, IntoExpr] | Iterable[tuple[str, IntoExpr]],
+    *,
+    nullable: bool = False,
+) -> Expr: ...
+def merge(
+    exprs: Iterable[IntoExpr],
+    *,
+    duplicate_handling: Literal["error", "rightmost"] = "error",
+) -> Expr: ...
+
+# Lists
+def list_contains(child: IntoExpr, value: IntoExpr) -> Expr: ...
+def list_length(child: IntoExpr) -> Expr: ...
+def list_sum(child: IntoExpr, *, skip_nans: bool = True) -> Expr: ...
+
+# Conditionals and misc
+def case_when(
+    when_then: Iterable[tuple[IntoExpr, IntoExpr]],
+    else_value: IntoExpr | None = None,
+) -> Expr: ...
+def zip_(mask: IntoExpr, if_true: IntoExpr, if_false: IntoExpr) -> Expr: ...
+def mask(child: IntoExpr, mask: IntoExpr) -> Expr: ...
+def cast(child: IntoExpr, dtype: DType) -> Expr: ...
+def ext_storage(child: IntoExpr) -> Expr: ...
+def variant_get(child: IntoExpr, path: VariantPath, dtype: DType | None = None) -> Expr: ...
+
+# Serialization
+def deserialize(data: bytes) -> Expr: ...