pyspiral 0.8.9__cp311-abi3-macosx_11_0_arm64.whl

spiral/expressions/base.py

@@ -0,0 +1,186 @@
+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):
+            if item == "*":
+                # Select all fields when "*" is used in getitem.
+                # getitem("*") is not a valid expression, so we translate it to select(all).
+                expr = se.select(expr, names=[item])
+            else:
+                # 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 = expr.select(*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 paths:
+            if exclude:
+                raise ValueError("Cannot specify both selection and exclusion.")
+
+            # 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 p == "*":
+                        # This is handled later.
+                        continue
+
+                    if "." not in p:
+                        fields[p] = self[p]
+                        continue
+
+                    parent, child = p.split(".", 1)
+                    node = self[parent].select(child)
+                    if parent in fields:
+                        fields[parent] = se.merge(fields[parent], node)
+                    else:
+                        fields[parent] = node
+
+                packed = se.pack(fields)
+                if "*" in paths:
+                    packed = se.merge(self.select("*"), packed)
+
+                return packed
+
+            return se.select(self, names=list(paths))
+
+        if exclude:
+            if any("." in p for p in exclude):
+                raise ValueError("Exclusion of nested fields is not supported yet.")
+            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,77 @@
+from spiral import _lib
+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`
+    ...
+
+
+def contains(expr: ExprLike, value: ExprLike) -> bool:
+    from spiral import expressions as se
+
+    expr = se.lift(expr)
+    value = se.lift(value)
+    return _lib.expr.list.contains(expr.__expr__, value.__expr__)

spiral/expressions/pushdown.py

@@ -0,0 +1,12 @@
+from spiral import _lib
+from spiral.expressions.base import Expr, ExprLike
+
+
+def expensive(expr: ExprLike) -> Expr:
+    """Minimise the chance of evaluation the expression over the old version of data,
+    applying it as late as possible."""
+    from spiral import expressions as se
+
+    expr = se.lift(expr)
+
+    return Expr(_lib.expr.pushdown.expensive(expr.__expr__))

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))

spiral/expressions/text.py

@@ -0,0 +1,62 @@
+from spiral.expressions.base import Expr, ExprLike
+
+
+def field(expr: ExprLike, field_name: str | None = None, tokenizer: str | None = None) -> Expr:
+    """Configure a column for text indexing.
+
+    Args:
+        expr: An input column. The expression must either evaluate to a UTF-8,
+            or, if a `field_name` is provided, to a struct with a field of that name.
+        field_name: If provided, the expression must evaluate to a struct with a field of that name.
+            The given field will be indexed.
+        tokenizer: If provided, the text will be tokenized using the given tokenizer.
+
+    Returns:
+        An expression that can be used to construct a text index.
+    """
+    from spiral import _lib
+    from spiral.expressions import getitem, lift, merge, pack
+
+    expr = lift(expr)
+    if field_name is None:
+        return Expr(_lib.expr.text.field(expr.__expr__, tokenizer))
+
+    child = _lib.expr.text.field(getitem(expr, field_name).__expr__)
+    return merge(
+        expr,
+        pack({field_name: child}),
+    )
+
+
+def find(expr: ExprLike, term: str) -> Expr:
+    """Search for a term in the text.
+
+    Args:
+        expr: An index field.
+        term: The term to search for.
+
+    Returns:
+        An expression that can be used in ranking for text search.
+    """
+    from spiral import _lib
+    from spiral.expressions import lift
+
+    expr = lift(expr)
+    return Expr(_lib.expr.text.find(expr.__expr__, term))
+
+
+def boost(expr: ExprLike, factor: float) -> Expr:
+    """Boost the relevance of a ranking expression.
+
+    Args:
+        expr: Rank by expression.
+        factor: The factor by which to boost the relevance.
+
+    Returns:
+        An expression that can be used in ranking for text search.
+    """
+    from spiral import _lib
+    from spiral.expressions import lift
+
+    expr = lift(expr)
+    return Expr(_lib.expr.text.boost(expr.__expr__, factor))