pyspiral 0.2.5__cp310-abi3-macosx_11_0_arm64.whl → 0.4.0__cp310-abi3-macosx_11_0_arm64.whl

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

spiral/expressions/tiff.py

@@ -1,42 +1,44 @@
 import numpy as np
 import pyarrow as pa
 
-from spiral.expressions.base import ExprLike
+from spiral.expressions.base import Expr, ExprLike
 from spiral.expressions.udf import RefUDF
 
+_TIFF_RES_DTYPE: pa.DataType = pa.struct(
+    [
+        pa.field("pixels", pa.large_binary()),
+        pa.field("height", pa.uint32()),
+        pa.field("width", pa.uint32()),
+        pa.field("channels", pa.uint8()),
+        pa.field("channel_bit_depth", pa.uint8()),
+    ]
+)
+
 
 def read(
     expr: ExprLike,
-    indexes: ExprLike | int | list[int] | None = None,
+    indexes: ExprLike | int | None = None,
     window: ExprLike | tuple[tuple[int, int], tuple[int, int]] | None = None,
     boundless: ExprLike | bool | None = None,
-):
+) -> Expr:
     """
     Read referenced cell in a `TIFF` format. Requires `rasterio` to be installed.
 
     Args:
         expr: The referenced `TIFF` bytes.
-        indexes: The band indexes to read. Defaults to first band. The first dimension of the result's `shape` field
-            is either 1 or the number of indexes.
+        indexes: The band indexes to read. Defaults to all.
         window: The window to read. In format (row_range_tuple, col_range_tuple). Defaults to full window.
         boundless: If `True`, windows that extend beyond the dataset's extent
             are permitted and partially or completely filled arrays will be returned as appropriate.
 
     Returns:
-        An array where each element is a NumPy array represented as a struct with fields:
-            bytes: Array bytes with type `pa.large_binary()`.
-            shape: Array shape with type `pa.list_(pa.uint32(), 3)`.
-            dtype: String representation of NumPy dtype with type `pa.string()`.
-
-    Example:
-        A way to get the i-th element in the result as NumPy array:
-
-        ```
-        array: np.ndarray = np.frombuffer(
-            result["bytes"][i].as_py(),
-            dtype=np.dtype(result["dtype"][i].as_py()),
-        ).reshape(tuple(result["shape"][i].as_py()))
-        ```
+        An array where each element is a decoded image with fields:
+            pixels: bytes of shape (channels, width, height).
+            width: Width of the image with type `pa.uint32()`.
+            height: Height of the image with type `pa.uint32()`.
+            channels: Number of channels of the image with type `pa.uint8()`.
+                If `indexes` is not None, this is the length of `indexes` or 1 if `indexes` is an int.
+            channel_bit_depth: Bit depth of the channel with type `pa.uint8()`.
     """
     try:
         import rasterio  # noqa: F401
@@ -46,55 +48,42 @@ def read(
     return TiffReadUDF()(expr, indexes, window, boundless)
 
 
-def crop(
+def select(
     expr: ExprLike,
-    shape: ExprLike,
-):
+    shape: ExprLike | dict,
+    indexes: ExprLike | int | None = None,
+) -> Expr:
     """
-    Crop shapes out of the referenced cell in a `TIFF` format. Requires `rasterio` to be installed.
+    Select the shape out of the referenced cell in a `TIFF` format. Requires `rasterio` to be installed.
 
     Args:
         expr: The referenced `TIFF` bytes.
         shape: [GeoJSON-like](https://geojson.org/) shape.
+        indexes: The band indexes to read. Defaults to all.
 
     Returns:
-        An array where each element is a NumPy array represented as a struct with fields:
-            bytes: Array bytes with type `pa.large_binary()`.
-            shape: Array shape with type `pa.list_(pa.uint32(), 3)`.
-            dtype: String representation of NumPy dtype with type `pa.string()`.
-
-    Example:
-        A way to get the i-th element in the result as NumPy array:
-
-        ```
-        array: np.ndarray = np.frombuffer(
-            result["bytes"][i].as_py(),
-            dtype=np.dtype(result["dtype"][i].as_py()),
-        ).reshape(tuple(result["shape"][i].as_py()))
-        ```
+        An array where each element is a decoded image with fields:
+            pixels: bytes of shape (len(indexes) or 1, width, height).
+            width: Width of the image with type `pa.uint32()`.
+            height: Height of the image with type `pa.uint32()`.
+            channels: Number of channels of the image with type `pa.uint8()`.
+                If `indexes` is not None, this is the length of `indexes` or 1 if `indexes` is an int.
+            channel_bit_depth: Bit depth of the channel with type `pa.uint8()`.
     """
     try:
         import rasterio  # noqa: F401
     except ImportError:
-        raise ImportError("`rasterio` is required for tiff.crop")
+        raise ImportError("`rasterio` is required for tiff.select")
 
-    return TiffCropUDF()(expr, shape)
+    return TiffSelectUDF()(expr, shape, indexes)
 
 
 class TiffReadUDF(RefUDF):
-    RES_DTYPE: pa.DataType = pa.struct(
-        [
-            pa.field("bytes", pa.large_binary()),
-            pa.field("shape", pa.list_(pa.uint32(), 3)),
-            pa.field("dtype", pa.string()),
-        ]
-    )
-
     def __init__(self):
         super().__init__("tiff.read")
 
     def return_type(self, *input_types: pa.DataType) -> pa.DataType:
-        return TiffReadUDF.RES_DTYPE
+        return _TIFF_RES_DTYPE
 
     def invoke(self, fp, *input_args: pa.Array) -> pa.Array:
         try:
@@ -130,65 +119,76 @@ class TiffReadUDF(RefUDF):
             #   This matters more if we want to rewrite this function to work with multiple inputs at once, in which
             #   case we should first consider using Rust GDAL bindings - I believe rasterio uses GDAL under the hood.
             result: np.ndarray = src.read(indexes=indexes, window=window)
-            return pa.array(
-                [
-                    {
-                        "bytes": result.tobytes(),
-                        "shape": list(result.shape),
-                        "dtype": str(result.dtype),
-                    }
-                ],
-                type=TiffReadUDF.RES_DTYPE,
-            )
-
-
-class TiffCropUDF(RefUDF):
-    RES_DTYPE: pa.DataType = pa.struct(
-        [
-            pa.field("bytes", pa.large_binary()),
-            pa.field("shape", pa.list_(pa.uint32()), 3),
-            pa.field("dtype", pa.string()),
-        ]
-    )
+            return _return_result(result, indexes)
+
 
+class TiffSelectUDF(RefUDF):
     def __init__(self):
-        super().__init__("tiff.crop")
+        super().__init__("tiff.select")
 
     def return_type(self, *input_types: pa.DataType) -> pa.DataType:
-        return TiffCropUDF.RES_DTYPE
+        return _TIFF_RES_DTYPE
 
     def invoke(self, fp, *input_args: pa.Array) -> pa.Array:
         try:
             import rasterio
         except ImportError:
-            raise ImportError("`rasterio` is required for tiff.crop")
+            raise ImportError("`rasterio` is required for tiff.select")
 
-        from rasterio.mask import mask as rio_mask
+        from rasterio.mask import raster_geometry_mask
 
-        if len(input_args) != 2:
-            raise ValueError("tiff.crop expects exactly 2 arguments: expr, shape")
+        if len(input_args) != 3:
+            raise ValueError("tiff.select expects exactly 3 arguments: expr, shape, indexes")
 
-        _, shape = input_args
+        _, shape, indexes = input_args
 
         shape = shape[0].as_py()
         if shape is None:
-            raise ValueError("tiff.crop expects shape to be a GeoJSON-like shape")
+            raise ValueError("tiff.select expects shape to be a GeoJSON-like shape")
+
+        indexes = indexes[0].as_py()
+        if indexes is not None and not isinstance(indexes, int) and not isinstance(indexes, list):
+            raise ValueError(f"tiff.select expects indexes to be None or an int or a list, got {indexes}")
 
         opener = _VsiOpener(fp)
         with rasterio.open("ref", opener=opener) as src:
             src: rasterio.DatasetReader
-            result, _ = rio_mask(src, shapes=[shape], crop=True)
-            result: np.ndarray
-            return pa.array(
-                [
-                    {
-                        "bytes": result.tobytes(),
-                        "shape": list(result.shape),
-                        "dtype": str(result.dtype),
-                    }
-                ],
-                type=TiffCropUDF.RES_DTYPE,
-            )
+
+            shape_mask, _, window = raster_geometry_mask(src, [shape], crop=True)
+            out_shape = (src.count,) + shape_mask.shape
+
+            result: np.ndarray = src.read(window=window, indexes=indexes, out_shape=out_shape, masked=True)
+            return _return_result(result, indexes)
+
+
+def _return_result(result: np.ndarray, indexes) -> pa.Array:
+    channels = result.shape[0]
+    if indexes is None:
+        pass
+    elif isinstance(indexes, int):
+        assert channels == 1, f"Expected 1 channel, got {channels}"
+    else:
+        assert channels == len(indexes), f"Expected {len(indexes)} channels, got {channels}"
+
+    if result.dtype == np.uint8:
+        channel_bit_depth = 8
+    elif result.dtype == np.uint16:
+        channel_bit_depth = 16
+    else:
+        raise ValueError(f"Unsupported bit width: {result.dtype}")
+
+    return pa.array(
+        [
+            {
+                "pixels": result.tobytes(),
+                "height": result.shape[1],
+                "width": result.shape[2],
+                "channels": channels,
+                "channel_bit_depth": channel_bit_depth,
+            }
+        ],
+        type=_TIFF_RES_DTYPE,
+    )
 
 
 class _VsiOpener:

spiral/expressions/udf.py

@@ -25,7 +25,7 @@ class UDF(BaseUDF):
     """A User-Defined Function (UDF)."""
 
     def __init__(self, name: str):
-        super().__init__(_lib.spql.expr.udf.create(name, return_type=self.return_type, invoke=self.invoke))
+        super().__init__(_lib.expr.udf.create(name, return_type=self.return_type, invoke=self.invoke))
 
     @abc.abstractmethod
     def invoke(self, *input_args: pa.Array) -> pa.Array: ...
@@ -35,10 +35,10 @@ class RefUDF(BaseUDF):
     """A UDF over a single ref cell, and therefore can access the file object."""
 
     def __init__(self, name: str):
-        super().__init__(_lib.spql.expr.udf.create(name, return_type=self.return_type, invoke=self.invoke, scope="ref"))
+        super().__init__(_lib.expr.udf.create(name, return_type=self.return_type, invoke=self.invoke, scope="ref"))
 
     @abc.abstractmethod
-    def invoke(self, fp, *input_args: pa.Array) -> pa.Array:
+    def invoke(self, fp: _lib.FileObject, *input_args: pa.Array) -> pa.Array:
         """Invoke the UDF with the given arguments.
 
         NOTE: The first argument is always the ref cell. All array input args will be sliced to the appropriate row.

spiral/iceberg/__init__.py

@@ -0,0 +1,3 @@
+from spiral.iceberg.client import Iceberg
+
+__all__ = ["Iceberg"]

spiral/iceberg/client.py

@@ -0,0 +1,33 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pyiceberg.catalog import Catalog
+
+    from spiral.client import Spiral
+
+
+class Iceberg:
+    """
+    Apache Iceberg is a powerful open-source table format designed for high-performance data lakes.
+    Iceberg brings reliability, scalability, and advanced features like time travel, schema evolution,
+    and ACID transactions to your warehouse.
+    """
+
+    def __init__(self, spiral: "Spiral", *, project_id: str | None = None):
+        self._spiral = spiral
+        self._project_id = project_id
+
+        self._api = self._spiral.config.api
+
+    def catalog(self) -> "Catalog":
+        """Open the Iceberg catalog."""
+        from pyiceberg.catalog import load_catalog
+
+        return load_catalog(
+            "default",
+            **{
+                "type": "rest",
+                "uri": self._spiral.config.spiraldb.uri + "/iceberg",
+                "token": self._spiral.config.authn.token().expose_secret(),
+            },
+        )

spiral/indexes/__init__.py

@@ -0,0 +1,5 @@
+from spiral.indexes.client import Indexes
+from spiral.indexes.index import TextIndex
+from spiral.indexes.scan import SearchScan
+
+__all__ = ["Indexes", "SearchScan", "TextIndex"]

spiral/indexes/client.py

@@ -0,0 +1,137 @@
+import datetime
+
+from spiral.api import SpiralAPI
+from spiral.api.projects import TextIndexResource
+from spiral.core.client import Spiral as CoreSpiral
+from spiral.expressions.base import ExprLike
+from spiral.indexes.index import TextIndex
+from spiral.indexes.scan import SearchScan
+from spiral.types_ import Uri
+
+
+class Indexes:
+    def __init__(self, api: SpiralAPI, spiral: CoreSpiral, *, project_id: str | None = None):
+        self._api = api
+        self._spiral = spiral
+        self._project_id = project_id
+
+    def index(self, identifier: str) -> TextIndex:
+        """Returns the index with the given identifier."""
+        project_id, index_name = self._parse_identifier(identifier)
+        if project_id is None:
+            raise ValueError("Must provide a fully qualified index identifier.")
+
+        res = list(self._api.project.list_text_indexes(project_id, name=index_name))
+        if len(res) == 0:
+            raise ValueError(f"Index not found: {project_id}.{index_name}")
+        res = res[0]
+
+        return TextIndex(self, self._spiral.get_text_index(res.id), index_name)
+
+    def list_indexes(self) -> list[TextIndexResource]:
+        project_id = self._project_id
+        if project_id is None:
+            raise ValueError("Must provide a project ID to list indexes.")
+        return list(self._api.project.list_text_indexes(project_id))
+
+    def create_text_index(
+        self,
+        identifier: str,
+        # At least one projection is required. All projections must reference the same table!
+        # NOTE(marko): Indexes are currently independent of tables.
+        #   That will likely change with the new root resource such as documents.
+        *projections: ExprLike,
+        where: ExprLike | None = None,
+        root_uri: Uri | None = None,
+        exist_ok: bool = False,
+    ) -> TextIndex:
+        """Creates a text index over the table projection.
+
+        See `se.text.field` for how to create and configure indexable fields.
+
+        Args:
+            identifier: The index identifier, in the form `project.index` or `index`.
+            projections: At least one projection expression is required.
+                All projections must reference the same table.
+            where: An optional filter expression to apply to the index.
+            root_uri: The root URI for the index.
+            exist_ok: If True, do not raise an error if the index already exists.
+        """
+        from spiral import expressions as se
+
+        project_id, index_name = self._parse_identifier(identifier)
+        if project_id is None:
+            raise ValueError("Must provide a fully qualified index identifier.")
+
+        if not projections:
+            raise ValueError("At least one projection is required.")
+        projection = se.merge(*projections)
+        if where is not None:
+            where = se.lift(where)
+
+        core_index = self._spiral.create_text_index(
+            project_id,
+            index_name,
+            projection.__expr__,
+            where.__expr__ if where else None,
+            root_uri=root_uri,
+            # TODO(marko): Validate that if an index exists, it's the same?
+            exist_ok=exist_ok,
+        )
+
+        return TextIndex(self, core_index, index_name)
+
+    def _parse_identifier(self, identifier: str) -> tuple[str | None, str]:
+        parts = identifier.split(".")
+        if len(parts) == 1:
+            return self._project_id, parts[0]
+        elif len(parts) == 2:
+            return parts[0], parts[1]
+        else:
+            raise ValueError(f"Invalid index identifier: {identifier}")
+
+    def search(
+        self,
+        *rank_by: ExprLike,
+        where: ExprLike | None = None,
+        top_k: int = 10,
+        # Do not refresh the index if freshness does not exceed the freshness window.
+        # NOTE(marko): The current implementation fails the query if the index is stale.
+        freshness_window: datetime.timedelta | None = None,
+    ) -> SearchScan:
+        """Queries the index with the given rank by and where clauses.
+
+        Rank by expressions are combined for scoring.
+            See `se.text.find` and `se.text.boost` for scoring expressions.
+        The `where` expression is used to filter the results.
+            It must return a boolean value and use only conjunctions (ANDs). Expressions in where statement
+            are considered either a `must` or `must_not` clause in search terminology.
+
+        Args:
+            rank_by: At least one rank by expression is required.
+                These expressions are used to score the results.
+            where: An optional filter expression to apply to the index.
+                It must return a boolean value and use only conjunctions (ANDs).
+            top_k: The number of top results to return.
+            freshness_window: If provided, the index will not be refreshed if its freshness does not exceed this window.
+        """
+        from spiral import expressions as se
+
+        if not rank_by:
+            raise ValueError("At least one rank by expression is required.")
+        rank_by = se.or_(*rank_by)
+        if where is not None:
+            where = se.lift(where)
+
+        if freshness_window is None:
+            freshness_window = datetime.timedelta(seconds=0)
+        freshness_window_s = int(freshness_window.total_seconds())
+
+        return SearchScan(
+            self._spiral.open_search_scan(
+                rank_by.__expr__,
+                top_k=top_k,
+                freshness_window_s=freshness_window_s,
+                filter=where.__expr__ if where else None,
+            )
+        )

spiral/indexes/index.py

@@ -0,0 +1,34 @@
+import datetime
+from typing import TYPE_CHECKING
+
+from spiral.core.index import TextIndex as CoreTextIndex
+from spiral.expressions import Expr
+
+if TYPE_CHECKING:
+    from spiral.indexes import Indexes
+
+
+class TextIndex(Expr):
+    def __init__(self, indexes: "Indexes", index: CoreTextIndex, name: str):
+        super().__init__(index.__expr__)
+
+        self._indexes = indexes
+        self._index = index
+        self._name = name
+
+    @property
+    def client(self) -> "Indexes":
+        return self._indexes
+
+    @property
+    def index_id(self) -> str:
+        return self._index.id
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def status(self) -> (str, datetime.timedelta | None):
+        """Fetch the status of the index. If status is ready, returns the staleness of the index."""
+        status = self._index.status()
+        return status.status, datetime.timedelta(seconds=status.staleness_s) if status.staleness_s is not None else None

spiral/indexes/scan.py

@@ -0,0 +1,22 @@
+import pyarrow as pa
+
+from spiral.core.index import SearchScan as CoreSearchScan
+from spiral.settings import CI, DEV
+
+
+class SearchScan:
+    def __init__(self, scan: CoreSearchScan):
+        self._scan = scan
+
+    def to_record_batches(self) -> pa.RecordBatchReader:
+        """Read all results as a record batch reader."""
+        return self._scan.to_record_batches()
+
+    def to_table(self) -> pa.Table:
+        """Read all results as a table."""
+        # NOTE: Evaluates fully on Rust side which improved debuggability.
+        if DEV and not CI:
+            rb = self._scan.to_record_batch()
+            return pa.Table.from_batches([rb])
+
+        return self.to_record_batches().read_all()