pyspiral 0.3.1__cp310-abi3-macosx_11_0_arm64.whl → 0.4.1__cp310-abi3-macosx_11_0_arm64.whl

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

spiral/project.py

@@ -1,26 +1,18 @@
-from typing import TYPE_CHECKING, Any
-
-import pyarrow as pa
-
-from spiral import Table
-from spiral.api.tables import CreateTable, FindTable
-from spiral.core.core import Table as CoreTable
-from spiral.core.metastore import PyMetastore
-from spiral.core.spec import Schema
-from spiral.types_ import Uri
+from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from spiral.catalog import Spiral
+    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_db: "Spiral", id: str, name: str | None = None):
-        self._spiral_db = spiral_db
+    def __init__(self, spiral: "Spiral", id: str, name: str | None = None):
+        self._spiral = spiral
         self._id = id
         self._name = name
 
-        self._api = self._spiral_db.config.api
-
     def __str__(self):
         return self._id
 
@@ -35,103 +27,20 @@ class Project:
     def name(self) -> str:
         return self._name or self._id
 
-    def list_table_names(self) -> list[(str, str)]:
-        """List tuples of (dataset, table) names in the project."""
-        return [(t.dataset, t.table) for t in self._api.table.list(FindTable.Request(project_id=self.id))]
-
-    def list_tables(self) -> list[Table]:
-        """List tables in the project."""
-        return [
-            Table(
-                CoreTable(
-                    PyMetastore.http(
-                        table_id=t.id,
-                        root_uri=t.metadata.root_uri,
-                        key_schema=Schema.from_arrow(t.metadata.key_schema),
-                        base_url=self._api.base_url + "/metastore/",
-                        token_provider=self._spiral_db.config.authn.token,
-                    ),
-                ),
-                name=f"{self.id}.{t.dataset}.{t.table}",
-            )
-            for t in self._api.table.list(FindTable.Request(project_id=self.id))
-        ]
-
-    def create_table(
-        self,
-        identifier: str,
-        *,
-        key_schema: pa.Schema | Any,
-        uri: Uri | None = None,
-        exist_ok: bool = False,
-    ) -> Table:
-        """Create a new table in the project."""
-        dataset, table = self._parse_identifier(identifier)
-
-        if not isinstance(key_schema, pa.Schema):
-            key_schema = pa.schema(key_schema)
-
-        res = self._api.table.create(
-            CreateTable.Request(
-                project_id=self.id,
-                dataset=dataset,
-                table=table,
-                key_schema=key_schema,
-                root_uri=uri,
-                exist_ok=exist_ok,
-            )
-        )
-
-        # Must have the same schema as provided, even if the table already exists.
-        expected_key_schema = res.table.metadata.key_schema
-        if key_schema != expected_key_schema:
-            raise ValueError(f"Table already exists with different key schema: {expected_key_schema} != {key_schema}")
-        if uri and res.table.metadata.root_uri != uri:
-            raise ValueError(f"Table already exists with different root URI: {res.table.metadata.root_uri} != {uri}")
-
-        # Set up a metastore backed by SpiralDB
-        metastore = PyMetastore.http(
-            table_id=res.table.id,
-            root_uri=res.table.metadata.root_uri,
-            key_schema=Schema.from_arrow(res.table.metadata.key_schema),
-            base_url=self._api.base_url + "/metastore/",
-            token_provider=self._spiral_db.config.authn.token,
-        )
-
-        return Table(CoreTable(metastore), name=f"{self.id}.{res.table.dataset}.{res.table.table}")
+    @property
+    def tables(self) -> "Tables":
+        from spiral.tables import Tables
 
-    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_identifier(identifier)
+        return Tables(self._spiral._api, self._spiral._core, project_id=self.id)
 
-        # TODO(ngates): why does the client _need_ this information? Can we defer it?
-        res = self._api.table.find(
-            FindTable.Request(
-                project_id=self.id,
-                dataset=dataset,
-                table=table,
-            )
-        )
-        if res.table is None:
-            raise ValueError(f"Table not found: {self.id}.{dataset}.{table}")
+    @property
+    def indexes(self) -> "Indexes":
+        from spiral.indexes.client import Indexes
 
-        # Set up a metastore backed by SpiralDB
-        metastore = PyMetastore.http(
-            table_id=res.table.id,
-            root_uri=res.table.metadata.root_uri,
-            key_schema=Schema.from_arrow(res.table.metadata.key_schema),
-            base_url=self._api.base_url + "/metastore/",
-            token_provider=self._spiral_db.config.authn.token,
-        )
+        return Indexes(self._spiral._api, self._spiral._core, project_id=self._id)
 
-        return Table(CoreTable(metastore), name=f"{self.id}.{res.table.dataset}.{res.table.table}")
+    @property
+    def iceberg(self) -> "Iceberg":
+        from spiral.iceberg import Iceberg
 
-    @staticmethod
-    def _parse_identifier(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: {identifier}")
+        return Iceberg(self._spiral, project_id=self._id)

spiral/{proto → protogen}/_/scandal/__init__.py

@@ -30,11 +30,31 @@ class Source(betterproto.Message):
     parquet: "MetadataParquet" = betterproto.message_field(10, group="metadata")
 
 
+@dataclass(eq=False, repr=False)
+class Metadata(betterproto.Message):
+    pass
+
+
+@dataclass(eq=False, repr=False)
+class MetadataParquet(betterproto.Message):
+    pass
+
+
 @dataclass(eq=False, repr=False)
 class Sink(betterproto.Message):
     url: str = betterproto.string_field(1)
 
 
+@dataclass(eq=False, repr=False)
+class Connectivity(betterproto.Message):
+    """
+    Information about the client's perceived connectivity to a FileSystem.
+    """
+
+    unreachable: bool = betterproto.bool_field(1)
+    round_trip_time_us: Optional[int] = betterproto.int32_field(2, optional=True)
+
+
 @dataclass(eq=False, repr=False)
 class Fetch(betterproto.Message):
     """Let's make "fetch" happen."""
@@ -52,15 +72,10 @@ class FetchRequest(betterproto.Message):
      spfs://<fsid>/path?token=<jwt&gt URI.
     """
 
-    headers: Dict[str, str] = betterproto.map_field(
-        2, betterproto.TYPE_STRING, betterproto.TYPE_STRING
-    )
-    """Custom headers to sign into the request."""
-
-    connectivity: "Connectivity" = betterproto.message_field(3)
+    connectivity: "Connectivity" = betterproto.message_field(2)
     """Declares whether the client has connectivity to the FileSystem."""
 
-    accepts: List[str] = betterproto.string_field(4)
+    accepts: List[str] = betterproto.string_field(3)
     """Declares the MIME types the client can read directly."""
 
 
@@ -81,10 +96,7 @@ class Put(betterproto.Message):
 @dataclass(eq=False, repr=False)
 class PutRequest(betterproto.Message):
     uri: str = betterproto.string_field(1)
-    headers: Dict[str, str] = betterproto.map_field(
-        2, betterproto.TYPE_STRING, betterproto.TYPE_STRING
-    )
-    connectivity: "Connectivity" = betterproto.message_field(3)
+    connectivity: "Connectivity" = betterproto.message_field(2)
 
 
 @dataclass(eq=False, repr=False)
@@ -92,64 +104,6 @@ class PutResponse(betterproto.Message):
     sinks: List["Sink"] = betterproto.message_field(1)
 
 
-@dataclass(eq=False, repr=False)
-class Head(betterproto.Message):
-    pass
-
-
-@dataclass(eq=False, repr=False)
-class HeadRequest(betterproto.Message):
-    uri: str = betterproto.string_field(1)
-    headers: Dict[str, str] = betterproto.map_field(
-        2, betterproto.TYPE_STRING, betterproto.TYPE_STRING
-    )
-
-
-@dataclass(eq=False, repr=False)
-class HeadResponse(betterproto.Message):
-    url: str = betterproto.string_field(1)
-    """Returns signed URL to head the resource."""
-
-
-@dataclass(eq=False, repr=False)
-class Delete(betterproto.Message):
-    pass
-
-
-@dataclass(eq=False, repr=False)
-class DeleteRequest(betterproto.Message):
-    uri: str = betterproto.string_field(1)
-    headers: Dict[str, str] = betterproto.map_field(
-        2, betterproto.TYPE_STRING, betterproto.TYPE_STRING
-    )
-
-
-@dataclass(eq=False, repr=False)
-class DeleteResponse(betterproto.Message):
-    url: str = betterproto.string_field(1)
-    """Returns signed URL to delete the resource."""
-
-
-@dataclass(eq=False, repr=False)
-class Connectivity(betterproto.Message):
-    """
-    Information about the client's perceived connectivity to a FileSystem.
-    """
-
-    unreachable: bool = betterproto.bool_field(1)
-    round_trip_time_us: Optional[int] = betterproto.int32_field(2, optional=True)
-
-
-@dataclass(eq=False, repr=False)
-class Metadata(betterproto.Message):
-    pass
-
-
-@dataclass(eq=False, repr=False)
-class MetadataParquet(betterproto.Message):
-    pass
-
-
 class ScandalServiceStub(betterproto.ServiceStub):
     async def fetch(
         self,
@@ -185,40 +139,6 @@ class ScandalServiceStub(betterproto.ServiceStub):
             metadata=metadata,
         )
 
-    async def head(
-        self,
-        head_request: "HeadRequest",
-        *,
-        timeout: Optional[float] = None,
-        deadline: Optional["Deadline"] = None,
-        metadata: Optional["MetadataLike"] = None
-    ) -> "HeadResponse":
-        return await self._unary_unary(
-            "/scandal.ScandalService/Head",
-            head_request,
-            HeadResponse,
-            timeout=timeout,
-            deadline=deadline,
-            metadata=metadata,
-        )
-
-    async def delete(
-        self,
-        delete_request: "DeleteRequest",
-        *,
-        timeout: Optional[float] = None,
-        deadline: Optional["Deadline"] = None,
-        metadata: Optional["MetadataLike"] = None
-    ) -> "DeleteResponse":
-        return await self._unary_unary(
-            "/scandal.ScandalService/Delete",
-            delete_request,
-            DeleteResponse,
-            timeout=timeout,
-            deadline=deadline,
-            metadata=metadata,
-        )
-
 
 class ScandalServiceBase(ServiceBase):
     async def fetch(self, fetch_request: "FetchRequest") -> "FetchResponse":
@@ -227,12 +147,6 @@ class ScandalServiceBase(ServiceBase):
     async def put(self, put_request: "PutRequest") -> "PutResponse":
         raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
 
-    async def head(self, head_request: "HeadRequest") -> "HeadResponse":
-        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
-
-    async def delete(self, delete_request: "DeleteRequest") -> "DeleteResponse":
-        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
-
     async def __rpc_fetch(
         self, stream: "grpclib.server.Stream[FetchRequest, FetchResponse]"
     ) -> None:
@@ -247,20 +161,6 @@ class ScandalServiceBase(ServiceBase):
         response = await self.put(request)
         await stream.send_message(response)
 
-    async def __rpc_head(
-        self, stream: "grpclib.server.Stream[HeadRequest, HeadResponse]"
-    ) -> None:
-        request = await stream.recv_message()
-        response = await self.head(request)
-        await stream.send_message(response)
-
-    async def __rpc_delete(
-        self, stream: "grpclib.server.Stream[DeleteRequest, DeleteResponse]"
-    ) -> None:
-        request = await stream.recv_message()
-        response = await self.delete(request)
-        await stream.send_message(response)
-
     def __mapping__(self) -> Dict[str, grpclib.const.Handler]:
         return {
             "/scandal.ScandalService/Fetch": grpclib.const.Handler(
@@ -275,16 +175,4 @@ class ScandalServiceBase(ServiceBase):
                 PutRequest,
                 PutResponse,
             ),
-            "/scandal.ScandalService/Head": grpclib.const.Handler(
-                self.__rpc_head,
-                grpclib.const.Cardinality.UNARY_UNARY,
-                HeadRequest,
-                HeadResponse,
-            ),
-            "/scandal.ScandalService/Delete": grpclib.const.Handler(
-                self.__rpc_delete,
-                grpclib.const.Cardinality.UNARY_UNARY,
-                DeleteRequest,
-                DeleteResponse,
-            ),
         }

spiral/protogen/_/spiral/table/__init__.py

@@ -0,0 +1,22 @@
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# sources: spiral/table/statistics.proto
+# plugin: python-betterproto
+# This file has been @generated
+
+from dataclasses import dataclass
+
+import betterproto
+
+
+@dataclass(eq=False, repr=False)
+class ApproximateSetMembership(betterproto.Message):
+    bloom_filter: "BloomFilter" = betterproto.message_field(
+        2, group="membership_strategy"
+    )
+
+
+@dataclass(eq=False, repr=False)
+class BloomFilter(betterproto.Message):
+    bit_vec: bytes = betterproto.bytes_field(1)
+    bitmap_bits: int = betterproto.uint64_field(2)
+    k_num: int = betterproto.uint32_field(3)