pyspiral 0.4.0__pp310-pypy310_pp73-macosx_10_12_x86_64.whl

spiral/expressions/__init__.py

@@ -0,0 +1,245 @@
+import builtins
+import functools
+import operator
+import warnings
+from typing import Any
+
+import pyarrow as pa
+
+from spiral import _lib, arrow_
+
+from . import http as http
+from . import io as io
+from . import list_ as list
+from . import mp4 as mp4
+from . import png as png
+from . import qoi as qoi
+from . import refs as refs
+from . import str_ as str
+from . import struct as struct
+from . import text as text
+from . import tiff as tiff
+from .base import Expr, ExprLike, NativeExpr
+
+__all__ = [
+    "Expr",
+    "add",
+    "and_",
+    "deref",
+    "divide",
+    "eq",
+    "getitem",
+    "gt",
+    "gte",
+    "http",
+    "io",
+    "is_not_null",
+    "is_null",
+    "lift",
+    "list",
+    "lt",
+    "lte",
+    "merge",
+    "modulo",
+    "multiply",
+    "negate",
+    "neq",
+    "not_",
+    "or_",
+    "pack",
+    "keyed",
+    "ref",
+    "refs",
+    "scalar",
+    "select",
+    "str",
+    "struct",
+    "subtract",
+    "tiff",
+    "xor",
+    "png",
+    "qoi",
+    "mp4",
+    "text",
+]
+
+# Inline some of the struct expressions since they're so common
+getitem = struct.getitem
+merge = struct.merge
+pack = struct.pack
+select = struct.select
+ref = refs.ref
+deref = refs.deref
+
+
+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):
+        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 key(name: builtins.str) -> Expr:
+    """Create a variable expression referencing a key column.
+
+    Args:
+        name: variable name
+    """
+    return Expr(_lib.expr.keyed(name))
+
+
+def keyed(name: builtins.str, dtype: pa.DataType) -> Expr:
+    """Create a variable expression referencing a column in the key table.
+
+    Key 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 key table.
+    """
+    return Expr(_lib.expr.keyed(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,149 @@
+import datetime
+from typing import TypeAlias
+
+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) -> "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.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 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.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.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`
+    ...