pyspiral 0.7.18__cp312-abi3-manylinux_2_28_x86_64.whl

spiral/expressions/__init__.py

@@ -0,0 +1,274 @@
+import builtins
+import functools
+import operator
+import warnings
+from typing import Any
+
+import pyarrow as pa
+
+from spiral import _lib, arrow_
+
+from . import file as file
+from . import http as http
+from . import list_ as list
+from . import s3 as s3
+from . import str_ as str
+from . import struct as struct
+from . import text as text
+from .base import Expr, ExprLike, NativeExpr
+from .udf import UDF
+
+__all__ = [
+    "Expr",
+    "add",
+    "and_",
+    "divide",
+    "eq",
+    "getitem",
+    "gt",
+    "gte",
+    "is_not_null",
+    "is_null",
+    "lift",
+    "list",
+    "lt",
+    "lte",
+    "merge",
+    "modulo",
+    "multiply",
+    "negate",
+    "neq",
+    "not_",
+    "or_",
+    "pack",
+    "aux",
+    "scalar",
+    "select",
+    "str",
+    "struct",
+    "subtract",
+    "xor",
+    "text",
+    "s3",
+    "http",
+    "file",
+    "UDF",
+]
+
+# Inline some of the struct expressions since they're so common
+getitem = struct.getitem
+merge = struct.merge
+pack = struct.pack
+select = struct.select
+
+
+def lift(expr: ExprLike) -> Expr:
+    # Convert an ExprLike into an Expr.
+
+    if isinstance(expr, Expr):
+        return expr
+    if isinstance(expr, NativeExpr):
+        return Expr(expr)
+
+    if isinstance(expr, dict):
+        # NOTE: we assume this is a struct expression. We could be smarter and be context aware to determine if
+        # this is in fact a struct scalar, but the user can always create one of those manually.
+
+        # First we un-nest any dot-separated field names
+        expr: dict = arrow_.nest_structs(expr)
+
+        return pack({k: lift(v) for k, v in expr.items()})
+
+    if isinstance(expr, builtins.list):
+        return lift(pa.array(expr))
+
+    # Unpack tables and chunked arrays
+    if isinstance(expr, pa.Table | pa.RecordBatch):
+        expr = expr.to_struct_array()
+    if isinstance(expr, pa.ChunkedArray):
+        expr = expr.combine_chunks()
+
+    # If the value is struct-like, we un-nest any dot-separated field names
+    if isinstance(expr, pa.StructArray | pa.StructScalar):
+        # TODO(marko): Figure out what to do with nullable struct arrays when unpacking them.
+        #   We need to merge struct validity into the child validity?
+        if isinstance(expr, pa.StructArray) and expr.null_count != 0:
+            # raise ValueError("lift: cannot lift a struct array with nulls.")
+            warnings.warn("found a struct array with nulls", stacklevel=2)
+        if isinstance(expr, pa.StructScalar) and not expr.is_valid:
+            # raise ValueError("lift: cannot lift a struct scalar with nulls.")
+            warnings.warn("found a struct scalar with nulls", stacklevel=2)
+        return lift(arrow_.nest_structs(expr))
+
+    if isinstance(expr, pa.Array):
+        return Expr(_lib.expr.array_lit(expr))
+
+    # Otherwise, assume it's a scalar.
+    return scalar(expr)
+
+
+def evaluate(expr: ExprLike) -> pa.RecordBatchReader:
+    # TODO(marko): This implementation is currently minimal and most ExprLike-s fail.
+    if isinstance(expr, pa.RecordBatchReader):
+        return expr
+    if isinstance(expr, pa.Table):
+        return expr.to_reader()
+    if isinstance(expr, pa.RecordBatch):
+        return pa.RecordBatchReader.from_batches(expr.schema, [expr])
+    if isinstance(expr, pa.StructArray):
+        return pa.Table.from_struct_array(expr).to_reader()
+
+    if isinstance(expr, pa.ChunkedArray):
+        if not pa.types.is_struct(expr.type):
+            raise ValueError("Arrow chunked array must be a struct type.")
+
+        def _iter_batches():
+            for chunk in expr.chunks:
+                yield pa.RecordBatch.from_struct_array(chunk)
+
+        return pa.RecordBatchReader.from_batches(pa.schema(expr.type.fields), _iter_batches())
+
+    if isinstance(expr, pa.Array):
+        raise ValueError("Arrow array must be a struct array.")
+
+    if isinstance(expr, Expr) or isinstance(expr, NativeExpr):
+        raise NotImplementedError(
+            "Expr evaluation not supported yet. Use Arrow to write instead. Reach out if you require this feature."
+        )
+
+    if isinstance(expr, dict):
+        # NOTE: we assume this is a struct expression. We could be smarter and be context aware to determine if
+        # this is in fact a struct scalar, but the user can always create one of those manually.
+
+        # First we un-nest any dot-separated field names
+        expr: dict = arrow_.nest_structs(expr)
+        return evaluate(arrow_.dict_to_table(expr))
+
+    if isinstance(expr, builtins.list):
+        return evaluate(pa.array(expr))
+
+    if isinstance(expr, pa.Scalar):
+        return evaluate(pa.array([expr]))
+
+    # Otherwise, try scalar.
+    return evaluate(scalar(expr))
+
+
+def aux(name: builtins.str, dtype: pa.DataType) -> Expr:
+    """Create a variable expression referencing a column in the auxiliary table.
+
+    Auxiliary table is optionally given to `Scan#to_record_batches` function when reading only specific keys
+    or doing cell pushdown.
+
+    Args:
+        name: variable name
+        dtype: must match dtype of the column in the auxiliary table.
+    """
+    return Expr(_lib.expr.aux(name, dtype))
+
+
+def scalar(value: Any) -> Expr:
+    """Create a scalar expression."""
+    if not isinstance(value, pa.Scalar):
+        value = pa.scalar(value)
+    # TODO(marko): Use Vortex scalar instead of passing as array.
+    return Expr(_lib.expr.scalar(pa.array([value.as_py()], type=value.type)))
+
+
+def cast(expr: ExprLike, dtype: pa.DataType) -> Expr:
+    """Cast an expression into another PyArrow DataType."""
+    expr = lift(expr)
+    return Expr(_lib.expr.cast(expr.__expr__, dtype))
+
+
+def and_(expr: ExprLike, *exprs: ExprLike) -> Expr:
+    """Create a conjunction of one or more expressions."""
+
+    return functools.reduce(operator.and_, [lift(e) for e in exprs], lift(expr))
+
+
+def or_(expr: ExprLike, *exprs: ExprLike) -> Expr:
+    """Create a disjunction of one or more expressions."""
+    return functools.reduce(operator.or_, [lift(e) for e in exprs], lift(expr))
+
+
+def eq(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Create an equality comparison."""
+    return operator.eq(lift(lhs), rhs)
+
+
+def neq(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Create a not-equal comparison."""
+    return operator.ne(lift(lhs), rhs)
+
+
+def xor(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Create a XOR comparison."""
+    return operator.xor(lift(lhs), rhs)
+
+
+def lt(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Create a less-than comparison."""
+    return operator.lt(lift(lhs), rhs)
+
+
+def lte(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Create a less-than-or-equal comparison."""
+    return operator.le(lift(lhs), rhs)
+
+
+def gt(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Create a greater-than comparison."""
+    return operator.gt(lift(lhs), rhs)
+
+
+def gte(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Create a greater-than-or-equal comparison."""
+    return operator.ge(lift(lhs), rhs)
+
+
+def negate(expr: ExprLike) -> Expr:
+    """Negate the given expression."""
+    return operator.neg(lift(expr))
+
+
+def not_(expr: ExprLike) -> Expr:
+    """Negate the given expression."""
+    expr = lift(expr)
+    return Expr(_lib.expr.not_(expr.__expr__))
+
+
+def is_null(expr: ExprLike) -> Expr:
+    """Check if the given expression is null."""
+    expr = lift(expr)
+    return Expr(_lib.expr.is_null(expr.__expr__))
+
+
+def is_not_null(expr: ExprLike) -> Expr:
+    """Check if the given expression is not null."""
+    return not_(is_null(expr))
+
+
+def add(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Add two expressions."""
+    return operator.add(lift(lhs), rhs)
+
+
+def subtract(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Subtract two expressions."""
+    return operator.sub(lift(lhs), rhs)
+
+
+def multiply(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Multiply two expressions."""
+    return operator.mul(lift(lhs), rhs)
+
+
+def divide(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Divide two expressions."""
+    return operator.truediv(lift(lhs), rhs)
+
+
+def modulo(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Modulo two expressions."""
+    return operator.mod(lift(lhs), rhs)

spiral/expressions/base.py

@@ -0,0 +1,167 @@
+import datetime
+from typing import TypeAlias, Union
+
+import pyarrow as pa
+
+from spiral import _lib
+
+NativeExpr: TypeAlias = _lib.expr.Expr
+
+
+class Expr:
+    """Base class for Spiral expressions. All expressions support comparison and basic arithmetic operations."""
+
+    def __init__(self, native: NativeExpr) -> None:
+        if not isinstance(native, NativeExpr):
+            raise TypeError(f"Expected a native expression, got {type(native)}")
+        self._native = native
+
+    @property
+    def __expr__(self) -> NativeExpr:
+        return self._native
+
+    def __str__(self):
+        return str(self.__expr__)
+
+    def __repr__(self):
+        return repr(self.__expr__)
+
+    def __getitem__(self, item: str | int | list[str]) -> "Expr":
+        """
+        Get an item from a struct or list.
+
+        Args:
+            item: The key or index to get.
+                If item is a string, it is assumed to be a field in a struct. Dot-separated string is supported
+                to access nested fields.
+                If item is a list of strings, it is assumed to be a list of fields in a struct.
+                If item is an integer, it is assumed to be an index in a list.
+        """
+        from spiral import expressions as se
+
+        expr = self
+
+        if isinstance(item, int):
+            # Assume list and get an element.
+            expr = se.list_.element_at(expr, item)
+        elif isinstance(item, str):
+            # Walk into the struct.
+            for part in item.split("."):
+                expr = se.getitem(expr, part)
+        elif isinstance(item, list) and all(isinstance(i, str) for i in item):
+            expr = se.pack({k: expr[k] for k in item})
+        else:
+            raise TypeError(f"Invalid item type: {type(item)}")
+
+        return expr
+
+    def __eq__(self, other: "ExprLike") -> "Expr":
+        return self._binary("eq", other)
+
+    def __ne__(self, other: "ExprLike") -> "Expr":
+        return self._binary("neq", other)
+
+    def __lt__(self, other: "ExprLike") -> "Expr":
+        return self._binary("lt", other)
+
+    def __le__(self, other: "ExprLike") -> "Expr":
+        return self._binary("lte", other)
+
+    def __gt__(self, other: "ExprLike") -> "Expr":
+        return self._binary("gt", other)
+
+    def __ge__(self, other: "ExprLike") -> "Expr":
+        return self._binary("gte", other)
+
+    def __and__(self, other: "ExprLike") -> "Expr":
+        return self._binary("and", other)
+
+    def __or__(self, other: "ExprLike") -> "Expr":
+        return self._binary("or", other)
+
+    def __xor__(self, other: "ExprLike") -> "Expr":
+        raise NotImplementedError
+
+    def __add__(self, other: "ExprLike") -> "Expr":
+        return self._binary("add", other)
+
+    def __sub__(self, other: "ExprLike") -> "Expr":
+        return self._binary("sub", other)
+
+    def __mul__(self, other: "ExprLike") -> "Expr":
+        return self._binary("mul", other)
+
+    def __truediv__(self, other: "ExprLike") -> "Expr":
+        return self._binary("div", other)
+
+    def __mod__(self, other: "ExprLike") -> "Expr":
+        return self._binary("mod", other)
+
+    def __neg__(self):
+        return Expr(_lib.expr.unary("neg", self.__expr__))
+
+    def in_(self, other: "ExprLike") -> "Expr":
+        from spiral import expressions as se
+
+        other = se.lift(other)
+        return Expr(_lib.expr.list.contains(other.__expr__, self.__expr__))
+
+    def contains(self, other: "ExprLike") -> "Expr":
+        from spiral import expressions as se
+
+        return se.lift(other).in_(self)
+
+    def cast(self, dtype: pa.DataType) -> "Expr":
+        """Cast the expression result to a different data type."""
+        return Expr(_lib.expr.cast(self.__expr__, dtype))
+
+    def select(self, *paths: str, exclude: list[str] = None) -> "Expr":
+        """Select fields from a struct-like expression.
+
+        Args:
+            *paths: Field names to select. If a path contains a dot, it is assumed to be a nested struct field.
+            exclude: List of field names to exclude from result.
+        """
+        from spiral import expressions as se
+
+        # If any of the paths contain nested fields, then we re-pack nested select statements.
+        if any("." in p for p in paths):
+            fields = {}
+            for p in paths:
+                if "." in p:
+                    parent, child = p.split(".", 1)
+                    fields[parent] = self[parent].select(child)
+                else:
+                    fields[p] = self[p]
+            packed = se.pack(fields)
+            if exclude:
+                packed = packed.select(exclude=exclude)
+            return packed
+
+        if paths:
+            if exclude:
+                raise ValueError("Cannot specify both selection and exclusion.")
+            return se.select(self, names=list(paths))
+
+        if exclude:
+            return se.select(self, exclude=exclude)
+
+        return self
+
+    def _binary(self, op: str, rhs: "ExprLike") -> "Expr":
+        """Create a comparison expression."""
+        from spiral import expressions as se
+
+        rhs = se.lift(rhs)
+        return Expr(_lib.expr.binary(op, self.__expr__, rhs.__expr__))
+
+
+ScalarLike: TypeAlias = bool | int | float | str | list["ScalarLike"] | datetime.datetime | None
+ArrowLike: TypeAlias = Union[
+    pa.RecordBatch,
+    "pa.Array[pa.Scalar[pa.DataType]]",
+    "pa.ChunkedArray[pa.Scalar[pa.DataType]]",
+    "pa.Scalar[pa.DataType]",
+    pa.Table,
+]
+ExprLike: TypeAlias = Expr | dict[str, "ExprLike"] | list["ExprLike"] | ArrowLike | ScalarLike

spiral/expressions/file.py

@@ -0,0 +1,17 @@
+from spiral import _lib
+from spiral.expressions.base import Expr, ExprLike
+
+
+def get(expr: ExprLike, abort_on_error: bool = False) -> Expr:
+    """Read data from the local filesystem by the file:// URL.
+
+    Args:
+        expr: URLs of the data that needs to be read.
+        abort_on_error: Should the expression abort on errors or just collect them.
+    """
+    from spiral import expressions as se
+
+    expr = se.lift(expr)
+
+    # This just works :)
+    return Expr(_lib.expr.s3.get(expr.__expr__, abort_on_error))

spiral/expressions/http.py

@@ -0,0 +1,17 @@
+from spiral import _lib
+from spiral.expressions.base import Expr, ExprLike
+
+
+def get(expr: ExprLike, abort_on_error: bool = False) -> Expr:
+    """Read data from the URL.
+
+    Args:
+        expr: URLs of the data that needs to be read.
+        abort_on_error: Should the expression abort on errors or just collect them.
+    """
+    from spiral import expressions as se
+
+    expr = se.lift(expr)
+
+    # This just works :)
+    return Expr(_lib.expr.s3.get(expr.__expr__, abort_on_error))

spiral/expressions/list_.py

@@ -0,0 +1,68 @@
+from spiral.expressions.base import Expr, ExprLike
+
+
+def in_(expr: ExprLike, values: ExprLike) -> Expr:
+    """Check if a value is in a list.
+
+    Args:
+        expr: The value to check.
+        values: The list array expression to check against.
+    """
+    # `se.list.in_(Array[2, 4], Array[[1, 2], [1, 2]]) -> Array[True, False]`
+    from spiral.expressions import lift
+
+    expr = lift(expr)
+    return expr.in_(values)
+
+
+def element_at(expr: ExprLike, index: ExprLike) -> Expr:
+    """Get the element at the given index.
+
+    Args:
+        expr: The list array expression.
+        index: The index to get.
+    """
+    # e.g. `se.list.element_at([1, 2, 3], 1) -> 2`
+    ...
+    from spiral import _lib
+    from spiral.expressions import lift
+
+    expr = lift(expr)
+    index = lift(index)
+    return Expr(_lib.expr.list.element_at(expr.__expr__, index.__expr__))
+
+
+def of(*expr: ExprLike) -> Expr:
+    # Creates an array or scalar list from a series of expressions, all values must be of the same type.
+    # The expressions must all also have the same length (1 for scalars).
+    #
+    # e.g. `se.list.of(1+3, 2, 3) -> [4, 2, 3]`
+    ...
+
+
+def zip(*lists: ExprLike) -> Expr:
+    # Merge the given lists, with duplicates.
+    #
+    # e.g. `se.list.merge([1, 2], [3, 4]) -> [(1, 2), (3, 4)]`
+    ...
+
+
+def concat(*lists: ExprLike) -> Expr:
+    # Concatenate the given lists. The types of all the lists must be the same.
+    #
+    # e.g. `se.list.concat([1, 2], [3, 4]) -> [1, 2, 3, 4]`
+    ...
+
+
+def slice_(expr: ExprLike, start: int | None = None, stop: int | None = None) -> Expr:
+    # Slice a list.
+    #
+    # e.g. `se.list.slice_([0, 1, 2], slice(0,2)) -> [0, 1]`
+    ...
+
+
+def length(expr: ExprLike) -> Expr:
+    # Get the length of a list.
+    #
+    # e.g. `se.list.length([1, 2, 3]) -> 3`
+    ...

spiral/expressions/s3.py

@@ -0,0 +1,16 @@
+from spiral import _lib
+from spiral.expressions.base import Expr, ExprLike
+
+
+def get(expr: ExprLike, abort_on_error: bool = False) -> Expr:
+    """Read data from object storage by the s3:// URL.
+
+    Args:
+        expr: URLs of the data that needs to be read from object storage.
+        abort_on_error: Should the expression abort on errors or just collect them.
+    """
+    from spiral import expressions as se
+
+    expr = se.lift(expr)
+
+    return Expr(_lib.expr.s3.get(expr.__expr__, abort_on_error))

spiral/expressions/str_.py

@@ -0,0 +1,39 @@
+import pyarrow as pa
+import pyarrow.compute as pc
+import re2 as re
+
+from spiral import _lib
+from spiral.expressions.base import Expr, ExprLike
+
+# TODO(ngates): we can add a symmetric "ascii" expression namespace in the future if
+#  the performance is required.
+
+
+def substr(expr: ExprLike = None, *, begin: int = 0, end: int | None = None) -> Expr:
+    """Slice a string.
+
+    Args:
+        expr: The string expression to slice.
+        begin: The starting index of the slice.
+        end: The ending index of the slice.
+    """
+    from spiral import expressions as se
+
+    expr = se.lift(expr)
+    return Expr(_lib.spql.str.substr(expr.__expr__, begin=begin, end=end))
+
+
+def extract_regex(pattern: str, *, strings: ExprLike) -> Expr:
+    # Extract the first occurrence of a regex pattern from a string.
+    raise NotImplementedError
+
+
+def _extract_regex(arg: pa.Array | pa.Scalar, pattern: str) -> pa.Array | pa.Scalar:
+    # Compute the return type based on the regex groups
+    m = re.compile(pattern)
+    dtype = pa.struct([pa.field(k, type=pa.string()) for k in m.groupindex.keys()])
+
+    if pa.types.is_string(arg.type):
+        return pc.extract_regex(arg, pattern=pattern).cast(dtype)
+
+    raise TypeError("Input argument does not have the expected type")

spiral/expressions/struct.py

@@ -0,0 +1,59 @@
+from spiral import _lib
+from spiral.expressions.base import Expr, ExprLike
+
+
+def getitem(expr: ExprLike, field: str) -> Expr:
+    """Get field from a struct.
+
+    Args:
+        expr: The struct expression to get the field from.
+        field: The field to get. Dot-separated string is supported to access nested fields.
+    """
+    from spiral import expressions as se
+
+    expr = se.lift(expr)
+    return Expr(_lib.expr.struct.getitem(expr.__expr__, field))
+
+
+def pack(fields: dict[str, ExprLike], *, nullable: bool = False) -> Expr:
+    """Assemble a new struct from the given named fields.
+
+    Args:
+        fields: A dictionary of field names to expressions. The field names will be used as the struct field names.
+    """
+    from spiral import expressions as se
+
+    return Expr(
+        _lib.expr.struct.pack(list(fields.keys()), [se.lift(expr).__expr__ for expr in fields.values()], nullable)
+    )
+
+
+def merge(*structs: "ExprLike") -> Expr:
+    """Merge fields from the given structs into a single struct.
+
+    Args:
+        *structs: Each expression must evaluate to a struct.
+
+    Returns:
+        A single struct containing all the fields from the input structs.
+        If a field is present in multiple structs, the value from the last struct is used.
+    """
+    from spiral import expressions as se
+
+    if len(structs) == 1:
+        return se.lift(structs[0])
+    return Expr(_lib.expr.struct.merge([se.lift(struct).__expr__ for struct in structs]))
+
+
+def select(expr: ExprLike, names: list[str] = None, exclude: list[str] = None) -> Expr:
+    """Select fields from a struct.
+
+    Args:
+        expr: The struct-like expression to select fields from.
+        names: Field names to select. If a path contains a dot, it is assumed to be a nested struct field.
+        exclude: List of field names to exclude from result. Exactly one of `names` or `exclude` must be provided.
+    """
+    from spiral import expressions as se
+
+    expr = se.lift(expr)
+    return Expr(_lib.expr.struct.select(expr.__expr__, names, exclude))