added parallel download function

dataflux_core/download.py

@@ -21,6 +21,10 @@
 
 import uuid
 import logging
+import multiprocessing
+import itertools
+import math
+
 import signal
 import sys
 
@@ -156,6 +160,52 @@ def __init__(self, max_composite_object_size):
         self.max_composite_object_size = max_composite_object_size
 
 
+def dataflux_download_parallel(
+    project_name: str,
+    bucket_name: str,
+    objects: list[tuple[str, int]],
+    storage_client: object = None,
+    dataflux_download_optimization_params: DataFluxDownloadOptimizationParams = None,
+    parallelization: int = 1,
+) -> list[bytes]:
+    """Perform the DataFlux download algorithm in parallel to download the object contents as bytes and return.
+
+    Args:
+        project_name: the name of the GCP project.
+        bucket_name: the name of the GCS bucket that holds the objects to compose.
+            The function uploads the the composed object to this bucket too.
+        objects: A list of tuples which indicate the object names and sizes (in bytes) in the bucket.
+            Example: [("object_name_A", 1000), ("object_name_B", 2000)]
+        storage_client: the google.cloud.storage.Client initialized with the project.
+            If not defined, the function will initialize the client with the project_name.
+        dataflux_download_optimization_params: the paramemters used to optimize the download performance.
+        parallelization: The number of parallel processes that will simultaneously execute the download.
+    Returns:
+        the contents of the object in bytes.
+    """
+    chunk_size = math.ceil(len(objects) / parallelization)
+    chunks = []
+    for i in range(parallelization):
+        chunk = objects[i * chunk_size : (i + 1) * chunk_size]
+        if chunk:
+            chunks.append(chunk)
+    with multiprocessing.Pool(processes=len(chunks)) as pool:
+        results = pool.starmap(
+            dataflux_download,
+            (
+                (
+                    project_name,
+                    bucket_name,
+                    chunk,
+                    storage_client,
+                    dataflux_download_optimization_params,
+                )
+                for chunk in chunks
+            ),
+        )
+        return list(itertools.chain.from_iterable(results))
+
+
 def dataflux_download(
     project_name: str,
     bucket_name: str,

dataflux_core/tests/test_download.py

@@ -73,6 +73,37 @@ def test_dataflux_download(self):
                 f"expected only 3 objects in bucket, but found {len(bucket.blobs)}"
             )
 
+    def test_dataflux_download_parallel(self):
+        test_cases = [
+            {"name": "exceed number of items", "threads": 4},
+            {"name": "single thread", "threads": 1},
+            {"name": "standard", "threads": 2},
+        ]
+        bucket_name = "test_bucket"
+        objects = [("one", 3), ("two", 3), ("three", 5)]
+        client = fake_gcs.Client()
+        bucket = client.bucket(bucket_name)
+        bucket._add_file("one", bytes("one", "utf-8"))
+        bucket._add_file("two", bytes("two", "utf-8"))
+        bucket._add_file("three", bytes("three", "utf-8"))
+        params = download.DataFluxDownloadOptimizationParams(32)
+        expected_result = [b"one", b"two", b"three"]
+        for tc in test_cases:
+            result = download.dataflux_download_parallel(
+                "",
+                bucket_name,
+                objects,
+                client,
+                params,
+                tc["threads"],
+            )
+            self.assertEqual(result, expected_result)
+            # This checks for succesful deletion of the composed object.
+            if len(bucket.blobs) != 3:
+                self.fail(
+                    f"{tc['name']} expected only 3 objects in bucket, but found {len(bucket.blobs)}"
+                )
+
     def test_clean_composed_object(self):
         class ComposedObj:
             def __init__(self):