pyspiral 0.7.9__cp312-abi3-manylinux_2_28_x86_64.whl → 0.7.10__cp312-abi3-manylinux_2_28_x86_64.whl

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

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

{pyspiral-0.7.9.dist-info → pyspiral-0.7.10.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=ydlpo1_yEJ2g1Axq3LoOd_OfioJa2swc2j5IDCa4uho,107
-pyspiral-0.7.9.dist-info/entry_points.txt,sha256=R96Y3FpYX6XbQu9qMPfUTgiCcf4qM9OBQQZTDdBkZwA,74
+pyspiral-0.7.10.dist-info/METADATA,sha256=rCgXce3dHmwg5oIvXTHJnnMMEX93NFLvTRbUY0Ns9-Y,1875
+pyspiral-0.7.10.dist-info/WHEEL,sha256=ydlpo1_yEJ2g1Axq3LoOd_OfioJa2swc2j5IDCa4uho,107
+pyspiral-0.7.10.dist-info/entry_points.txt,sha256=R96Y3FpYX6XbQu9qMPfUTgiCcf4qM9OBQQZTDdBkZwA,74
 spiral/__init__.py,sha256=PwaYBWFBtB7cYi7peMmhk_Lm5XzjRoLwOtLbUhc1ZDo,1449
-spiral/_lib.abi3.so,sha256=jpQ5vqK5-SgdNp-oHrMgE2oJSuPxJamtjVNAmAfn8oI,68064024
+spiral/_lib.abi3.so,sha256=fbdXuFFBMXuFd6-B9k-9LmTtt4eSV5r-sYfRqSbK8bA,68064904
 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,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=t3CFnidG1kGHeJk1zIMVyImKapUJQx1OXvGn88brOo4,11059
 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.9.dist-info/RECORD,,
+pyspiral-0.7.10.dist-info/RECORD,,

spiral/enrichment.py

@@ -160,15 +160,48 @@ 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 not isinstance(e, KilledWorker):
+                # Re-raise other exceptions
+                raise e
+
+            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:
@@ -212,7 +245,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 +252,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"]
 
@@ -255,12 +284,12 @@ def _enrichment_task(
 
     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