pyspiral 0.7.1__cp312-abi3-manylinux_2_28_x86_64.whl → 0.7.2__cp312-abi3-manylinux_2_28_x86_64.whl

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

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

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

@@ -1,8 +1,8 @@
-pyspiral-0.7.1.dist-info/METADATA,sha256=WQ0LCz8jtz9qHpU6Zim1I_134EwD3Ve8V4bW0nU_t-0,1874
-pyspiral-0.7.1.dist-info/WHEEL,sha256=ydlpo1_yEJ2g1Axq3LoOd_OfioJa2swc2j5IDCa4uho,107
-pyspiral-0.7.1.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=ydlpo1_yEJ2g1Axq3LoOd_OfioJa2swc2j5IDCa4uho,107
+pyspiral-0.7.2.dist-info/entry_points.txt,sha256=R96Y3FpYX6XbQu9qMPfUTgiCcf4qM9OBQQZTDdBkZwA,74
 spiral/__init__.py,sha256=gAysTwG_oEeKVMdCOfOzDhl0bM2miiK8Ds2vvUihBWw,1153
-spiral/_lib.abi3.so,sha256=WqIyQu2RIN27T4K326ZBRMIVXysE6p-GW2ENnNueOm0,67478232
+spiral/_lib.abi3.so,sha256=Zgzetu00z_NPMToA5146UIUniIBSmokc1ER53GRxCcA,67496792
 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
@@ -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=4j5W68YEJOABw5CXrt0A5gMILdR8KaXPjeMLlAQ4Gi4,6072
+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
@@ -106,4 +106,4 @@ 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.1.dist-info/RECORD,,
+pyspiral-0.7.2.dist-info/RECORD,,

spiral/enrichment.py

@@ -4,11 +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.expressions import Expr
 
 if TYPE_CHECKING:
-    from spiral import KeySpaceIndex, Scan, Table
+    import dask.distributed
+
+    from spiral import KeySpaceIndex, Table
 
 logger = logging.getLogger(__name__)
 
@@ -47,20 +50,37 @@ 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(self, *, batch_readahead: int | None = None, partition_size_bytes: int | None = None) -> None:
+    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.
         """
-        with self._table.txn() as txn:
-            txn.writeback(
-                self._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.
@@ -70,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.
@@ -87,16 +108,19 @@ 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()
@@ -116,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))