pyspiral 0.7.7__cp312-abi3-manylinux_2_28_aarch64.whl → 0.7.9__cp312-abi3-manylinux_2_28_aarch64.whl

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

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

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

@@ -1,8 +1,8 @@
-pyspiral-0.7.7.dist-info/METADATA,sha256=5KbbTylag9ZfdffltxdGDcIK2fdzoMZHkQ3ASWi0ZRM,1874
-pyspiral-0.7.7.dist-info/WHEEL,sha256=I5JYpyYzeAl2SOerY_wvkm-HJti0rDQc6zMeJs35MpM,108
-pyspiral-0.7.7.dist-info/entry_points.txt,sha256=R96Y3FpYX6XbQu9qMPfUTgiCcf4qM9OBQQZTDdBkZwA,74
+pyspiral-0.7.9.dist-info/METADATA,sha256=La_MkKyTCJ_qrYYT2eGLaPYkaf9OSpNa36nYOfMaN_4,1874
+pyspiral-0.7.9.dist-info/WHEEL,sha256=I5JYpyYzeAl2SOerY_wvkm-HJti0rDQc6zMeJs35MpM,108
+pyspiral-0.7.9.dist-info/entry_points.txt,sha256=R96Y3FpYX6XbQu9qMPfUTgiCcf4qM9OBQQZTDdBkZwA,74
 spiral/__init__.py,sha256=PwaYBWFBtB7cYi7peMmhk_Lm5XzjRoLwOtLbUhc1ZDo,1449
-spiral/_lib.abi3.so,sha256=0_bgv3p4eChTxdrSfd6v1KZexnYVjwi8SdUXdlfru8I,61716472
+spiral/_lib.abi3.so,sha256=mAjWINjLl5rrbGEDln3oZW4RJUV77O48x4mVDJuYr2g,61819128
 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
@@ -30,7 +30,7 @@ spiral/cli/orgs.py,sha256=fmOuLxpeIFfKqePRi292Gv9k-EF5pPn_tbKd2BLl2Ig,2869
 spiral/cli/printer.py,sha256=aosc763hDFgoXJGkiANmNyO3kAsecAS1JWgjEhn8GCM,1784
 spiral/cli/projects.py,sha256=1M1nGrBT-t0aY9RV5Cnmzy7YrhIvmHwdkpa3y9j8rG8,5756
 spiral/cli/state.py,sha256=10wTIVQ0SJkY67Z6-KQ1LFlt3aVIPmZhoHFdTwp4kNA,130
-spiral/cli/tables.py,sha256=gp7v8K-6EUTa0P7ZcHXD9gYPPpp4nHABloo_elvUrBw,7757
+spiral/cli/tables.py,sha256=6vt6EBGt7I9b0kAQ6sQORbmWiKbRdH4ubQYjjuNBXEg,6900
 spiral/cli/telemetry.py,sha256=Uxo1Q1FkKJ6n6QNGOUmL3j_pRRWRx0qWIhoP-U9BuR0,589
 spiral/cli/text.py,sha256=DlWGe4JrkdERAiqyITNpk91Wqb63Re99rNYlIFsIamc,4031
 spiral/cli/types.py,sha256=XYzo1GgX7dBBItoBSrHI4vO5C2lLmS2sktb-2GnGH3E,1362
@@ -39,7 +39,7 @@ spiral/client.py,sha256=53dVv8wxYMmozUfR8MVcUufKGqdVIdb0yZ0gchczBoQ,6426
 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=MmuIzluwaxAFFamr_J8iT15qP8AkA4smkz4zrF-yGnM,6993
+spiral/core/client/__init__.pyi,sha256=YgDM-MoIt3J-QKxvsfs5gRiaTBtOA6TphbNBIAnrFCw,6956
 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
@@ -49,18 +49,18 @@ spiral/core/expr/struct_/__init__.pyi,sha256=MXckd98eV_x3X0RhEWvlkA3DcDXRtLs5pNn
 spiral/core/expr/text/__init__.pyi,sha256=ed83n1xcsGY7_QDhMmJGnSQ20UrJFXcdv1AveSEcS1c,175
 spiral/core/expr/udf/__init__.pyi,sha256=zsZs081KVhY3-1JidqTkWMW81Qd_ScoTGZvasIhIK-4,358
 spiral/core/expr/video/__init__.pyi,sha256=nQJEcSsigZuRpMjkI_O4EEtMK_n2zRvorcL_KEeD5vU,95
-spiral/core/table/__init__.pyi,sha256=zcf4GripPZtiwh6uHkPgVyDij1g2nYL1DogN83z5ISU,4037
+spiral/core/table/__init__.pyi,sha256=h84QDg6hLuPcmRpavx5zOZM77ZCi2-YwIlrrUZJp1sE,4374
 spiral/core/table/manifests/__init__.pyi,sha256=eVfDpmhYSjafIvvALqAkZe5baN3Y1HpKpxYEbjwd4gQ,1043
 spiral/core/table/metastore/__init__.pyi,sha256=rc3u9MwEKRvL2kxOc8lBorddFRnM8o_o1frqtae86a4,1697
 spiral/core/table/spec/__init__.pyi,sha256=fVuc2j3uoTdWfYNm720OfUIgrLYw9fRwj44maI5bgdY,5709
 spiral/dataloader.py,sha256=W9siY4BF4p_rwTTSS4KgsaQsPLxxza6XmQhrdBzzMJ8,10592
-spiral/dataset.py,sha256=PMLoXnXuEUciP6-NXqTmQLXu0UIH7OcC4-iZtY_iuO8,7973
+spiral/dataset.py,sha256=S8pdiBXIhwMxQiJYgF7UI_8HkN7pZO798UzlO1LNXy4,8409
 spiral/datetime_.py,sha256=elXaUWtZuuLVcu9E0aXnvYRPB9XWqZbLDToozQYQYjU,950
 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=w0MrZ93wDuvS4sazw_8dPmnhzkQ4SAU5A1CGE7WF-F8,7046
+spiral/enrichment.py,sha256=j8CzWJqLRq0Zko-qz3NzicsFOAvdzQSRJ58vdmqxPsk,9879
 spiral/expressions/__init__.py,sha256=ZsD8g7vB0G7xy19GUiH4m79kw7KEkTQRwJl5Gn1cgtw,8049
 spiral/expressions/base.py,sha256=PvhJkcUSsPSIaxirHVzM9zlqyBXiaiia1HXohXdOmL4,5377
 spiral/expressions/file.py,sha256=7D9jIENJcoT0KFharBLkzK9dZgO4DYn5K_KCt0twefg,518
@@ -106,4 +106,4 @@ spiral/table.py,sha256=p95AYv6b7e14F3t7j-B-r45k9CtG84ngikdlAhh9WxA,12260
 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.7.dist-info/RECORD,,
+pyspiral-0.7.9.dist-info/RECORD,,

spiral/cli/tables.py

@@ -1,6 +1,5 @@
-import datetime
 from collections.abc import Callable
-from typing import Annotated, Literal
+from typing import Annotated
 
 import questionary
 import rich
@@ -61,6 +60,21 @@ def ls(
     CONSOLE.print(rich_table)
 
 
+@app.command(help="Show the leading rows of the table.")
+def head(
+    project: ProjectArg,
+    table: Annotated[str | None, Option(help="Table name.")] = None,
+    dataset: Annotated[str | None, Option(help="Dataset name.")] = None,
+    n: Annotated[int, Option("-n", help="Maximum number of rows to show. Defaults to 10.")] = 10,
+):
+    import polars as pl
+
+    _, t = get_table(project, table, dataset)
+
+    with pl.Config(tbl_rows=-1):
+        CONSOLE.print(t.to_polars().limit(n).collect())
+
+
 def validate_non_empty_str(text: str) -> bool | str:
     if len(text) > 0:
         return True
@@ -137,40 +151,9 @@ def flush(
     project: ProjectArg,
     table: Annotated[str | None, Option(help="Table name.")] = None,
     dataset: Annotated[str | None, Option(help="Dataset name.")] = None,
-    keep: Annotated[
-        Literal["1h", "2h", "4h"] | None,
-        Option(help="Duration string that indicates how much WAL to keep. Defaults to 24h."),
-    ] = None,
-    full: Annotated[bool, Option(help="Flush full Write-Ahead-Log.")] = False,
 ):
-    # TODO(marko): Use some human-readable duration parsing library.
-    duration = None
-    if keep is not None:
-        if full:
-            raise ValueError("Cannot specify both --keep and --full")
-        match keep:
-            case "1h":
-                duration = datetime.timedelta(hours=1)
-            case "2h":
-                duration = datetime.timedelta(hours=2)
-            case "4h":
-                duration = datetime.timedelta(hours=4)
-            case _:
-                raise ValueError(f"Invalid duration string: {keep}")
-
-    if full:
-        # Warn and wait for confirmation.
-        ERR_CONSOLE.print("[bold yellow]Warning: All currently open transaction will fail to commit.[/bold yellow]")
-        if not questionary.confirm("Are you sure you want to continue?", default=False).ask():  # pyright: ignore[reportAny]
-            ERR_CONSOLE.print("Aborting.")
-            raise typer.Exit(1)
-
-        duration = datetime.timedelta(hours=0)
-
-    keep_latest_s = int(duration.total_seconds()) if duration is not None else None
-
     identifier, t = get_table(project, table, dataset)
-    state.spiral.internal.flush_wal(t.core, keep_latest_s=keep_latest_s)  # pyright: ignore[reportPrivateUsage]
+    state.spiral.internal.flush_wal(t.core)  # pyright: ignore[reportPrivateUsage]
     CONSOLE.print(f"Flushed WAL for table {identifier} in project {project}.")
 
 

spiral/core/client/__init__.pyi

@@ -175,7 +175,7 @@ class ShuffleConfig:
     ): ...
 
 class Internal:
-    def flush_wal(self, table: Table, *, keep_latest_s: int | None = None) -> None:
+    def flush_wal(self, table: Table) -> None:
         """
         Flush the write-ahead log of the table.
         """

spiral/core/table/__init__.pyi

@@ -60,6 +60,13 @@ class ScanState:
 class MaterializablePlan:
     pass
 
+class EvaluatedExecutablePlan:
+    pass
+
+class EvaluatedPlanStream:
+    def __next__(self) -> EvaluatedExecutablePlan: ...
+    def __iter__(self) -> EvaluatedPlanStream: ...
+
 class Scan:
     def key_schema(self) -> Schema: ...
     def schema(self) -> Schema: ...
@@ -90,6 +97,10 @@ class Scan:
         # If `infinite` is True, shards are shuffled after exhausted but not before the first pass.
         # Otherwise, shards are not shuffle and shuffle config is only used for shuffle buffer.
         ...
+
+    def evaluate_analyze(
+        self, key_table: pa.Table | pa.RecordBatch | None = None, batch_readahead: int | None = None
+    ) -> EvaluatedPlanStream: ...
     def metrics(self) -> dict[str, Any]: ...
 
 class KeySpaceState:

spiral/dataset.py

@@ -226,7 +226,16 @@ class TableScanner(ds.Scanner):
 
     def head(self, num_rows: int):
         """Return the first `num_rows` rows of the dataset."""
-        reader = self.to_reader()
+
+        kwargs = {}
+        if num_rows <= 10_000:
+            # We are unlikely to need more than a couple batches
+            kwargs["batch_readahead"] = 1
+            # The progress bar length is the total number of splits in this dataset. We will likely
+            # stop streaming early. As a result, the progress bar is misleading.
+            kwargs["hide_progress_bar"] = True
+
+        reader = self._scan.to_record_batches(key_table=self.key_table, **kwargs)
         batches = []
         row_count = 0
         for batch in reader:

spiral/enrichment.py

@@ -4,14 +4,14 @@ from functools import partial
 from typing import TYPE_CHECKING, Optional
 
 from spiral.core.client import Shard
-from spiral.core.table import Scan
-from spiral.core.table.spec import Operation
+from spiral.core.table import KeyRange
+from spiral.core.table.spec import Key, Operation
 from spiral.expressions import Expr
 
 if TYPE_CHECKING:
     import dask.distributed
 
-    from spiral import KeySpaceIndex, Table
+    from spiral import KeySpaceIndex, Scan, Table
 
 logger = logging.getLogger(__name__)
 
@@ -50,7 +50,7 @@ class Enrichment:
         """The filter expression."""
         return self._where
 
-    def _scan(self) -> Scan:
+    def _scan(self) -> "Scan":
         return self._table.spiral.scan(self._projection, where=self._where)
 
     def apply(
@@ -90,6 +90,7 @@ class Enrichment:
         index: Optional["KeySpaceIndex"] = None,
         partition_size_bytes: int | None = None,
         tx_dump: str | None = None,
+        checkpoint_dump: str | None = None,
         client: Optional["dask.distributed.Client"] = None,
         **kwargs,
     ) -> None:
@@ -109,6 +110,7 @@ class Enrichment:
             partition_size_bytes: The maximum partition size in bytes.
                 If not provided, the default partition size is used.
             tx_dump: Optional path to dump the transaction JSON for debugging.
+            checkpoint_dump: Optional path to dump intermediate checkpoints for incremental progress.
             client: Optional Dask distributed client. If not provided, a new client will be created
             **kwargs: Additional keyword arguments to pass to `dask.distributed.Client`
                 such as `address` to connect to an existing cluster.
@@ -126,11 +128,23 @@ class Enrichment:
         tx = self._table.txn()
         plan_scan = self._scan()
 
-        # Determine the "tasks". Use the index if provided.
-        shards = plan_scan.shards()
-        if index is not None:
+        # Determine the "tasks".
+        shards = None
+        # Use checkpoint, if provided.
+        if checkpoint_dump is not None:
+            checkpoint: list[KeyRange] | None = _checkpoint_load_key_ranges(checkpoint_dump)
+            if checkpoint is None:
+                logger.info(f"No existing checkpoint found at {checkpoint_dump}. Starting from scratch.")
+            else:
+                logger.info(f"Resuming enrichment from checkpoint at {checkpoint_dump} with {len(checkpoint)} ranges.")
+                shards = [Shard(kr, None) for kr in checkpoint]
+        # Fallback to index-based sharding.
+        if shards is None and index is not None:
             # TODO(marko): This will use index's asof automatically.
             shards = self._table.spiral.internal.compute_shards(index.core)
+        # Fallback to default sharding.
+        if shards is None:
+            shards = plan_scan.shards()
 
         # Partially bind the enrichment function.
         _compute = partial(
@@ -139,14 +153,28 @@ class Enrichment:
             state_json=plan_scan.core.plan_state().to_json(),
             output_table_id=self._table.table_id,
             partition_size_bytes=partition_size_bytes,
+            incremental=checkpoint_dump is not None,
         )
         enrichments = client.map(_compute, shards)
 
         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)
 
+        # 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}.")
+            _checkpoint_dump_key_ranges(checkpoint_dump, failed_ranges)
+
         if tx.is_empty():
             logger.warning("Transaction not committed. No rows were read for enrichment.")
             return
@@ -155,20 +183,62 @@ class Enrichment:
         tx.commit(compact=True, tx_dump=tx_dump)
 
 
+def _checkpoint_load_key_ranges(checkpoint_dump: str) -> list[KeyRange] | None:
+    import json
+    import os
+
+    if not os.path.exists(checkpoint_dump):
+        return None
+
+    with open(checkpoint_dump) as f:
+        data = json.load(f)
+        return [
+            KeyRange(begin=Key(bytes.fromhex(r["begin"])), end=Key(bytes.fromhex(r["end"])))
+            for r in data.get("key_ranges", [])
+        ]
+
+
+def _checkpoint_dump_key_ranges(checkpoint_dump: str, ranges: list[KeyRange]):
+    import json
+    import os
+
+    os.makedirs(os.path.dirname(checkpoint_dump), exist_ok=True)
+    with open(checkpoint_dump, "w") as f:
+        json.dump(
+            {"key_ranges": [{"begin": bytes(r.begin).hex(), "end": bytes(r.end).hex()} for r in ranges]},
+            f,
+        )
+
+
 @dataclasses.dataclass
 class EnrichmentTaskResult:
+    key_range: KeyRange
     ops: list[Operation]
+    error: str | None = None
 
     def __getstate__(self):
-        return {"ops": [op.to_json() for op in self.ops]}
+        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"]
 
 
 # NOTE(marko): This function must be picklable!
 def _enrichment_task(
-    shard: Shard, *, settings_dict, state_json, output_table_id, partition_size_bytes: int | None
+    shard: Shard,
+    *,
+    settings_dict,
+    state_json,
+    output_table_id,
+    partition_size_bytes: int | None,
+    incremental: bool,
 ) -> EnrichmentTaskResult:
     # Returns operations that can be included in a transaction.
     from spiral import Scan, Spiral
@@ -182,5 +252,15 @@ def _enrichment_task(
     table = sp.table(output_table_id)
 
     task_tx = table.txn()
-    task_tx.writeback(task_scan, key_range=shard.key_range, partition_size_bytes=partition_size_bytes)
-    return EnrichmentTaskResult(ops=task_tx.take())
+
+    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())
+    except Exception as e:
+        task_tx.abort()
+
+        if incremental:
+            return EnrichmentTaskResult(key_range=shard.key_range, ops=[], error=str(e))
+
+        logger.error(f"Enrichment task failed for shard {shard}: {e}")
+        raise e