pyspiral 0.2.5__pp310-pypy310_pp73-macosx_10_13_x86_64.whl

spiral/expressions/base.py

@@ -0,0 +1,149 @@
+import datetime
+from typing import TypeAlias
+
+import pyarrow as pa
+
+from spiral import _lib
+
+NativeExpr: TypeAlias = _lib.spql.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) -> "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 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)
+        else:
+            # Walk into the struct.
+            for part in item.split("."):
+                expr = se.getitem(expr, part)
+
+        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.spql.expr.unary("neg", self.__expr__))
+
+    def in_(self, other: "ExprLike") -> "Expr":
+        from spiral import expressions as se
+
+        other = se.lift(other)
+        return Expr(_lib.spql.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.spql.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 not paths:
+            return self
+
+        return se.select(self, names=list(paths), exclude=exclude)
+
+    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.spql.expr.binary(op, self.__expr__, rhs.__expr__))
+
+
+ScalarLike: TypeAlias = bool | int | float | str | list | datetime.datetime | None
+ExprLike: TypeAlias = Expr | dict | ScalarLike

spiral/expressions/http.py

@@ -0,0 +1,86 @@
+import hishel
+import httpx
+import pyarrow as pa
+
+from spiral.expressions.base import Expr, ExprLike
+from spiral.expressions.struct import pack
+from spiral.expressions.udf import UDF
+from spiral.settings import APP_DIR
+
+
+def get(url: ExprLike, headers: ExprLike = None, force_cache: bool = False) -> Expr:
+    """Submit a GET request to either a scalar of vector of URLs."""
+    to_pack = {"url": url}
+    if headers is not None:
+        to_pack["headers"] = headers
+    return HttpGet(force_cache)(pack(to_pack))
+
+
+class HttpGet(UDF):
+    RES_DTYPE: pa.DataType = pa.struct(
+        [
+            pa.field("bytes", pa.large_binary()),
+            pa.field("status", pa.int32()),
+            pa.field("headers", pa.map_(pa.string(), pa.string())),
+        ]
+    )
+
+    def __init__(self, force_cache: bool = False):
+        super().__init__("http.get")
+        self._force_cache = force_cache
+
+    def return_type(self, *input_types: pa.DataType) -> pa.DataType:
+        return HttpGet.RES_DTYPE
+
+    def invoke(self, *input_args: pa.Array) -> pa.Array:
+        if len(input_args) != 1:
+            raise ValueError(f"Expected 1 argument, got {len(input_args)}")
+        result = _http_request(input_args[0], self._force_cache)
+        if isinstance(result, pa.ChunkedArray):
+            result = result.combine_chunks()
+        return result
+
+
+def _http_request(arg: pa.Array, force_cache: bool) -> pa.Array:
+    client = _HttpClient()
+
+    if isinstance(arg, pa.StructArray):
+        # We assume a vector of requests, but with potentially many arguments
+        return pa.array(
+            [
+                _response_dict(
+                    client.request(
+                        req.get("method", "GET").upper(),
+                        req["url"],
+                        headers=req.get("headers", {}),
+                        extensions={"force_cache": force_cache},
+                    )
+                )
+                for req in arg.to_pylist()
+            ],
+            type=HttpGet.RES_DTYPE,
+        )
+
+    raise TypeError(f"Unsupported argument: {arg} ({type(arg)})")
+
+
+def _response_dict(response: httpx.Response) -> dict:
+    if response.status_code != 200:
+        raise ValueError(f"Request failed with status {response.status_code}")
+    return {
+        "bytes": response.read(),
+        "status": response.status_code,
+        "headers": dict(response.headers),
+    }
+
+
+class _HttpClient(hishel.CacheClient):
+    _instance: "_HttpClient" = None
+
+    def __new__(cls, *args, **kwargs):
+        if not cls._instance:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self):
+        super().__init__(storage=hishel.FileStorage(base_path=APP_DIR / "http.cache", ttl=3600))

spiral/expressions/io.py

@@ -0,0 +1,100 @@
+import tarfile
+from io import BytesIO
+
+import pyarrow as pa
+
+from spiral.expressions.base import Expr, ExprLike
+from spiral.expressions.struct import pack
+from spiral.expressions.udf import UDF
+
+
+def read_file(path: ExprLike) -> Expr:
+    """
+    Read file path(s) from disk into a struct with a single field "bytes" containing the file contents.
+
+    Args:
+        path: Expression evaluating to an array of strings representing local disk paths.
+    """
+    to_pack = {"path": path}
+    return FileRead()(pack(to_pack))
+
+
+class FileRead(UDF):
+    RES_DTYPE: pa.DataType = pa.struct(
+        [
+            pa.field("bytes", pa.large_binary()),
+        ]
+    )
+
+    def __init__(self):
+        super().__init__("file.read")
+
+    def return_type(self, *input_types: pa.DataType) -> pa.DataType:
+        return FileRead.RES_DTYPE
+
+    def invoke(self, *input_args: pa.Array) -> pa.Array:
+        if len(input_args) != 1:
+            raise ValueError(f"Expected 1 argument, got {len(input_args)}")
+        arg = input_args[0]
+
+        res = []
+        for req in arg:
+            with open(req["path"].as_py(), "rb") as f:
+                res.append({"bytes": f.read()})
+
+        return pa.array(res, type=FileRead.RES_DTYPE)
+
+
+def read_tar(path: ExprLike = None, bytes_: ExprLike = None) -> "Expr":
+    # Untar a vector of paths / byte arrays representing tarballs.
+    if path is None and bytes_ is None:
+        raise ValueError("Expected either path or bytes_ to be provided")
+    to_pack = {}
+    if path is not None:
+        to_pack["path"] = path
+    if bytes_ is not None:
+        to_pack["bytes"] = bytes_
+    return TarRead()(pack(to_pack))
+
+
+class TarRead(UDF):
+    RES_DTYPE = pa.list_(
+        pa.struct(
+            [
+                pa.field("name", pa.string()),
+                pa.field("bytes", pa.large_binary()),
+            ]
+        )
+    )
+
+    def __init__(self):
+        super().__init__("tar.read")
+
+    def return_type(self, *input_types: pa.DataType) -> pa.DataType:
+        return TarRead.RES_DTYPE
+
+    def invoke(self, *input_args: pa.Array) -> pa.Array:
+        if len(input_args) != 1:
+            raise ValueError(f"Expected 1 argument, got {len(input_args)}")
+        arg = input_args[0]
+
+        res = []
+        for req in arg:
+            if "path" in req:
+                kwargs = {"name": req["path"].as_py()}
+            elif "bytes" in req:
+                kwargs = {"fileobj": BytesIO(req["bytes"].as_py())}
+            else:
+                raise ValueError("Expected path or bytes_ to be provided")
+
+            files = []
+            with tarfile.open(**kwargs) as f:
+                for m in f.getmembers():
+                    m: tarfile.TarInfo
+                    if m.type == tarfile.DIRTYPE:
+                        continue
+                    # TODO(ngates): skip other types too maybe? Why are we even skipping directories?
+                    files.append({"name": m.name, "bytes": f.extractfile(m).read()})
+            res.append(files)
+
+        return pa.array(res, type=TarRead.RES_DTYPE)

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.spql.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/refs.py

@@ -0,0 +1,44 @@
+from typing import TYPE_CHECKING
+
+from spiral.expressions.base import Expr, ExprLike
+
+if TYPE_CHECKING:
+    from spiral import Table
+
+
+def ref(expr: ExprLike, field: str | None = None) -> Expr:
+    """Store binary values as references. This expression can only be used on write.
+
+    It is often better to store large cell values, such as bytes columns, that aren't used in filter expressions as
+    references. This enables more efficient scan pruning. Many of the Spiral's cell pushdown expressions work
+    over references.
+
+    Args:
+        expr: The expression to store as a reference.
+        field: If the expr evaluates into struct, the field name of that struct that should be referenced.
+            If `None`, the expr must evaluate into a type that supports referencing.
+    """
+    from spiral import _lib
+    from spiral.expressions import lift
+
+    expr = lift(expr)
+    return Expr(_lib.spql.expr.ref(expr.__expr__, field))
+
+
+def deref(expr: ExprLike, field: str | None = None, table: "Table" = None) -> Expr:
+    """De-reference referenced values.
+
+    See `ref` for more information on Spiral's reference values. This expression is used to de-reference referenced
+    column back into their original form, e.g. binary.
+
+    Args:
+        expr: The expression to de-reference.
+        field: If the expr evaluates into struct, the field name of that struct that should be de-referenced.
+            If `None`, the expr must evaluate into a reference type.
+        table (optional): The table to de-reference from, if not available in input expression.
+    """
+    from spiral import _lib
+    from spiral.expressions import lift
+
+    expr = lift(expr)
+    return Expr(_lib.spql.expr.deref(expr.__expr__, field, table._table if table is not None else None))

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,57 @@
+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.spql.expr.struct.getitem(expr.__expr__, field))
+
+
+def pack(fields: dict[str, ExprLike]) -> 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.spql.expr.struct.pack(list(fields.keys()), [se.lift(expr).__expr__ for expr in fields.values()]))
+
+
+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.spql.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.spql.expr.struct.select(expr.__expr__, names, exclude))