pyspiral 0.4.4__cp310-abi3-macosx_11_0_arm64.whl → 0.6.0__cp310-abi3-macosx_11_0_arm64.whl

spiral/{tables/debug → debug}/scan.py

@@ -1,12 +1,12 @@
 from datetime import datetime
 
-from spiral.core.table import TableScan
+from spiral.core.table import Scan
 from spiral.core.table.manifests import FragmentFile, FragmentManifest
 from spiral.core.table.spec import Key
 from spiral.types_ import Timestamp
 
 
-def show_scan(scan: TableScan):
+def show_scan(scan: Scan):
     """Displays a scan in a way that is useful for debugging."""
     table_ids = scan.table_ids()
     if len(table_ids) > 1:
@@ -14,18 +14,18 @@ def show_scan(scan: TableScan):
     table_id = table_ids[0]
     column_groups = scan.column_groups()
 
-    splits = scan.split()
-    key_space_scan = scan.key_space_scan(table_id)
+    splits = scan.splits()
+    key_space_state = scan.key_space_state(table_id)
 
     # Collect all key bounds from all manifests. This makes sure all visualizations are aligned.
     key_points = set()
-    key_space_manifest = key_space_scan.manifest
+    key_space_manifest = key_space_state.manifest
     for i in range(len(key_space_manifest)):
         fragment_file = key_space_manifest[i]
         key_points.add(fragment_file.key_extent.min)
         key_points.add(fragment_file.key_extent.max)
     for cg in column_groups:
-        cg_scan = scan.column_group_scan(cg)
+        cg_scan = scan.column_group_state(cg)
         cg_manifest = cg_scan.manifest
         for i in range(len(cg_manifest)):
             fragment_file = cg_manifest[i]
@@ -39,7 +39,7 @@ def show_scan(scan: TableScan):
 
     show_manifest(key_space_manifest, scope="Key space", key_points=key_points, splits=splits)
     for cg in scan.column_groups():
-        cg_scan = scan.column_group_scan(cg)
+        cg_scan = scan.column_group_state(cg)
         # Skip table id from the start of the column group.
         show_manifest(cg_scan.manifest, scope=".".join(cg.path[1:]), key_points=key_points, splits=splits)
 

spiral/expressions/base.py

@@ -132,7 +132,7 @@ class Expr:
                 packed = packed.select(exclude=exclude)
             return packed
 
-        if not paths:
+        if not paths and not exclude:
             return self
 
         return se.select(self, names=list(paths), exclude=exclude)
@@ -145,5 +145,5 @@ class Expr:
         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
+ScalarLike: TypeAlias = bool | int | float | str | list["ScalarLike"] | datetime.datetime | None
+ExprLike: TypeAlias = Expr | dict["ExprLike", "ExprLike"] | ScalarLike

spiral/expressions/udf.py

@@ -38,7 +38,7 @@ class RefUDF(BaseUDF):
         super().__init__(_lib.expr.udf.create(name, return_type=self.return_type, invoke=self.invoke, scope="ref"))
 
     @abc.abstractmethod
-    def invoke(self, fp: _lib.FileObject, *input_args: pa.Array) -> pa.Array:
+    def invoke(self, fp, *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/client.py → iceberg.py}

@@ -13,10 +13,8 @@ class Iceberg:
     and ACID transactions to your warehouse.
     """
 
-    def __init__(self, spiral: "Spiral", *, project_id: str | None = None):
+    def __init__(self, spiral: "Spiral"):
         self._spiral = spiral
-        self._project_id = project_id
-
         self._api = self._spiral.config.api
 
     def catalog(self) -> "Catalog":

spiral/key_space_index.py

@@ -0,0 +1,44 @@
+from spiral.core.client import KeySpaceIndex as CoreKeySpaceIndex
+from spiral.expressions import Expr
+from spiral.types_ import Timestamp
+
+
+class KeySpaceIndex:
+    """
+    KeysIndex represents an optionally materialized key space, defined by a projection and a filter over a table.
+    It can be used to efficiently and precisely shard the table for parallel processing or distributed training.
+
+    An index is defined by:
+    - A granularity that defines the target size of key ranges in the index.
+        IMPORTANT: Actual key ranges may be smaller, but will not exceed twice the granularity.
+    - A projection expression that defines which columns are included in the resulting key space.
+    - An optional filter expression that defines which rows are included in the index.
+    """
+
+    def __init__(self, core: CoreKeySpaceIndex, *, name: str | None = None):
+        self.core = core
+        self._name = name
+
+    @property
+    def index_id(self) -> str:
+        return self.core.id
+
+    @property
+    def table_id(self) -> str:
+        return self.core.table_id
+
+    @property
+    def name(self) -> str:
+        return self._name or self.index_id
+
+    @property
+    def asof(self) -> Timestamp:
+        return self.core.asof
+
+    @property
+    def projection(self) -> Expr:
+        return Expr(self.core.projection)
+
+    @property
+    def filter(self) -> Expr | None:
+        return Expr(self.core.filter) if self.core.filter is not None else None

spiral/project.py

@@ -1,16 +1,23 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
+
+import pyarrow as pa
+
+from spiral.api.projects import KeySpaceIndexResource, TableResource, TextIndexResource
+from spiral.core.table.spec import Schema
+from spiral.expressions import ExprLike
+from spiral.key_space_index import KeySpaceIndex
+from spiral.table import Table
+from spiral.text_index import TextIndex
+from spiral.types_ import Uri
 
 if TYPE_CHECKING:
     from spiral.client import Spiral
-    from spiral.iceberg import Iceberg
-    from spiral.indexes import Indexes
-    from spiral.tables import Tables
 
 
 class Project:
-    def __init__(self, spiral: "Spiral", id: str, name: str | None = None):
+    def __init__(self, spiral: "Spiral", project_id: str, name: str | None = None):
         self._spiral = spiral
-        self._id = id
+        self._id = project_id
         self._name = name
 
     def __str__(self):
@@ -27,20 +34,166 @@ class Project:
     def name(self) -> str:
         return self._name or self._id
 
-    @property
-    def tables(self) -> "Tables":
-        from spiral.tables import Tables
+    def list_tables(self) -> list[TableResource]:
+        return list(self._spiral.api.project.list_tables(self._id))
 
-        return Tables(self._spiral._api, self._spiral._core, project_id=self.id)
+    def list_text_indexes(self) -> list[TextIndexResource]:
+        return list(self._spiral.api.project.list_text_indexes(self._id))
 
-    @property
-    def indexes(self) -> "Indexes":
-        from spiral.indexes.client import Indexes
+    def list_key_space_indexes(self) -> list[KeySpaceIndexResource]:
+        return list(self._spiral.api.project.list_key_space_indexes(self._id))
 
-        return Indexes(self._spiral._api, self._spiral._core, project_id=self._id)
+    def table(self, identifier: str) -> Table:
+        """Open a table with a `dataset.table` identifier, or `table` name using the `default` dataset."""
+        dataset, table = self._parse_table_identifier(identifier)
 
-    @property
-    def iceberg(self) -> "Iceberg":
-        from spiral.iceberg import Iceberg
+        res = list(self._spiral.api.project.list_tables(project_id=self._id, dataset=dataset, table=table))
+        if len(res) == 0:
+            raise ValueError(f"Table not found: {self._id}.{dataset}.{table}")
+        res = res[0]
+
+        return Table(
+            self._spiral, self._spiral._core.table(res.id), identifier=f"{res.project_id}.{res.dataset}.{res.table}"
+        )
+
+    def create_table(
+        self,
+        identifier: str,
+        *,
+        key_schema: pa.Schema | Any,
+        root_uri: Uri | None = None,
+        exist_ok: bool = False,
+    ) -> Table:
+        """Create a new table in the project.
+
+        Args:
+            identifier: The table identifier, in the form `dataset.table` or `table`.
+            key_schema: The schema of the table's keys.
+            root_uri: The root URI for the table.
+            exist_ok: If True, do not raise an error if the table already exists.
+        """
+        dataset, table = self._parse_table_identifier(identifier)
+
+        if not isinstance(key_schema, pa.Schema):
+            key_schema = pa.schema(key_schema)
+        key_schema = Schema.from_arrow(key_schema)
+
+        core_table = self._spiral._core.create_table(
+            project_id=self._id,
+            dataset=dataset,
+            table=table,
+            key_schema=key_schema,
+            root_uri=root_uri,
+            exist_ok=exist_ok,
+        )
+
+        return Table(self._spiral, core_table, identifier=f"{self._id}.{dataset}.{table}")
+
+    def _parse_table_identifier(self, identifier: str) -> tuple[str, str]:
+        parts = identifier.split(".")
+        if len(parts) == 1:
+            return "default", parts[0]
+        elif len(parts) == 2:
+            return parts[0], parts[1]
+        else:
+            raise ValueError(f"Invalid table identifier: {self._id}.{identifier}")
+
+    def text_index(self, name: str) -> TextIndex:
+        """Returns the index with the given name."""
+        res = list(self._spiral.api.project.list_text_indexes(project_id=self._id, name=name))
+        if len(res) == 0:
+            raise ValueError(f"Index not found: {name}")
+        res = res[0]
+
+        return TextIndex(self._spiral._core.text_index(res.id), name=name)
+
+    def create_text_index(
+        self,
+        name: str,
+        *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:
+            name: The index name. Must be unique within the project.
+            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
+
+        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._core.create_text_index(
+            project_id=self._id,
+            name=name,
+            projection=projection.__expr__,
+            filter=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(core_index, name=name)
+
+    def key_space_index(self, name: str) -> KeySpaceIndex:
+        """Returns the index with the given name."""
+        res = list(self._spiral.api.project.list_key_space_indexes(project_id=self._id, name=name))
+        if len(res) == 0:
+            raise ValueError(f"Index not found: {name}")
+        res = res[0]
+
+        return KeySpaceIndex(self._spiral._core.key_space_index(res.id), name=name)
+
+    def create_key_space_index(
+        self,
+        name: str,
+        granularity: int,
+        *projections: ExprLike,
+        where: ExprLike | None = None,
+        root_uri: Uri | None = None,
+        exist_ok: bool = False,
+    ) -> KeySpaceIndex:
+        """Creates a key space index over the table projection.
+
+        Args:
+            name: The index name. Must be unique within the project.
+            granularity: The granularity at which to store keys, i.e. the size of desired key ranges.
+                The key ranges will not be greater than 2x the granularity, but may be smaller.
+            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
+
+        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._core.create_key_space_index(
+            project_id=self._id,
+            name=name,
+            granularity=granularity,
+            projection=projection.__expr__,
+            filter=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 Iceberg(self._spiral, project_id=self._id)
+        return KeySpaceIndex(core_index, name=name)