decentriq_util 0.1.0.dev2__tar.gz → 0.1.0.dev3__tar.gz

{decentriq_util-0.1.0.dev2 → decentriq_util-0.1.0.dev3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: decentriq_util
-Version: 0.1.0.dev2
+Version: 0.1.0.dev3
 Summary: Utility functions for Python code run in Decentriq workers
 Author: hafenr
 Author-email: robin.hafen@gmail.com
@@ -13,4 +13,5 @@ Classifier: Programming Language :: Python :: 3.12
 Requires-Dist: pandas (>=1.4.4)
 Requires-Dist: protobuf (>=4.21,<5.0)
 Requires-Dist: psutil (>=5.9.6,<6.0.0)
+Requires-Dist: pyarrow (>=14.0.0,<15.0.0)
 Requires-Dist: pyspark (>=3.4.0,<4.0.0)

{decentriq_util-0.1.0.dev2 → decentriq_util-0.1.0.dev3}/decentriq_util/spark/__init__.py

@@ -7,6 +7,7 @@ import os
 import shutil
 import psutil
 
+import pyarrow.parquet as pq
 from pyspark.sql import SparkSession
 
 
@@ -36,6 +37,33 @@ def _get_cgroup_memory() -> int:
         return 4 * 1024 * 1024 * 1024
 
 
+def _get_uncompressed_parquet_size(path: str):
+    def is_parquet(filename: str):
+        return filename.endswith(".parquet")
+    def get_uncompressed_size(file_path: str):
+        parquet_file = pq.ParquetFile(file_path)
+        metadata = parquet_file.metadata
+        uncompressed_size = 0
+        for i in range(metadata.num_row_groups):
+            row_group = metadata.row_group(i)
+            for j in range(row_group.num_columns):
+                column = row_group.column(j)
+                uncompressed_size += column.total_uncompressed_size
+        return uncompressed_size
+    if os.path.isfile(path) and is_parquet(path):
+        uncompressed_size = get_uncompressed_size(path)
+        return uncompressed_size
+    # Path might be a subdir as obtained by `partitionBy`
+    elif os.path.isdir(path):
+        total_uncompressed_size = 0
+        for file in os.listdir(path):
+            file_path = os.path.join(path, file)
+            total_uncompressed_size += _get_uncompressed_parquet_size(file_path)
+        return total_uncompressed_size
+    else:
+        return 0
+
+
 def _get_num_partitions(
         paths: list[str],
         available_memory: int,
@@ -44,36 +72,43 @@ def _get_num_partitions(
     """Determine the number of partitions s.t. each partition is as big as possible
     while still making sure all cores are occupied."""
     def get_file_or_dir_size(path):
-        if os.path.isfile(path):
-            size = os.path.getsize(path)
-        elif os.path.isdir(path):
-            total_size = 0
-            for dirpath, dirnames, filenames in os.walk(path):
-                for f in filenames:
-                    fp = os.path.join(dirpath, f)
-                    total_size += os.path.getsize(fp)
-            size = total_size
+        if path.endswith(".parquet"):
+            return _get_uncompressed_parquet_size(path)
         else:
-            raise ValueError(f"The path {path} is neither a file nor a directory.")
+            if os.path.isfile(path):
+                size = os.path.getsize(path)
+            elif os.path.isdir(path):
+                total_size = 0
+                for dirpath, dirnames, filenames in os.walk(path):
+                    for f in filenames:
+                        fp = os.path.join(dirpath, f)
+                        total_size += os.path.getsize(fp)
+                size = total_size
+            else:
+                raise ValueError(f"The path {path} is neither a file nor a directory.")
         return size
 
     file_size = max([get_file_or_dir_size(path) for path in paths])
+    if file_size == 0:
+        return 200  # default number of partitions in spark
 
     # Don't make the partitions smaller than 36MiB
     megabyte_in_bytes = 1024 * 1024
     min_partition_size = 32 * megabyte_in_bytes
 
     # Try to make partitions as big as possible s.t. each CPU is saturated,
-    # while accounting for 20% overhead in spark.
-    target_partition_size = max(
-        int(0.8 * available_memory / cpu_count) // megabyte_in_bytes * megabyte_in_bytes,
+    # while accounting for 50% overhead in spark.
+    task_per_cpu = 3  # try to get at least this many tasks/partitions per CPU
+    max_target_partition_size = max(
+        int(0.5 * available_memory / (task_per_cpu * cpu_count)) // megabyte_in_bytes * megabyte_in_bytes,
         min_partition_size
     )
 
+    target_partition_size = max_target_partition_size
     while (
         # If the file is small enough that it cannot be split across all cores,
         # divide the partition size further.
-        int(file_size / target_partition_size) < cpu_count
+        int(file_size / target_partition_size) < (task_per_cpu * cpu_count)
     ):
         # Don't go smaller than the minimum parition size
         if target_partition_size / 2 < min_partition_size:
@@ -81,6 +116,7 @@ def _get_num_partitions(
         target_partition_size = target_partition_size / 2
 
     print(f"File size: {file_size}")
+    print(f"Max target partition size: {max_target_partition_size}")
     print(f"Target partition size: {target_partition_size}")
 
     num_partitions = int(math.ceil(file_size / target_partition_size))
@@ -90,7 +126,7 @@ def _get_num_partitions(
 
 def _determine_optimal_spark_settings(
         files: list[str],
-        heap_size_extra_room: Optional[int] = 512 * 1024 * 1024
+        heap_size_extra_room: Optional[int] = 2 * 512 * 1024 * 1024
 ) -> list[Tuple[str, str]]:
     def get_spark_memory() -> int:
         """Determine the amount of memory available to spark"""
@@ -100,16 +136,14 @@ def _determine_optimal_spark_settings(
 
     cpu_count = multiprocessing.cpu_count()
     spark_memory = get_spark_memory()
+    spark_memory_4096 = (spark_memory // 4096) * 4096
     num_partitions = _get_num_partitions(files, spark_memory, cpu_count)
 
-    print(f"CPU count: {cpu_count}")
-    print(f"Spark memory: {spark_memory}")
-    print(f"Num partitions: {num_partitions}")
-
     settings = [
         ("spark.sql.shuffle.partitions", str(num_partitions)),
         ("spark.default.parallelism", str(num_partitions)),
         ("spark.driver.cores", str(cpu_count)),
+        ("spark.driver.memory", str(spark_memory_4096)),
     ]
     print("Spark settings:\n" + f"\n".join([str(x) for x in settings]))
 

{decentriq_util-0.1.0.dev2 → decentriq_util-0.1.0.dev3}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "decentriq_util"
-version = "0.1.0-dev2"
+version = "0.1.0-dev3"
 description = "Utility functions for Python code run in Decentriq workers"
 authors = ["hafenr <robin.hafen@gmail.com>"]
 
@@ -10,6 +10,7 @@ protobuf = "^4.21"
 pandas = ">=1.4.4"
 pyspark = "^3.4.0"
 psutil = "^5.9.6"
+pyarrow = "^14.0.0"
 
 [tool.poetry.dev-dependencies]
 pytest = "^7.2.0"