pyspiral 0.6.18__cp312-abi3-macosx_11_0_arm64.whl → 0.6.20__cp312-abi3-macosx_11_0_arm64.whl

{pyspiral-0.6.18.dist-info → pyspiral-0.6.20.dist-info}/METADATA

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

{pyspiral-0.6.18.dist-info → pyspiral-0.6.20.dist-info}/RECORD

@@ -1,8 +1,8 @@
-pyspiral-0.6.18.dist-info/METADATA,sha256=25UcwQPTisubAD48g4I_FVeNQ7PhcM5NFyE1loBsGxo,1875
-pyspiral-0.6.18.dist-info/WHEEL,sha256=KQvxBiy7GLcML6Ad3w_ZPrgSvER1uXd7aYb6wy6b44Y,103
-pyspiral-0.6.18.dist-info/entry_points.txt,sha256=uft7u-a6g40NLt4Q6BleWbK4NY0M8nZuYPpP8DV0EOk,45
+pyspiral-0.6.20.dist-info/METADATA,sha256=1lIl3Ge03ZbmZSes7OFOIAreoQnR8jrYm7je4lmXVS4,1875
+pyspiral-0.6.20.dist-info/WHEEL,sha256=KQvxBiy7GLcML6Ad3w_ZPrgSvER1uXd7aYb6wy6b44Y,103
+pyspiral-0.6.20.dist-info/entry_points.txt,sha256=R96Y3FpYX6XbQu9qMPfUTgiCcf4qM9OBQQZTDdBkZwA,74
 spiral/__init__.py,sha256=gAysTwG_oEeKVMdCOfOzDhl0bM2miiK8Ds2vvUihBWw,1153
-spiral/_lib.abi3.so,sha256=pLhJoDJIT_f0Ibvp2FwaTnI6VWXBdb-GPABakjheHHM,70263280
+spiral/_lib.abi3.so,sha256=OjvPM5bRKKcb2gqVePTP3-DLnCQvLkv_ZrZBGdgQ-sQ,70404048
 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
@@ -49,7 +49,7 @@ 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=YBL12_JPTWz2mNbqlDqbT1exxVJYzwfXdHCi6Z37JxA,3841
+spiral/core/table/__init__.pyi,sha256=szCtZqZ_L0vF_99x7d3olwazJdEn5LwkgRK-8QEOzrI,3914
 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
@@ -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=YDaXcJPtmJzpLrYmn2pdllVcRIkXlb578KKgkIb38Eo,6518
+spiral/enrichment.py,sha256=aXkKgV24w_0XmhTel_SOshJ2xCgkmXsFGZvaE7yXQ8k,6075
 spiral/expressions/__init__.py,sha256=vMNFeeozkWph3dBpEkHPThUhZdT9ZZzxHe71HnkWlDU,8020
 spiral/expressions/base.py,sha256=PvhJkcUSsPSIaxirHVzM9zlqyBXiaiia1HXohXdOmL4,5377
 spiral/expressions/file.py,sha256=HRzGjc3goIlUlKjysoirexDaflNdnj9OoZ6j2uTKZnA,388
@@ -94,7 +94,7 @@ spiral/protogen/_/substrait/extensions/__init__.py,sha256=nhnEnho70GAT8WPj2xtwJU
 spiral/protogen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 spiral/protogen/util.py,sha256=smnvVo6nYH3FfDm9jqhNLaXz4bbTBaQezHQDCTvZyiQ,1486
 spiral/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-spiral/scan.py,sha256=csbk5ePbU-RlEVIF7isccF2zRBB8L8ZY_HEpalMjgLY,12340
+spiral/scan.py,sha256=0ZgzPX5gZ3uvbCyifpjHKXYk_ERphf9xHbrOVKRNkXc,12574
 spiral/server.py,sha256=ztBmB5lBnUz-smQxR_tC8AI5SOhz17wH0MI3GuzDUdM,600
 spiral/settings.py,sha256=sUhMMBCXaPvUYztN_gztD9TjeUYJwVeEcJrq4FLy6M0,3232
 spiral/snapshot.py,sha256=cTobi5jtiANxalGA-isokQHblNmXGtuUvgUGGNVybsI,1555
@@ -104,6 +104,6 @@ spiral/streaming_/stream.py,sha256=DM1hBDHnWm1ZFKZ-hZ4zxeSXITcUI6kWzwdJZvywI8o,5
 spiral/substrait_.py,sha256=AKeOD4KIXvz2J4TYxnIneOiHddtBIyOhuNxVO_uH0eg,12592
 spiral/table.py,sha256=prjDBcm6Qerdq3ypXzfbXb7ngAcO0j-Z9aTeZvzKoqs,12209
 spiral/text_index.py,sha256=FQ9rgIEGLSJryS9lFdMhKtPFey18BXoWbPXyvZPJJ04,442
-spiral/transaction.py,sha256=hQm6DfCklMDpIYJ9qA2wR45cCuUPGCiJy1tHGE3AsEY,3418
+spiral/transaction.py,sha256=KQhx3DvQyxG2C8md-YGsF_PgBRfayI0r_7ebMItDHdI,3938
 spiral/types_.py,sha256=W_jyO7F6rpPiH69jhgSgV7OxQZbOlb1Ho3InpKUP6Eo,155
-pyspiral-0.6.18.dist-info/RECORD,,
+pyspiral-0.6.20.dist-info/RECORD,,

{pyspiral-0.6.18.dist-info → pyspiral-0.6.20.dist-info}/entry_points.txt

@@ -1,2 +1,3 @@
 [console_scripts]
 spiral=spiral.cli.app:main
+pyspiral=spiral.cli.app:main

spiral/core/table/__init__.pyi

@@ -73,6 +73,7 @@ class Scan:
         key_range: KeyRange | None = None,
         key_table: pa.Table | pa.RecordBatch | None = None,
         batch_readahead: int | None = None,
+        progress: bool = True,
     ) -> pa.RecordBatchReader: ...
     def to_shuffled_record_batches(
         self,
@@ -111,6 +112,7 @@ class Transaction:
         batch_readahead: int | None = None,
     ): ...
     def drop_columns(self, column_paths: list[str]): ...
+    def ops(self) -> list[Operation]: ...
     def take(self) -> list[Operation]: ...
     def include(self, ops: list[Operation]): ...
     def commit(self): ...

spiral/enrichment.py

@@ -86,7 +86,6 @@ 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 operations as a JSON file for debugging.
             **kwargs: Additional keyword arguments to pass to `dask.distributed.Client`
                 such as `address` to connect to an existing cluster.
         """
@@ -100,7 +99,6 @@ class Enrichment:
 
         # Start a transaction BEFORE the planning scan.
         tx = self._table.txn()
-        backup_ops = []
         plan_scan = self._table.spiral.scan(self._projection, where=self._where)
 
         # Determine the "tasks". Use the index if provided.
@@ -123,19 +121,12 @@ class Enrichment:
         for result in dask_client.gather(enrichments):
             result: EnrichmentTaskResult
             tx.include(result.ops)
-            backup_ops.extend(result.ops)
 
         if tx.is_empty():
             logger.warning("Transaction not committed. No rows were read for enrichment.")
             return
 
-        # TODO(marko): We can remove this when I have more trust in very large tx commits.
-        if tx_dump is not None:
-            with open(tx_dump, "w") as f:
-                f.writelines([op.to_json() for op in backup_ops])
-            logger.info(f"Transaction dumped to {tx_dump}")
-
-        tx.commit()
+        tx.commit(tx_dump=tx_dump)
 
 
 @dataclasses.dataclass

spiral/scan.py

@@ -58,6 +58,7 @@ class Scan:
         key_table: pa.Table | pa.RecordBatchReader | None = None,
         batch_size: int | None = None,
         batch_readahead: int | None = None,
+        hide_progress_bar: bool = False,
     ) -> pa.RecordBatchReader:
         """Read as a stream of RecordBatches.
 
@@ -71,6 +72,7 @@ class Scan:
                 IMPORTANT: This is currently only respected when the key_table is used. If key table is a
                     RecordBatchReader, the batch_size argument must be None, and the existing batching is respected.
             batch_readahead: the number of batches to prefetch in the background.
+            hide_progress_bar: If True, disables the progress bar during reading.
         """
         if key_range is not None and key_table is not None:
             raise ValueError("Only one of key_range or key_table can be provided.")
@@ -83,7 +85,9 @@ class Scan:
         elif isinstance(key_table, pa.Table):
             key_table = key_table.to_reader(max_chunksize=batch_size)
 
-        return self.core.to_record_batches(key_range=key_range, key_table=key_table, batch_readahead=batch_readahead)
+        return self.core.to_record_batches(
+            key_range=key_range, key_table=key_table, batch_readahead=batch_readahead, progress=(not hide_progress_bar)
+        )
 
     def to_table(
         self,
@@ -325,4 +329,4 @@ def _read_shard_task(shard: Shard, *, settings_dict, state_json) -> "pd.DataFram
     state = ScanState.from_json(state_json)
     task_scan = Scan(sp, sp.core.load_scan(state))
 
-    return task_scan.to_pandas(key_range=shard.key_range)
+    return task_scan.to_record_batches(key_range=shard.key_range, hide_progress_bar=True).read_all().to_pandas()

spiral/transaction.py

@@ -1,9 +1,13 @@
+import logging
+
 from spiral.core.table import KeyRange
 from spiral.core.table import Transaction as CoreTransaction
 from spiral.core.table.spec import Operation
 from spiral.expressions.base import ExprLike
 from spiral.scan import Scan
 
+logger = logging.getLogger(__name__)
+
 
 class Transaction:
     """Spiral table transaction.
@@ -91,8 +95,17 @@ class Transaction:
         """
         self._core.include(ops)
 
-    def commit(self):
+    def commit(self, *, tx_dump: str | None = None):
         """Commit the transaction."""
+        # TODO(marko): We can remove this when I have more trust in large tx commits.
+        if tx_dump is not None:
+            try:
+                with open(tx_dump, "w") as f:
+                    f.writelines([op.to_json() for op in self._core.ops()])
+                logger.info(f"Transaction dumped to {tx_dump}")
+            except Exception as e:
+                logger.error(f"Failed to dump transaction to {tx_dump}: {e}")
+
         self._core.commit()
 
     def abort(self):