neo4j-etl-lib 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

etl_lib/task/data_loading/ParallelSQLLoad2Neo4jTask.py

@@ -0,0 +1,122 @@
+from abc import ABC, abstractmethod
+from typing import Callable, Union
+
+from etl_lib.core.ClosedLoopBatchProcessor import ClosedLoopBatchProcessor
+from etl_lib.core.ETLContext import ETLContext
+from etl_lib.core.ParallelBatchProcessor import ParallelBatchProcessor
+from etl_lib.core.SplittingBatchProcessor import SplittingBatchProcessor, dict_id_extractor
+from etl_lib.core.Task import Task, TaskReturn
+from etl_lib.data_sink.CypherBatchSink import CypherBatchSink
+from etl_lib.data_source.SQLBatchSource import SQLBatchSource
+from sqlalchemy import text
+
+
+class ParallelSQLLoad2Neo4jTask(Task, ABC):
+    """
+    Parallelized version of SQLLoad2Neo4jTask: reads via SQLBatchSource,
+    splits into non-overlapping partitions (grid), processes each partition
+    in parallel through a CypherBatchSink, and closes the loop.
+
+    Subclasses must implement:
+        - _sql_query()
+        - _cypher_query()
+        - optionally override _count_query() and _id_extractor().
+
+    Control parameters:
+        batch_size: max items per partition batch
+        table_size: dimension of the splitting grid
+        max_workers: parallel threads per partition group (defaults to table_size)
+        prefetch: number of partition-groups to prefetch
+    """
+
+    def __init__(
+            self,
+            context: ETLContext,
+            batch_size: int = 5000,
+            table_size: int = 10,
+            max_workers: int = None,
+            prefetch: int = 4
+    ):
+        super().__init__(context)
+        self.context = context
+        self.batch_size = batch_size
+        self.table_size = table_size
+        # default max_workers to table_size for full parallelism
+        self.max_workers = max_workers or table_size
+        self.prefetch = prefetch
+
+    @abstractmethod
+    def _sql_query(self) -> str:
+        """
+        Return the SQL query to load source rows.
+        """
+        pass
+
+    @abstractmethod
+    def _cypher_query(self) -> str:
+        """
+        Return the Cypher query to write rows into Neo4j.
+        """
+        pass
+
+    def _count_query(self) -> Union[str, None]:
+        """
+        Optional SQL to count source rows for progress reporting.
+        """
+        return None
+
+    def _id_extractor(self) -> Callable:
+        """
+        Extractor mapping each row item to a (row,col) partition index.
+        Default expects dict rows with 'start' and 'end' keys.
+        Override to customize.
+        """
+        return dict_id_extractor()
+
+    def run_internal(self) -> TaskReturn:
+        # total count for ClosedLoopBatchProcessor
+        total_count = self.__get_source_count()
+        # source of raw rows
+        source = SQLBatchSource(self.context, self, self._sql_query())
+
+        # splitter: non-overlapping partitions as defined by the id_extractor
+        splitter = SplittingBatchProcessor(
+            context=self.context,
+            task=self,
+            predecessor=source,
+            table_size=self.table_size,
+            id_extractor=self._id_extractor()
+        )
+
+        # parallel processor: runs CypherBatchSink on each partition concurrently
+        parallel = ParallelBatchProcessor(
+            context=self.context,
+            task=self,
+            worker_factory=lambda: CypherBatchSink(context=self.context, task=self, predecessor=None,
+                                                   query=self._cypher_query()),
+            predecessor=splitter,
+            max_workers=self.max_workers,
+            prefetch=self.prefetch
+        )
+
+        # close loop: drives the pipeline and reports progress
+        closing = ClosedLoopBatchProcessor(
+            context=self.context,
+            task=self,
+            predecessor=parallel,
+            expected_rows=total_count
+        )
+
+        # run once to completion and return aggregated stats
+        result = next(closing.get_batch(self.batch_size))
+        return TaskReturn(True, result.statistics)
+
+    def __get_source_count(self):
+        count_query = self._count_query()
+        if count_query is None:
+            return None
+        with self.context.sql.engine.connect() as conn:
+            with conn.begin():
+                result = conn.execute(text(count_query))
+                row = result.fetchone()
+                return row[0] if row else None

etl_lib/task/data_loading/SQLLoad2Neo4jTask.py

@@ -1,8 +1,8 @@
 from abc import abstractmethod
 
+from etl_lib.core.ETLContext import ETLContext
 from sqlalchemy import text
 
-from etl_lib.core import ETLContext
 from etl_lib.core.ClosedLoopBatchProcessor import ClosedLoopBatchProcessor
 from etl_lib.core.Task import Task, TaskReturn
 from etl_lib.data_sink.CypherBatchSink import CypherBatchSink

{neo4j_etl_lib-0.2.0.dist-info → neo4j_etl_lib-0.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: neo4j-etl-lib
-Version: 0.2.0
+Version: 0.3.0
 Summary: Building blocks for ETL pipelines.
 Keywords: etl,graph,database
 Author-email: Bert Radke <bert.radke@pm.me>
@@ -15,7 +15,7 @@ Classifier: Topic :: Database
 Classifier: Development Status :: 4 - Beta
 License-File: LICENSE
 Requires-Dist: pydantic>=2.10.5; python_version >= '3.8'
-Requires-Dist: neo4j>=5.27.0; python_version >= '3.7'
+Requires-Dist: neo4j-rust-ext>=5.27.0; python_version >= '3.7'
 Requires-Dist: python-dotenv>=1.0.1; python_version >= '3.8'
 Requires-Dist: tabulate>=0.9.0; python_version >= '3.7'
 Requires-Dist: click>=8.1.8; python_version >= '3.7'

{neo4j_etl_lib-0.2.0.dist-info → neo4j_etl_lib-0.3.0.dist-info}/RECORD

@@ -1,32 +1,36 @@
-etl_lib/__init__.py,sha256=m8CiP9nDQUZCy4l5Crx09vtYrpAtqWkstIOj_Wlp1K8,65
+etl_lib/__init__.py,sha256=FyaTAuElsn3y3j1g15X141PnLTYxPrSUVU_YaVmiyPs,65
 etl_lib/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 etl_lib/cli/run_tools.py,sha256=KIar-y22P4kKm-yoJjecYsPwqC7U76M71dEgFO5-ZBo,8561
-etl_lib/core/BatchProcessor.py,sha256=6quNPE9Dp8hYJDQDTqxQtxbQ3KCmb56Mko34EIsNhyI,3352
+etl_lib/core/BatchProcessor.py,sha256=mRpdxZ6ZMKI8XsY3TPuy4dVcvRqLKCO-p63KeOhFyKE,3417
 etl_lib/core/ClosedLoopBatchProcessor.py,sha256=WzML1nldhZRbP8fhlD6utuK5SBYRl1cJgEobVDIdBP4,1626
-etl_lib/core/ETLContext.py,sha256=CHRNfjTmIIrPUp2F7l4MJlTsRLwfuBsNJdF_SEeShvI,6695
-etl_lib/core/ProgressReporter.py,sha256=FNmunA8q1vkAcmDF5z8ExH-mTotQMX0fyf2VAz-zI8o,8414
-etl_lib/core/Task.py,sha256=3e8iVXSfXaeBecvgTcs2LiIf2JwpKETRFhH4ig6lock,9202
-etl_lib/core/ValidationBatchProcessor.py,sha256=EhO6PFQB-4PZgIOTXP4PwkbAl5HRK0zgTeKMseRU5QU,3261
+etl_lib/core/ETLContext.py,sha256=wmEnbs3n_80B6La9Py_-MHG8BN0FajE9MjGPej0A3To,8045
+etl_lib/core/ParallelBatchProcessor.py,sha256=jNo1Xv1Ts34UZIseoQLDZOhHOVeEr8dUibKUt0FJ4Hw,7318
+etl_lib/core/ProgressReporter.py,sha256=UvWAPCuOrqyUcb5_kosIsCg1dyVQL-tnjgqnzs2cwZA,9372
+etl_lib/core/SplittingBatchProcessor.py,sha256=OIRMUVFpUoZc0w__JJjUr7B9QC3sBlqQp41xghrQzC0,11616
+etl_lib/core/Task.py,sha256=muQFY5qj2n-ZVV8F6vlHqo2lVSvB3wtGdIgkSXVpOFM,9365
+etl_lib/core/ValidationBatchProcessor.py,sha256=U1M2Qp9Ledt8qFiHAg8zMxE9lLRkBrr51NKs_Y8skK8,3400
 etl_lib/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-etl_lib/core/utils.py,sha256=wwfyvy78fL6sqHdV0IFqAVyEkp6vo5Yo8gRZua2dulw,816
+etl_lib/core/utils.py,sha256=9gezAateAmtWZ5lr6ZWYlTtgrj4-CUKJSGz7g1upnRg,2293
 etl_lib/data_sink/CSVBatchSink.py,sha256=oq4VJwnA4WSyJzdvwstGv73vOEuWmPSfCynhVmxBByU,2204
 etl_lib/data_sink/CypherBatchSink.py,sha256=nBH4bzN1IvdSFcKgiAIrAY5IauB565sdyVrnRc1hg_4,1566
 etl_lib/data_sink/SQLBatchSink.py,sha256=vyGrrxpdmCLUZMI2_W2ORej3FLGbwN9-b2GMYHd-k9g,1451
 etl_lib/data_sink/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 etl_lib/data_source/CSVBatchSource.py,sha256=HILkaQAFua1OM7xgSNKY6teXpcJjWUPaS4Aol-GLYL8,2767
 etl_lib/data_source/CypherBatchSource.py,sha256=06WuW11BqYjAXBZqL96Qr9MR8JrcjujDpxXe8cI-SYY,2238
-etl_lib/data_source/SQLBatchSource.py,sha256=rfIpzMU9IG5xrftZss7W-KpSIINCUdPfJBEXa5QDFts,2323
+etl_lib/data_source/SQLBatchSource.py,sha256=O3ZA2GXvo5j_KGwOILzguYZMPY_FJkV5j8FIa3-d9oM,4067
 etl_lib/data_source/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 etl_lib/task/CreateReportingConstraintsTask.py,sha256=nTcHLBIgXz_h2OQg-SHjQr68bhH974u0MwrtWPnVwng,762
 etl_lib/task/ExecuteCypherTask.py,sha256=thE8YTZzv1abxNhhDcb4p4ke6qmI6kWR4XQ-GrCBBBU,1284
 etl_lib/task/GDSTask.py,sha256=X1E83wYa-N7AXy43WPEqIy77d__z-2wpBjWNhGNXJzA,1781
 etl_lib/task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-etl_lib/task/data_loading/CSVLoad2Neo4jTask.py,sha256=bD1dWX3guY0f_lFqQP_pzPa03AYveLhjDQzOqkyJVQc,3840
-etl_lib/task/data_loading/SQLLoad2Neo4jTask.py,sha256=N8pQFx5vYKJ3UDn6sw7yBmCJao_b5vBOT3tJ8pWVVLk,2942
+etl_lib/task/data_loading/CSVLoad2Neo4jTask.py,sha256=9XiVdJHpABE-Hx1bsvTKLJWtChc8XMwXeO5RicaHDUo,3873
+etl_lib/task/data_loading/ParallelCSVLoad2Neo4jTask.py,sha256=2xN-5bHV9XgoaJLbbTEBuJFoZHV_CYi_hg6M1HQ-ffA,4030
+etl_lib/task/data_loading/ParallelSQLLoad2Neo4jTask.py,sha256=9xpCW5i8yGnUHyg475Z8rjN2P5NDshJGUEgTU6sm2Bs,4286
+etl_lib/task/data_loading/SQLLoad2Neo4jTask.py,sha256=HR3DcjOUkQN4SbCkgQYzljQCYhOcb-x2-DR5dBdapzU,2953
 etl_lib/task/data_loading/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 etl_lib/test_utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 etl_lib/test_utils/utils.py,sha256=kxWJqdRf1pg-4ByMfrtW3HDbgXIvyVtLndGDVvMCmoI,5641
-neo4j_etl_lib-0.2.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-neo4j_etl_lib-0.2.0.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-neo4j_etl_lib-0.2.0.dist-info/METADATA,sha256=FISyLgs53iHj3Uun7RgGFVtltLus_FiCkU591q9fJ3o,2497
-neo4j_etl_lib-0.2.0.dist-info/RECORD,,
+neo4j_etl_lib-0.3.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+neo4j_etl_lib-0.3.0.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+neo4j_etl_lib-0.3.0.dist-info/METADATA,sha256=GJcjdPvmzjEUq0pLndSSVzOg3c7CR6bIWz3sB_9tkVY,2506
+neo4j_etl_lib-0.3.0.dist-info/RECORD,,