feat: Wire override_num_blocks through full call chain for Ray Data read operations

data_juicer/config/config.py

@@ -208,6 +208,14 @@ def build_base_parser() -> ArgumentParser:
         "(e.g., block_size for JSON reading). This configuration is "
         "especially useful when reading large JSON files.",
     )
+    parser.add_argument(
+        "--override_num_blocks",
+        type=Optional[int],
+        default=None,
+        help="Override the number of output blocks for Ray Data read "
+        "operations. Useful for controlling parallelism when reading "
+        "large datasets.",
+    )
     parser.add_argument(
         "--work_dir",
         type=str,

data_juicer/core/data/load_strategy.py

@@ -201,6 +201,8 @@ class RayLocalJsonDataLoadStrategy(RayDataLoadStrategy):
     def load_data(self, **kwargs):
         from data_juicer.core.data.ray_dataset import RayDataset
 
+        override_num_blocks = kwargs.pop("override_num_blocks", None)
+
         path = self.ds_config["path"]
 
         # Convert to absolute path if relative
@@ -281,7 +283,7 @@ def load_data(self, **kwargs):
         else:
             logger.info(f"Loading {data_format} data.")
         try:
-            dataset = RayDataset.read(data_format, path)
+            dataset = RayDataset.read(data_format, path, override_num_blocks=override_num_blocks)
             return RayDataset(dataset, dataset_path=path, cfg=self.cfg)
         except Exception as e:
             if auto_detect:

data_juicer/core/data/ray_dataset.py

@@ -347,11 +347,11 @@ def count(self) -> int:
         return self.data.count()
 
     @classmethod
-    def read(cls, data_format: str, paths: Union[str, List[str]]) -> RayDataset:
+    def read(cls, data_format: str, paths: Union[str, List[str]], **kwargs) -> RayDataset:
         if data_format in {"json", "jsonl", "json.gz", "jsonl.gz", "json.zst", "jsonl.zst"}:
-            return RayDataset.read_json(paths)
+            return RayDataset.read_json(paths, **kwargs)
         elif data_format == "webdataset":
-            return RayDataset.read_webdataset(paths)
+            return RayDataset.read_webdataset(paths, **kwargs)
         elif data_format in {
             "parquet",
             "images",
@@ -369,23 +369,23 @@ def read(cls, data_format: str, paths: Union[str, List[str]]) -> RayDataset:
                 from data_juicer.utils.lazy_loader import LazyLoader
 
                 LazyLoader.check_packages(["pylance"])
-            return getattr(ray.data, f"read_{data_format}")(paths)
+            return getattr(ray.data, f"read_{data_format}")(paths, **kwargs)
 
     @classmethod
-    def read_json(cls, paths: Union[str, List[str]]) -> RayDataset:
+    def read_json(cls, paths: Union[str, List[str]], override_num_blocks: Optional[int] = None) -> RayDataset:
         # Note: a temp solution for reading json stream
         # TODO: replace with ray.data.read_json_stream once it is available
         import pyarrow.json as js
 
         try:
             js.open_json
-            return read_json_stream(paths)
+            return read_json_stream(paths, override_num_blocks=override_num_blocks)
         except AttributeError:
-            return ray.data.read_json(paths)
+            return ray.data.read_json(paths, override_num_blocks=override_num_blocks)
 
     @classmethod
-    def read_webdataset(cls, paths: Union[str, List[str]]) -> RayDataset:
-        return ray.data.read_webdataset(paths, decoder=partial(_custom_default_decoder, format="PIL"))
+    def read_webdataset(cls, paths: Union[str, List[str]], **kwargs) -> RayDataset:
+        return ray.data.read_webdataset(paths, decoder=partial(_custom_default_decoder, format="PIL"), **kwargs)
 
     def to_list(self) -> list:
         return self.data.to_pandas().to_dict(orient="records")

data_juicer/core/executor/ray_executor_partitioned.py

@@ -414,7 +414,10 @@ def _run_impl(self, load_data_np: Optional[PositiveInt] = None, skip_return=Fals
         # Load the full dataset using a single DatasetBuilder
         logger.info("Loading dataset with single DatasetBuilder...")
 
-        dataset = self.datasetbuilder.load_dataset(num_proc=load_data_np)
+        override_num_blocks = getattr(self.cfg, "override_num_blocks", None)
+        dataset = self.datasetbuilder.load_dataset(
+            num_proc=load_data_np, override_num_blocks=override_num_blocks
+        )
         columns = dataset.schema().columns
 
         # Prepare operations