pyspiral 0.7.0__cp312-abi3-macosx_11_0_arm64.whl → 0.7.2__cp312-abi3-macosx_11_0_arm64.whl

{pyspiral-0.7.0.dist-info → pyspiral-0.7.2.dist-info}/METADATA

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

{pyspiral-0.7.0.dist-info → pyspiral-0.7.2.dist-info}/RECORD

@@ -1,8 +1,8 @@
-pyspiral-0.7.0.dist-info/METADATA,sha256=zP3jypGjswSz8nqMxMj8N6a4udG_LoZZLQAf8Xepg64,1874
-pyspiral-0.7.0.dist-info/WHEEL,sha256=KQvxBiy7GLcML6Ad3w_ZPrgSvER1uXd7aYb6wy6b44Y,103
-pyspiral-0.7.0.dist-info/entry_points.txt,sha256=R96Y3FpYX6XbQu9qMPfUTgiCcf4qM9OBQQZTDdBkZwA,74
+pyspiral-0.7.2.dist-info/METADATA,sha256=f9Q9zUpltwDCVkaY3jVW_spoFHDt0PwAFYaRePg-joU,1874
+pyspiral-0.7.2.dist-info/WHEEL,sha256=KQvxBiy7GLcML6Ad3w_ZPrgSvER1uXd7aYb6wy6b44Y,103
+pyspiral-0.7.2.dist-info/entry_points.txt,sha256=R96Y3FpYX6XbQu9qMPfUTgiCcf4qM9OBQQZTDdBkZwA,74
 spiral/__init__.py,sha256=gAysTwG_oEeKVMdCOfOzDhl0bM2miiK8Ds2vvUihBWw,1153
-spiral/_lib.abi3.so,sha256=Kdun77wTqTZOUrdRMBewZfE-ti3AGBGW9uWzd16W83U,70381040
+spiral/_lib.abi3.so,sha256=HD0QUi6SdaOO9F3MbS2jrcgXg8oE7pBrwPfhWOClOt4,70358864
 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,7 +35,7 @@ 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=zMp-xXGL4R1Py_rYrC5o3jFLam1oA74azi50dvMP-_o,6329
+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
@@ -60,13 +60,13 @@ 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=aXkKgV24w_0XmhTel_SOshJ2xCgkmXsFGZvaE7yXQ8k,6075
+spiral/enrichment.py,sha256=iKZn4tLsRQZPtaY-WdJSqPZ3H5UMfVyavcwTKO_3aCw,6980
 spiral/expressions/__init__.py,sha256=vMNFeeozkWph3dBpEkHPThUhZdT9ZZzxHe71HnkWlDU,8020
 spiral/expressions/base.py,sha256=PvhJkcUSsPSIaxirHVzM9zlqyBXiaiia1HXohXdOmL4,5377
-spiral/expressions/file.py,sha256=HRzGjc3goIlUlKjysoirexDaflNdnj9OoZ6j2uTKZnA,388
-spiral/expressions/http.py,sha256=WfHVLqz_LjBr78mN3ARBRQqgBrkao7-S73JxjC4Xwvo,356
+spiral/expressions/file.py,sha256=7D9jIENJcoT0KFharBLkzK9dZgO4DYn5K_KCt0twefg,518
+spiral/expressions/http.py,sha256=OOHh0WBxg3vwza_m74-rkoQWSclRMI60aPAbQ6yKZi0,486
 spiral/expressions/list_.py,sha256=MMt5lf5H1M3O-x6N_PvqOLGq9NOk6Ukv0fPWwPC_uy4,1809
-spiral/expressions/s3.py,sha256=d6Z_nnBRXopoPWnwGxlgjwyeYlrIMmonX-83PkE8LPo,375
+spiral/expressions/s3.py,sha256=PhQhMP-d8PLsSRtGCZbytnm7lI9VbDAbuSs2LBM4G7Q,505
 spiral/expressions/str_.py,sha256=tY8RXW3JWvr1-bEfCZtk5FAf11wKJnXPuA9EoeJ9tA4,1265
 spiral/expressions/struct.py,sha256=pGAnCDh6AK0BK1XfZ1qG4ce4ranIQEE1HQsgmzBcfwQ,2038
 spiral/expressions/text.py,sha256=-02gBWYoyNQ3qQ1--9HTa8IryUDojYQVIp8C7rgnOWQ,1893
@@ -102,8 +102,8 @@ spiral/streaming_/__init__.py,sha256=s7MlW2ERsuZmZGExLFL6RcZon2e0tNBocBg5ANgki7k
 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=prjDBcm6Qerdq3ypXzfbXb7ngAcO0j-Z9aTeZvzKoqs,12209
+spiral/table.py,sha256=ep8ZYtl6POebkPViR2FrekhFazNmAbOAESoLUODlup8,12242
 spiral/text_index.py,sha256=FQ9rgIEGLSJryS9lFdMhKtPFey18BXoWbPXyvZPJJ04,442
 spiral/transaction.py,sha256=KQhx3DvQyxG2C8md-YGsF_PgBRfayI0r_7ebMItDHdI,3938
 spiral/types_.py,sha256=W_jyO7F6rpPiH69jhgSgV7OxQZbOlb1Ho3InpKUP6Eo,155
-pyspiral-0.7.0.dist-info/RECORD,,
+pyspiral-0.7.2.dist-info/RECORD,,

spiral/client.py

@@ -116,6 +116,8 @@ class Spiral:
             asof = timestamp_micros(asof)
 
         # Combine all projections into a single struct.
+        if not projections:
+            raise ValueError("At least one projection is required.")
         projection = se.merge(*projections)
         if where is not None:
             where = se.lift(where)

spiral/enrichment.py

@@ -4,10 +4,13 @@ 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.expressions import Expr
 
 if TYPE_CHECKING:
+    import dask.distributed
+
     from spiral import KeySpaceIndex, Table
 
 logger = logging.getLogger(__name__)
@@ -47,19 +50,37 @@ class Enrichment:
         """The filter expression."""
         return self._where
 
-    def apply(self, *, batch_readahead: int | None = None, partition_size_bytes: int | None = None) -> None:
+    def _scan(self) -> Scan:
+        return self._table.spiral.scan(self._projection, where=self._where)
+
+    def apply(
+        self, *, batch_readahead: int | None = None, partition_size_bytes: int | None = None, tx_dump: str | None = None
+    ) -> None:
         """Apply the enrichment onto the table in a streaming fashion.
 
         For large tables, consider using `apply_dask` for distributed execution.
+
+        Args:
+            index: Optional key space index to use for sharding the enrichment.
+                If not provided, the table's default sharding will be used.
+            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.
         """
-        scan = self._table.spiral.scan(self._projection, where=self._where)
 
-        with self._table.txn() as txn:
-            txn.writeback(
-                scan,
-                partition_size_bytes=partition_size_bytes,
-                batch_readahead=batch_readahead,
-            )
+        txn = self._table.txn()
+
+        txn.writeback(
+            self._scan(),
+            partition_size_bytes=partition_size_bytes,
+            batch_readahead=batch_readahead,
+        )
+
+        if txn.is_empty():
+            logger.warning("Transaction not committed. No rows were read for enrichment.")
+            return
+
+        txn.commit(tx_dump=tx_dump)
 
     # TODO(marko): Need to figure out this sharding with key space index in places.
     #   We could compute on-demand instead of requiring a resource.
@@ -69,6 +90,7 @@ class Enrichment:
         index: Optional["KeySpaceIndex"] = None,
         partition_size_bytes: int | None = None,
         tx_dump: str | None = None,
+        client: Optional["dask.distributed.Client"] = None,
         **kwargs,
     ) -> None:
         """Use distributed Dask to apply the enrichment. Requires `dask[distributed]` to be installed.
@@ -86,20 +108,23 @@ class Enrichment:
                 If not provided, the table's default sharding will be used.
             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.
+            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.
         """
-        try:
-            from dask.distributed import Client
-        except ImportError:
-            raise ImportError("dask is not installed, please install dask[distributed] to use this feature.")
+        if client is None:
+            try:
+                from dask.distributed import Client
+            except ImportError:
+                raise ImportError("dask is not installed, please install dask[distributed] to use this feature.")
 
-        # Connect before doing any work.
-        dask_client = Client(**kwargs)
+            # Connect before doing any work.
+            client = Client(**kwargs)
 
         # Start a transaction BEFORE the planning scan.
         tx = self._table.txn()
-        plan_scan = self._table.spiral.scan(self._projection, where=self._where)
+        plan_scan = self._scan()
 
         # Determine the "tasks". Use the index if provided.
         shards = plan_scan.shards()
@@ -115,10 +140,10 @@ class Enrichment:
             output_table_id=self._table.table_id,
             partition_size_bytes=partition_size_bytes,
         )
-        enrichments = dask_client.map(_compute, shards)
+        enrichments = client.map(_compute, shards)
 
-        logger.info(f"Applying enrichment with {len(shards)} shards. Follow progress at {dask_client.dashboard_link}")
-        for result in dask_client.gather(enrichments):
+        logger.info(f"Applying enrichment with {len(shards)} shards. Follow progress at {client.dashboard_link}")
+        for result in client.gather(enrichments):
             result: EnrichmentTaskResult
             tx.include(result.ops)
 

spiral/expressions/file.py

@@ -2,15 +2,16 @@ from spiral import _lib
 from spiral.expressions.base import Expr, ExprLike
 
 
-def get(expr: ExprLike) -> Expr:
+def get(expr: ExprLike, abort_on_error: bool = False) -> Expr:
     """Read data from the local filesystem by the file:// URL.
 
     Args:
         expr: URLs of the data that needs to be read.
+        abort_on_error: Should the expression abort on errors or just collect them.
     """
     from spiral import expressions as se
 
     expr = se.lift(expr)
 
     # This just works :)
-    return Expr(_lib.expr.s3.get(expr.__expr__))
+    return Expr(_lib.expr.s3.get(expr.__expr__, abort_on_error))

spiral/expressions/http.py

@@ -2,15 +2,16 @@ from spiral import _lib
 from spiral.expressions.base import Expr, ExprLike
 
 
-def get(expr: ExprLike) -> Expr:
+def get(expr: ExprLike, abort_on_error: bool = False) -> Expr:
     """Read data from the URL.
 
     Args:
         expr: URLs of the data that needs to be read.
+        abort_on_error: Should the expression abort on errors or just collect them.
     """
     from spiral import expressions as se
 
     expr = se.lift(expr)
 
     # This just works :)
-    return Expr(_lib.expr.s3.get(expr.__expr__))
+    return Expr(_lib.expr.s3.get(expr.__expr__, abort_on_error))

spiral/expressions/s3.py

@@ -2,14 +2,15 @@ from spiral import _lib
 from spiral.expressions.base import Expr, ExprLike
 
 
-def get(expr: ExprLike) -> Expr:
+def get(expr: ExprLike, abort_on_error: bool = False) -> Expr:
     """Read data from object storage by the s3:// URL.
 
     Args:
         expr: URLs of the data that needs to be read from object storage.
+        abort_on_error: Should the expression abort on errors or just collect them.
     """
     from spiral import expressions as se
 
     expr = se.lift(expr)
 
-    return Expr(_lib.expr.s3.get(expr.__expr__))
+    return Expr(_lib.expr.s3.get(expr.__expr__, abort_on_error))

spiral/table.py

@@ -134,9 +134,9 @@ class Table(Expr):
         """
         from spiral import expressions as se
 
-        # Combine table with all projections into a single struct.
-        # The table is included to ensure key columns are present in the scan output.
-        projection = se.merge(self, *projections)
+        # 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)
         if where is not None:
             where = se.lift(where)