pyspiral 0.7.9__cp312-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl → 0.7.11__cp312-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

{pyspiral-0.7.9.dist-info → pyspiral-0.7.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyspiral
-Version: 0.7.9
+Version: 0.7.11
 Classifier: Intended Audience :: Science/Research
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python

{pyspiral-0.7.9.dist-info → pyspiral-0.7.11.dist-info}/RECORD

@@ -1,8 +1,8 @@
-pyspiral-0.7.9.dist-info/METADATA,sha256=La_MkKyTCJ_qrYYT2eGLaPYkaf9OSpNa36nYOfMaN_4,1874
-pyspiral-0.7.9.dist-info/WHEEL,sha256=0ecHyBdkJfSXYIVmWsPh7S-4h4fSrB4FlXhlnIu9c_A,130
-pyspiral-0.7.9.dist-info/entry_points.txt,sha256=R96Y3FpYX6XbQu9qMPfUTgiCcf4qM9OBQQZTDdBkZwA,74
+pyspiral-0.7.11.dist-info/METADATA,sha256=Xok1DuwVcaDpgNdpQWSbppaxmOjqWVawQOSgRcIG96o,1875
+pyspiral-0.7.11.dist-info/WHEEL,sha256=0ecHyBdkJfSXYIVmWsPh7S-4h4fSrB4FlXhlnIu9c_A,130
+pyspiral-0.7.11.dist-info/entry_points.txt,sha256=R96Y3FpYX6XbQu9qMPfUTgiCcf4qM9OBQQZTDdBkZwA,74
 spiral/__init__.py,sha256=PwaYBWFBtB7cYi7peMmhk_Lm5XzjRoLwOtLbUhc1ZDo,1449
-spiral/_lib.abi3.so,sha256=lLD_ZKRMNTADiM96MnkJT-0rIeZFDSzx_yiNOk0UtrU,61819128
+spiral/_lib.abi3.so,sha256=NeGs3eqjki0bSsmAPuiieCjmLSxbUL_bzPCZ7bBk3GY,61740360
 spiral/adbc.py,sha256=7IxfWIeQN-fh0W5OdN_PP2x3pzQYg6ZUOLsHg3jktqw,14842
 spiral/api/__init__.py,sha256=ULBlVq3PnfNOO6T5naE_ULmmii-83--qTuN2PpAUQN0,2241
 spiral/api/admin.py,sha256=A1iVR1XYJSObZivPAD5UzmPuMgupXc9kaHNYYa_kwfs,585
@@ -35,11 +35,11 @@ spiral/cli/telemetry.py,sha256=Uxo1Q1FkKJ6n6QNGOUmL3j_pRRWRx0qWIhoP-U9BuR0,589
 spiral/cli/text.py,sha256=DlWGe4JrkdERAiqyITNpk91Wqb63Re99rNYlIFsIamc,4031
 spiral/cli/types.py,sha256=XYzo1GgX7dBBItoBSrHI4vO5C2lLmS2sktb-2GnGH3E,1362
 spiral/cli/workloads.py,sha256=2_SLfQTFN6y73R9H0i9dk8VIOVagKxSxOpHXC56yptY,2015
-spiral/client.py,sha256=53dVv8wxYMmozUfR8MVcUufKGqdVIdb0yZ0gchczBoQ,6426
+spiral/client.py,sha256=fkgauF1UhmQGkcWmKjBT1SVQ-4g9K6qVtjLBMpfHYKU,6479
 spiral/core/__init__.pyi,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 spiral/core/_tools/__init__.pyi,sha256=b2KLfTOQ67pjfbYt07o0IGiTu5o2bZw69lllV8v0Dps,143
 spiral/core/authn/__init__.pyi,sha256=z_GWyIS62fuiYQrYO8hzw4W8oGaiciqS1u5qtAt54VY,769
-spiral/core/client/__init__.pyi,sha256=YgDM-MoIt3J-QKxvsfs5gRiaTBtOA6TphbNBIAnrFCw,6956
+spiral/core/client/__init__.pyi,sha256=WfHQ8xL6x6NGq32h9ERMCdYEwBVtEWWbmk3A5icE0l4,7153
 spiral/core/expr/__init__.pyi,sha256=3HSKjkotiEkxBvGBALXEBIie0JiyI9bCpehwA3nMQkU,571
 spiral/core/expr/images/__init__.pyi,sha256=wnE_wZXq7a4iqTg3SVm-ssxGw1WQZyk5dGOPaP4Btko,73
 spiral/core/expr/list_/__init__.pyi,sha256=Q_9c87eIQfZbqlaw_rq3fvs93YEsW7K5VYk6VZ4g6mU,126
@@ -60,7 +60,7 @@ spiral/debug/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 spiral/debug/manifests.py,sha256=7f1O3ba9mrA5nXpOF9cEIQuUAteP5wiBkFy_diQJ7No,3216
 spiral/debug/metrics.py,sha256=XdRDcjggtsLNGCAjam6IxG9072pz_d2C8iLApNRFUtk,2044
 spiral/debug/scan.py,sha256=UEm_aRnql5pwDPTpZgakMLNjlzkKL4RurBFFqH_BLAQ,9526
-spiral/enrichment.py,sha256=j8CzWJqLRq0Zko-qz3NzicsFOAvdzQSRJ58vdmqxPsk,9879
+spiral/enrichment.py,sha256=hFK9wlWWiGkp3ST18WN7QTKwMc7EO7peUwW3AAiC0Gc,11650
 spiral/expressions/__init__.py,sha256=ZsD8g7vB0G7xy19GUiH4m79kw7KEkTQRwJl5Gn1cgtw,8049
 spiral/expressions/base.py,sha256=PvhJkcUSsPSIaxirHVzM9zlqyBXiaiia1HXohXdOmL4,5377
 spiral/expressions/file.py,sha256=7D9jIENJcoT0KFharBLkzK9dZgO4DYn5K_KCt0twefg,518
@@ -96,14 +96,14 @@ spiral/protogen/util.py,sha256=smnvVo6nYH3FfDm9jqhNLaXz4bbTBaQezHQDCTvZyiQ,1486
 spiral/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 spiral/scan.py,sha256=8IFuhqxzVGdXPW8uWWTFKiC38mTLpwkBZvK1YHqtbHM,12574
 spiral/server.py,sha256=ztBmB5lBnUz-smQxR_tC8AI5SOhz17wH0MI3GuzDUdM,600
-spiral/settings.py,sha256=sUhMMBCXaPvUYztN_gztD9TjeUYJwVeEcJrq4FLy6M0,3232
+spiral/settings.py,sha256=QXVyu5uU-sClICyoAtnbOtJuOzbhCAH9rIvbjAi0I5c,3274
 spiral/snapshot.py,sha256=cTobi5jtiANxalGA-isokQHblNmXGtuUvgUGGNVybsI,1555
 spiral/streaming_/__init__.py,sha256=s7MlW2ERsuZmZGExLFL6RcZon2e0tNBocBg5ANgki7k,61
 spiral/streaming_/reader.py,sha256=tl_lC9xgh1-QFhsZn4xQT7It3PVTzHCEUT2BG2dWBRQ,4166
 spiral/streaming_/stream.py,sha256=DM1hBDHnWm1ZFKZ-hZ4zxeSXITcUI6kWzwdJZvywI8o,5915
 spiral/substrait_.py,sha256=AKeOD4KIXvz2J4TYxnIneOiHddtBIyOhuNxVO_uH0eg,12592
-spiral/table.py,sha256=p95AYv6b7e14F3t7j-B-r45k9CtG84ngikdlAhh9WxA,12260
+spiral/table.py,sha256=g6y1iV2Esk6kYhsIJfW7Wje0EGs7jgA0bVuZaSzNX1A,12068
 spiral/text_index.py,sha256=FQ9rgIEGLSJryS9lFdMhKtPFey18BXoWbPXyvZPJJ04,442
 spiral/transaction.py,sha256=bI5oqBAmPMSF0yOOYcPfGbV37Xc1-_V-wQNKw1xOlTA,4136
 spiral/types_.py,sha256=W_jyO7F6rpPiH69jhgSgV7OxQZbOlb1Ho3InpKUP6Eo,155
-pyspiral-0.7.9.dist-info/RECORD,,
+pyspiral-0.7.11.dist-info/RECORD,,

spiral/client.py

@@ -6,7 +6,7 @@ import pyarrow as pa
 
 from spiral.api import SpiralAPI
 from spiral.api.projects import CreateProjectRequest, CreateProjectResponse
-from spiral.core.client import Internal
+from spiral.core.client import Internal, KeyColumns
 from spiral.core.client import Spiral as CoreSpiral
 from spiral.datetime_ import timestamp_micros
 from spiral.expressions import ExprLike
@@ -102,6 +102,7 @@ class Spiral:
         *projections: ExprLike,
         where: ExprLike | None = None,
         asof: datetime | int | None = None,
+        _key_columns: KeyColumns | None = None,
     ) -> Scan:
         """Starts a read transaction on the Spiral.
 
@@ -125,9 +126,7 @@ class Spiral:
         return Scan(
             self,
             self.core.scan(
-                projection.__expr__,
-                filter=where.__expr__ if where else None,
-                asof=asof,
+                projection.__expr__, filter=where.__expr__ if where else None, asof=asof, key_columns=_key_columns
             ),
         )
 

spiral/core/client/__init__.pyi

@@ -1,3 +1,4 @@
+from enum import Enum
 from typing import Any, Literal
 
 import pyarrow as pa
@@ -7,6 +8,12 @@ from spiral.core.table import ColumnGroupState, KeyRange, KeySpaceState, Scan, S
 from spiral.core.table.spec import ColumnGroup, Schema
 from spiral.expressions import Expr
 
+# Only for typing, the actual definition is in Rust.
+class KeyColumns(Enum):
+    IfProjected = 0
+    Included = 1
+    Only = 2
+
 class Spiral:
     """A client for Spiral database"""
     def __init__(
@@ -26,6 +33,7 @@ class Spiral:
         projection: Expr,
         filter: Expr | None = None,
         asof: int | None = None,
+        key_columns: KeyColumns | None = None,
     ) -> Scan:
         """Construct a table scan."""
         ...

spiral/enrichment.py

@@ -1,9 +1,11 @@
+from __future__ import annotations
+
 import dataclasses
 import logging
 from functools import partial
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
 
-from spiral.core.client import Shard
+from spiral.core.client import KeyColumns, Shard
 from spiral.core.table import KeyRange
 from spiral.core.table.spec import Key, Operation
 from spiral.expressions import Expr
@@ -27,7 +29,7 @@ class Enrichment:
 
     def __init__(
         self,
-        table: "Table",
+        table: Table,
         projection: Expr,
         where: Expr | None,
     ):
@@ -36,7 +38,7 @@ class Enrichment:
         self._where = where
 
     @property
-    def table(self) -> "Table":
+    def table(self) -> Table:
         """The table to write back into."""
         return self._table
 
@@ -50,8 +52,8 @@ class Enrichment:
         """The filter expression."""
         return self._where
 
-    def _scan(self) -> "Scan":
-        return self._table.spiral.scan(self._projection, where=self._where)
+    def _scan(self) -> Scan:
+        return self._table.spiral.scan(self._projection, where=self._where, _key_columns=KeyColumns.Included)
 
     def apply(
         self, *, batch_readahead: int | None = None, partition_size_bytes: int | None = None, tx_dump: str | None = None
@@ -87,11 +89,11 @@ class Enrichment:
     def apply_dask(
         self,
         *,
-        index: Optional["KeySpaceIndex"] = None,
+        index: KeySpaceIndex | None = None,
         partition_size_bytes: int | None = None,
         tx_dump: str | None = None,
         checkpoint_dump: str | None = None,
-        client: Optional["dask.distributed.Client"] = None,
+        client: dask.distributed.Client | None = None,
         **kwargs,
     ) -> None:
         """Use distributed Dask to apply the enrichment. Requires `dask[distributed]` to be installed.
@@ -146,10 +148,21 @@ class Enrichment:
         if shards is None:
             shards = plan_scan.shards()
 
+        # TODO(marko): This is temporary workaround. Passing token is a bad idea.
+        #   Token can expire during long-running enrichments.
+        #   Maybe if device code is used, we can pass something.
+        token = self._table.spiral.config.authn.token()
+        if token is None:
+            raise ValueError("Spiral client is not authenticated.")
+        settings_dict = self._table.spiral.config.model_dump()
+        spiraldb_dict = settings_dict.get("spiraldb", {})
+        spiraldb_dict["token"] = token.expose_secret()
+        settings_dict["spiraldb"] = spiraldb_dict
+
         # Partially bind the enrichment function.
         _compute = partial(
             _enrichment_task,
-            settings_dict=self._table.spiral.config.model_dump(),
+            settings_dict=settings_dict,
             state_json=plan_scan.core.plan_state().to_json(),
             output_table_id=self._table.table_id,
             partition_size_bytes=partition_size_bytes,
@@ -160,19 +173,51 @@ class Enrichment:
         logger.info(f"Applying enrichment with {len(shards)} shards. Follow progress at {client.dashboard_link}")
 
         failed_ranges = []
-        for result in client.gather(enrichments):
-            result: EnrichmentTaskResult
-
-            if result.error is not None:
-                logger.error(f"Enrichment task failed for range {result.key_range}: {result.error}")
-                failed_ranges.append(result.key_range)
-                continue
-
-            tx.include(result.ops)
+        try:
+            for result, shard in zip(client.gather(enrichments), shards):
+                result: EnrichmentTaskResult
+
+                if result.error is not None:
+                    logger.error(f"Enrichment task failed for range {shard.key_range}: {result.error}")
+                    failed_ranges.append(shard.key_range)
+                    continue
+
+                tx.include(result.ops)
+        except Exception as e:
+            # If not incremental, re-raise the exception.
+            if checkpoint_dump is None:
+                raise e
+
+            # Handle worker failures (e.g., KilledWorker from Dask)
+            from dask.distributed import KilledWorker
+
+            if isinstance(e, KilledWorker):
+                logger.error(f"Dask worker was killed during enrichment: {e}")
+
+            # Try to gather partial results and mark remaining tasks as failed
+            for future, shard in zip(enrichments, shards):
+                if future.done() and not future.exception():
+                    try:
+                        result = future.result()
+
+                        if result.error is not None:
+                            logger.error(f"Enrichment task failed for range {shard.key_range}: {result.error}")
+                            failed_ranges.append(shard.key_range)
+                            continue
+
+                        tx.include(result.ops)
+                    except Exception:
+                        # Task failed or incomplete, add to failed ranges
+                        failed_ranges.append(shard.key_range)
+                else:
+                    # Task didn't complete, add to failed ranges
+                    failed_ranges.append(shard.key_range)
 
         # Dump checkpoint of failed ranges, if any.
         if checkpoint_dump is not None:
-            logger.info(f"Dumping checkpoint with {len(failed_ranges)} failed ranges to {checkpoint_dump}.")
+            logger.info(
+                f"Dumping checkpoint with failed {len(failed_ranges)}/{len(shards)} ranges to {checkpoint_dump}."
+            )
             _checkpoint_dump_key_ranges(checkpoint_dump, failed_ranges)
 
         if tx.is_empty():
@@ -212,7 +257,6 @@ def _checkpoint_dump_key_ranges(checkpoint_dump: str, ranges: list[KeyRange]):
 
 @dataclasses.dataclass
 class EnrichmentTaskResult:
-    key_range: KeyRange
     ops: list[Operation]
     error: str | None = None
 
@@ -220,12 +264,9 @@ class EnrichmentTaskResult:
         return {
             "ops": [op.to_json() for op in self.ops],
             "error": self.error,
-            "begin": bytes(self.key_range.begin),
-            "end": bytes(self.key_range.end),
         }
 
     def __setstate__(self, state):
-        self.key_range = KeyRange(begin=Key(state["begin"]), end=Key(state["end"]))
         self.ops = [Operation.from_json(op_json) for op_json in state["ops"]]
         self.error = state["error"]
 
@@ -250,17 +291,16 @@ def _enrichment_task(
     state = ScanState.from_json(state_json)
     task_scan = Scan(sp, sp.core.load_scan(state))
     table = sp.table(output_table_id)
-
     task_tx = table.txn()
 
     try:
         task_tx.writeback(task_scan, key_range=shard.key_range, partition_size_bytes=partition_size_bytes)
-        return EnrichmentTaskResult(key_range=shard.key_range, ops=task_tx.take())
+        return EnrichmentTaskResult(ops=task_tx.take())
     except Exception as e:
         task_tx.abort()
 
         if incremental:
-            return EnrichmentTaskResult(key_range=shard.key_range, ops=[], error=str(e))
+            return EnrichmentTaskResult(ops=[], error=str(e))
 
         logger.error(f"Enrichment task failed for shard {shard}: {e}")
         raise e

spiral/settings.py

@@ -28,9 +28,11 @@ PACKAGE_NAME = "pyspiral"
 
 
 def validate_token(v, handler: ValidatorFunctionWrapHandler):
-    if not isinstance(v, str):
-        raise ValueError("Token value (SPIRAL__SPIRALDB__TOKEN) must be a string")
-    return Token(v)
+    if isinstance(v, Token):
+        return v
+    if isinstance(v, str):
+        return Token(v)
+    raise ValueError("Token value (SPIRAL__SPIRALDB__TOKEN) must be a string")
 
 
 TokenType = Annotated[

spiral/table.py

@@ -134,9 +134,7 @@ class Table(Expr):
         """
         from spiral import expressions as se
 
-        # TODO(marko): This shouldn't need to happen. We should be able to read keys from writeback scan.
-        # Include key columns in the projection.
-        projection = se.merge(self.select(*self.key_schema.names), *projections)
+        projection = se.merge(*projections)
         if where is not None:
             where = se.lift(where)