f9columnar 0.2.0__tar.gz → 0.2.2__tar.gz

{f9columnar-0.2.0 → f9columnar-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: f9columnar
-Version: 0.2.0
+Version: 0.2.2
 Summary: Columnar analysis utils.
 Home-page: https://gitlab.cern.ch/ijs-f9-ljubljana/f9columnar
 License: MIT
@@ -24,6 +24,7 @@ Requires-Dist: numba (==0.60.0)
 Requires-Dist: numpy (==2.0.2)
 Requires-Dist: pandas (==2.2.2)
 Requires-Dist: plothist (>=1.2.6,<2.0.0)
+Requires-Dist: psutil (>=6.1.0,<7.0.0)
 Requires-Dist: pydot (==3.0.1)
 Requires-Dist: pygraphviz (==1.13)
 Requires-Dist: requests (>=2.28.1,<3.0.0)

{f9columnar-0.2.0 → f9columnar-0.2.2}/f9columnar/hdf5_dataloader.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 import copy
 import json
 import logging
-import multiprocessing
 import os
 from collections.abc import Callable
 from itertools import product
@@ -11,8 +10,10 @@ from typing import Any
 
 import h5py
 import numpy as np
+import pandas as pd
 import torch
 from numpy.typing import ArrayLike
+from torch import multiprocessing
 from torch.utils.data import DataLoader, IterableDataset
 
 from f9columnar.processors import Processor, ProcessorsGraph
@@ -95,7 +96,7 @@ class Hdf5Iterator:
         return arrays, reports
 
 
-class Hdf5IteratorMaker:
+class Hdf5IteratorDfMaker:
     def __init__(
         self,
         name: str,
@@ -115,6 +116,7 @@ class Hdf5IteratorMaker:
         self.holdout_range = holdout_range
 
         self.total_num_entries = shape[0]
+        self.all_num_entries_dct: dict[str, int] = {}
 
     def _log_info(self) -> None:
         all_hdf5_files = list(self.hdf5_files_metadata.keys())
@@ -143,15 +145,15 @@ class Hdf5IteratorMaker:
             holdout_range=self.holdout_range,
         )
 
-    def make(self) -> list[list[Hdf5Iterator]]:
+    def _split(self) -> list[dict[str, list[int]]]:
         self._log_info()
 
         # how many entries each worker will process
         splits = [self.total_num_entries // self.num_workers] * self.num_workers
         splits[-1] += self.total_num_entries % self.num_workers
 
-        all_num_entries_dct = {file: metadata["shape"][0] for file, metadata in self.hdf5_files_metadata.items()}
-        num_entries_dct = copy.deepcopy(all_num_entries_dct)
+        self.all_num_entries_dct = {file: metadata["shape"][0] for file, metadata in self.hdf5_files_metadata.items()}
+        num_entries_dct = copy.deepcopy(self.all_num_entries_dct)
 
         # keep track of the start and stop entries for each root file
         hdf5_files_start_dct: dict[str, int] = {file: 0 for file in self.hdf5_files_metadata.keys()}
@@ -171,7 +173,7 @@ class Hdf5IteratorMaker:
                 total += num_entries
 
                 if total <= split:
-                    result[i][hdf5_file] = [start_entry, all_num_entries_dct[hdf5_file]]
+                    result[i][hdf5_file] = [start_entry, self.all_num_entries_dct[hdf5_file]]
                     done.append(hdf5_file)
 
                     if total == split:
@@ -186,39 +188,74 @@ class Hdf5IteratorMaker:
                     num_entries_dct[hdf5_file] -= delta
                     break
 
-        worker_hdf5_iterators: list[list] = [[] for _ in range(self.num_workers)]
-        check_total = 0
+        return result
 
-        for i, result_dct in enumerate(result):
-            for hdf5_file, start_end in result_dct.items():
-                entry_start, entry_end = start_end
-                check_total += entry_end - entry_start
+    def make(self) -> pd.DataFrame:
+        split_result = self._split()
 
-                iterator = self._make_iterator(hdf5_file, entry_start, entry_end)
-                worker_hdf5_iterators[i].append(iterator)
+        worker_df: dict[str, list] = {
+            "worker_id": [],
+            "file": [],
+            "start": [],
+            "stop": [],
+            "chunk_size": [],
+            "holdout_range": [],
+            "shuffle": [],
+        }
+
+        check_total = 0
+        for i, result_dct in enumerate(split_result):
+            for hdf5_file, start_stop in result_dct.items():
+                entry_start, entry_stop = start_stop
+                check_total += entry_stop - entry_start
+
+                worker_df["worker_id"].append(i)
+                worker_df["file"].append(hdf5_file)
+                worker_df["start"].append(entry_start)
+                worker_df["stop"].append(entry_stop)
+                worker_df["chunk_size"].append(self.chunk_size)
+                worker_df["holdout_range"].append(self.holdout_range)
+                worker_df["shuffle"].append(self.shuffle)
 
         if check_total != self.total_num_entries:
             raise ValueError("Total number of entries does not match.")
 
-        return worker_hdf5_iterators
+        return pd.DataFrame(worker_df)
 
 
 class Hdf5LoaderIterator:
     def __init__(
         self,
         name: str,
-        hdf5_iterators: list[Hdf5Iterator],
+        iterators_df: pd.DataFrame,
         worker_id: int,
         processors: list[Callable[[ArrayLike, dict], tuple[ArrayLike, dict]]] | ProcessorsGraph | None,
         hdf5_files_desc_dct: dict[str, dict[str, Any]] | None = None,
     ) -> None:
         self.name = name
-        self.hdf5_iterators = hdf5_iterators
+        self.iterators_df = iterators_df
         self.worker_id = worker_id
         self.processors = processors
         self.hdf5_files_desc_dct = hdf5_files_desc_dct
 
-        self.current_iterator_idx = 0
+        self.current_df_idx, self.current_iterator_idx = 0, 0
+
+    def _make_hdf5_iterator(self, df: pd.Series) -> Hdf5Iterator:
+        iterator = Hdf5Iterator(
+            df["file"],
+            chunk_size=df["chunk_size"],
+            start_entry=df["start"],
+            stop_entry=df["stop"],
+            shuffle=df["shuffle"],
+            holdout_range=df["holdout_range"],
+        )
+
+        return iterator
+
+    def _iterate_df(self) -> None:
+        df = self.iterators_df.iloc[self.current_iterator_idx]
+        self.iterator = self._make_hdf5_iterator(df)
+        self.current_df_idx += 1
 
     def _run_processors(self, arrays: ArrayLike, reports: dict) -> tuple[ArrayLike, dict] | dict[str, Processor]:
         if self.processors is None:
@@ -247,18 +284,22 @@ class Hdf5LoaderIterator:
 
     def __next__(self) -> tuple[ArrayLike, dict] | dict[str, Processor]:
         try:
-            current_iterator = self.hdf5_iterators[self.current_iterator_idx]
-            arrays, reports = next(current_iterator)
+            if self.current_df_idx == self.current_iterator_idx:
+                self._iterate_df()
+
+            arrays, reports = next(self.iterator)
 
         except StopIteration:
-            current_iterator.close()
+            self.iterator.close()
             self.current_iterator_idx += 1
 
-            if self.current_iterator_idx == len(self.hdf5_iterators):
+            if self.current_iterator_idx == len(self.iterators_df):
                 raise StopIteration
-            else:
-                current_iterator = self.hdf5_iterators[self.current_iterator_idx]
-                arrays, reports = next(current_iterator)
+
+            if self.current_df_idx == self.current_iterator_idx:
+                self._iterate_df()
+
+            arrays, reports = next(self.iterator)
 
         reports = self._make_report(reports)
 
@@ -308,7 +349,7 @@ class Hdf5IterableDataset(IterableDataset):
 
         self.shape = self._get_total_shape()
 
-        self.worker_hdf5_iterators = self._get_iterators()
+        self.worker_iterators_df = self._get_df_iterators()
 
     def _setup(self) -> None:
         for file, dataset_name in product(self.files, self.dataset_names):
@@ -384,9 +425,15 @@ class Hdf5IterableDataset(IterableDataset):
 
         return sum(shapes_0), shapes_1[0]
 
-    def _get_iterators(self) -> list[list[Hdf5Iterator]]:
-        return Hdf5IteratorMaker(
-            self.name, self.metadata, self.num_workers, self.shape, self.chunk_size, self.shuffle, self.holdout_range
+    def _get_df_iterators(self) -> pd.DataFrame:
+        return Hdf5IteratorDfMaker(
+            self.name,
+            self.metadata,
+            self.num_workers,
+            self.shape,
+            self.chunk_size,
+            self.shuffle,
+            self.holdout_range,
         ).make()
 
     def __iter__(self) -> Hdf5LoaderIterator:
@@ -397,17 +444,21 @@ class Hdf5IterableDataset(IterableDataset):
         else:
             worker_id = worker_info.id
 
-        hdf5_iterators = self.worker_hdf5_iterators[worker_id]
+        iterators_df = self.worker_iterators_df[self.worker_iterators_df["worker_id"] == worker_id].copy()
 
         return Hdf5LoaderIterator(
             self.name,
-            hdf5_iterators,
+            iterators_df,
             worker_id,
             self.processors,
             self.hdf5_files_desc_dct,
         )
 
 
+def default_collate_fn(batch: list[Any]) -> list[Any]:
+    return batch
+
+
 def get_hdf5_dataloader(
     name: str,
     files: str | list[str],
@@ -421,6 +472,10 @@ def get_hdf5_dataloader(
     processors: list[Callable[[ArrayLike, dict], tuple[ArrayLike, dict]]] | ProcessorsGraph | None = None,
     dataloader_kwargs: dict[str, Any] | None = None,
 ) -> tuple[DataLoader, int]:
+
+    if multiprocessing.get_start_method() == "fork" and num_workers > 0:
+        logging.debug("Using 'fork' start method. Consider using 'spawn' or 'forkserver'.")
+
     if num_workers == -1:
         num_workers = multiprocessing.cpu_count()
 
@@ -441,7 +496,7 @@ def get_hdf5_dataloader(
         name,
         files,
         dataset_names,
-        num_workers,
+        num_workers=num_workers if num_workers > 0 else 1,
         chunk_size=chunk_size,
         use_piles=use_piles,
         shuffle=shuffle,
@@ -454,7 +509,7 @@ def get_hdf5_dataloader(
         hdf5_dataset,
         batch_size=None,
         num_workers=num_workers,
-        collate_fn=lambda batch: batch,
+        collate_fn=default_collate_fn,
         **dataloader_kwargs,
     )
 

{f9columnar-0.2.0 → f9columnar-0.2.2}/f9columnar/root_dataloader.py

@@ -2,19 +2,21 @@ from __future__ import annotations
 
 import copy
 import logging
-import multiprocessing
 import os
 from collections.abc import Callable, Iterator
 from dataclasses import dataclass, field
 from typing import Any
 
 import awkward as ak
+import pandas as pd
 import torch
 import uproot
+from torch import multiprocessing
 from torch.utils.data import DataLoader, IterableDataset
 
 from f9columnar.processors import Processor, ProcessorsGraph
 from f9columnar.utils.helpers import get_file_size
+from f9columnar.utils.loggers import get_progress
 
 
 @dataclass
@@ -45,6 +47,10 @@ class RootFile:
         self.num_entries = self.tree.num_entries
         return self
 
+    def close(self) -> RootFile:
+        self.tree.close()
+        return self
+
 
 @dataclass
 class RootFiles:
@@ -60,6 +66,10 @@ class RootFiles:
     total_num_entries: int = 0
 
     def __post_init__(self) -> None:
+        progress = get_progress()
+        progress.start()
+        bar = progress.add_task("Preparing ROOT files", total=len(self.file_names))
+
         for i, file_name in enumerate(self.file_names):
             if type(self.key) is list:
                 root_file = RootFile(file_name, self.key[i])
@@ -68,89 +78,32 @@ class RootFiles:
             else:
                 raise ValueError(f"Key {self.key} is not a valid type.")
 
+            progress.update(bar, advance=1)
+
             root_file = root_file.open()
 
             self.files_dct[file_name] = root_file
-
             self.file_size_dct[file_name] = root_file.file_size
-            self.num_entries_dct[file_name] = root_file.num_entries
-
-            self.total_file_size += root_file.file_size
-            self.total_num_entries += root_file.num_entries
-
-    def __getitem__(self, file_name: str) -> RootFile:
-        return self.files_dct[file_name]
 
+            num_entries = root_file.num_entries
 
-class UprootIterator:
-    def __init__(
-        self,
-        root_file: RootFile,
-        step_size: int,
-        filter_name: Callable[[str], bool] | None = None,
-        entry_start: int | None = None,
-        entry_stop: int | None = None,
-        **kwargs: Any,
-    ) -> None:
-        if type(step_size) is not int:
-            raise ValueError(f"Step size {step_size} is not an integer.")
+            root_file.close()
 
-        self.root_file = root_file
-        self.step_size = step_size
-        self.filter_name = filter_name
-        self.kwargs = kwargs
+            if num_entries == 0:
+                continue
 
-        num_entries = self.root_file.num_entries
+            self.num_entries_dct[file_name] = num_entries
 
-        # assign start and stop entries
-        if entry_start is None:
-            self.entry_start = 0
-        else:
-            self.entry_start = entry_start
-
-        if entry_stop is None:
-            self.entry_stop = num_entries
-        else:
-            self.entry_stop = entry_stop
-
-        # handle cases where step size is larger than the number of entries
-        if self.entry_start is not None and self.entry_stop is not None:
-            delta_entry = self.entry_stop - self.entry_start
-
-            if self.step_size > delta_entry:
-                self.step_size = delta_entry
-        else:
-            if self.step_size > num_entries:
-                self.step_size = num_entries
-
-        self.tree: uproot.TTree = None
-        self.returned_iterator = False
-
-    def get_iterator(self) -> Iterator[tuple[ak.Array, Any]]:
-        if self.root_file.tree is None:
-            self.tree = uproot.open(self.root_file.file_name)
-        else:
-            self.tree = self.root_file.tree
-
-        uproot_iterator = self.tree.iterate(
-            library="ak",
-            report=True,
-            step_size=self.step_size,
-            filter_name=self.filter_name,
-            entry_start=self.entry_start,
-            entry_stop=self.entry_stop,
-            **self.kwargs,
-        )
-
-        self.returned_iterator = True
+            self.total_file_size += root_file.file_size
+            self.total_num_entries += num_entries
 
-        return uproot_iterator
+        progress.stop()
 
-    def close(self) -> None:
-        self.tree.close()
+    def __getitem__(self, file_name: str) -> RootFile:
+        return self.files_dct[file_name]
 
 
-class UprootIteratorMaker:
+class UprootIteratorsDfMaker:
     def __init__(
         self,
         name: str,
@@ -158,17 +111,14 @@ class UprootIteratorMaker:
         key: str | list[str],
         step_size: int,
         num_workers: int,
-        filter_name: Callable[[str], bool] | None = None,
-        **kwargs: Any,
     ) -> None:
         self.name = name
         self.files, self.key = files, key
         self.num_workers = num_workers
         self.step_size = step_size
-        self.filter_name = filter_name
-        self.kwargs = kwargs
 
         self.total_num_entries = 0
+        self.all_num_entries_dct: dict[str, int] = {}
 
     def _log_info(self, total_files_size: float, total_num_entries: int) -> None:
         info_str = "\n" + 15 * "=" + " info " + 15 * "="
@@ -180,17 +130,7 @@ class UprootIteratorMaker:
 
         logging.info(info_str)
 
-    def _make_iterator(self, root_file: RootFile, entry_start: int, entry_stop: int) -> UprootIterator:
-        return UprootIterator(
-            root_file,
-            step_size=self.step_size,
-            filter_name=self.filter_name,
-            entry_start=entry_start,
-            entry_stop=entry_stop,
-            **self.kwargs,
-        )
-
-    def make(self) -> list[list[UprootIterator]]:
+    def _split(self) -> list[dict[str, list[int]]]:
         root_files = RootFiles(self.files, self.key)
 
         self._log_info(root_files.total_file_size, root_files.total_num_entries)
@@ -202,13 +142,13 @@ class UprootIteratorMaker:
         splits = [total_num_entries // self.num_workers] * self.num_workers
         splits[-1] += total_num_entries % self.num_workers
 
-        all_num_entries_dct = root_files.num_entries_dct
+        self.all_num_entries_dct = root_files.num_entries_dct
         num_entries_dct = copy.deepcopy(root_files.num_entries_dct)
 
         # keep track of the start and stop entries for each root file
         root_files_start_dct: dict[str, int] = {root_file: 0 for root_file in self.files}
 
-        result: list[dict[str, list[int]]] = [{} for _ in range(len(splits))]
+        split_result: list[dict[str, list[int]]] = [{} for _ in range(len(splits))]
 
         done = []
         for i, split in enumerate(splits):
@@ -218,12 +158,12 @@ class UprootIteratorMaker:
                 if root_file in done:
                     continue
 
-                start_entry = root_files_start_dct[root_file]
+                entry_start = root_files_start_dct[root_file]
 
                 total += num_entries
 
                 if total <= split:
-                    result[i][root_file] = [start_entry, all_num_entries_dct[root_file]]
+                    split_result[i][root_file] = [entry_start, self.all_num_entries_dct[root_file]]
                     done.append(root_file)
 
                     if total == split:
@@ -233,45 +173,84 @@ class UprootIteratorMaker:
 
                 if total > split:
                     delta = num_entries - (total - split)
-                    result[i][root_file] = [start_entry, start_entry + delta]
+                    split_result[i][root_file] = [entry_start, entry_start + delta]
                     root_files_start_dct[root_file] += delta
                     num_entries_dct[root_file] -= delta
                     break
 
-        worker_uproot_iterators: list[list] = [[] for _ in range(self.num_workers)]
-        check_total = 0
+        return split_result
 
-        for i, result_dct in enumerate(result):
-            for root_file, start_end in result_dct.items():
-                entry_start, entry_end = start_end
-                check_total += entry_end - entry_start
+    def make(self) -> pd.DataFrame:
+        split_result = self._split()
 
-                iterator = self._make_iterator(root_files[root_file], entry_start, entry_end)
-                worker_uproot_iterators[i].append(iterator)
+        worker_df: dict[str, list] = {"worker_id": [], "file": [], "start": [], "stop": [], "step_size": []}
 
-        if check_total != total_num_entries:
+        check_total = 0
+        for i, result_dct in enumerate(split_result):
+            for root_file, start_stop in result_dct.items():
+                entry_start, entry_stop = start_stop
+                check_total += entry_stop - entry_start
+
+                delta_entry = entry_stop - entry_start
+                num_entries = self.all_num_entries_dct[root_file]
+
+                if self.step_size > delta_entry:
+                    step_size = delta_entry
+                elif self.step_size > num_entries:
+                    step_size = num_entries
+                else:
+                    step_size = self.step_size
+
+                worker_df["worker_id"].append(i)
+                worker_df["file"].append(f"{root_file}:{self.key}")
+                worker_df["start"].append(entry_start)
+                worker_df["stop"].append(entry_stop)
+                worker_df["step_size"].append(step_size)
+
+        if check_total != self.total_num_entries:
             raise ValueError("Total number of entries does not match.")
 
-        return worker_uproot_iterators
+        return pd.DataFrame(worker_df)
 
 
 class RootLoaderIterator:
     def __init__(
         self,
         name: str,
-        uproot_iterators: list[UprootIterator],
-        worker_id: int,
+        iterators_df: pd.DataFrame,
         processors: list[Callable[[ak.Array, dict], tuple[ak.Array, dict]]] | ProcessorsGraph | None = None,
+        filter_name: Callable[[str], bool] | None = None,
         root_files_desc_dct: dict[str, dict[str, Any]] | None = None,
     ) -> None:
         self.name = name
-        self.uproot_iterators = uproot_iterators
-        self.worker_id = worker_id
+        self.iterators_df = iterators_df
         self.processors = processors
+        self.filter_name = filter_name
         self.root_files_desc_dct = root_files_desc_dct
 
-        self.current_iterator: Iterator[tuple[ak.Array, Any]]
-        self.current_iterator_idx = 0
+        self.iterator: Iterator
+        self.tree: uproot.TTree
+
+        self.current_df_idx, self.current_iterator_idx = 0, 0
+
+    def _make_uproot_iterator(self, df: pd.Series) -> tuple[Iterator, uproot.TTree]:
+        tree = uproot.open(df["file"])
+
+        iterator = tree.iterate(
+            library="ak",
+            report=True,
+            step_size=df["step_size"],
+            filter_name=self.filter_name,
+            entry_start=df["start"],
+            entry_stop=df["stop"],
+        )
+
+        return iterator, tree
+
+    def _iterate_df(self) -> None:
+        df = self.iterators_df.iloc[self.current_iterator_idx]
+        self.iterator, self.tree = self._make_uproot_iterator(df)
+        self.current_df_idx += 1
 
     def _run_processors(self, arrays: ak.Array, reports: dict) -> tuple[ak.Array, dict] | dict[str, Processor]:
         if self.processors is None:
@@ -281,8 +260,7 @@ class RootLoaderIterator:
                 arrays, reports = proc(arrays, reports)
             return arrays, reports
         elif type(self.processors) is ProcessorsGraph:
-            processors = self.processors.fit(arrays, reports)
-            return processors
+            return self.processors.fit(arrays, reports)
         else:
             raise ValueError(f"Processors {self.processors} is not a valid type.")
 
@@ -293,7 +271,7 @@ class RootLoaderIterator:
 
         reports = {
             "name": self.name,
-            "worker_id": self.worker_id,
+            "worker_id": self.iterators_df.iloc[self.current_iterator_idx]["worker_id"],
             "file_path": file_path,
             "file": file_name,
             "start": start,
@@ -310,23 +288,22 @@ class RootLoaderIterator:
 
     def __next__(self) -> tuple[ak.Array, dict] | dict[str, Processor]:
         try:
-            uproot_iterator = self.uproot_iterators[self.current_iterator_idx]
+            if self.current_df_idx == self.current_iterator_idx:
+                self._iterate_df()
 
-            if uproot_iterator.returned_iterator is False:
-                self.current_iterator = uproot_iterator.get_iterator()
-
-            arrays, reports = next(self.current_iterator)
+            arrays, reports = next(self.iterator)
 
         except StopIteration:
-            uproot_iterator.close()
+            self.tree.close()
             self.current_iterator_idx += 1
 
-            if self.current_iterator_idx == len(self.uproot_iterators):
+            if self.current_iterator_idx == len(self.iterators_df):
                 raise StopIteration
-            else:
-                uproot_iterator = self.uproot_iterators[self.current_iterator_idx]
-                self.current_iterator = uproot_iterator.get_iterator()
-                arrays, reports = next(self.current_iterator)
+
+            if self.current_df_idx == self.current_iterator_idx:
+                self._iterate_df()
+
+            arrays, reports = next(self.iterator)
 
         reports = self._make_report(reports)
 
@@ -339,13 +316,15 @@ class RootIterableDataset(IterableDataset):
     def __init__(
         self,
         name: str,
-        worker_uproot_iterators: list[list[UprootIterator]],
+        worker_iterators_df: pd.DataFrame,
         processors: list[Callable[[ak.Array, dict], tuple[ak.Array, dict]]] | ProcessorsGraph | None = None,
+        filter_name: Callable[[str], bool] | None = None,
         root_files_desc_dct: dict[str, dict[str, Any]] | None = None,
     ) -> None:
         self.name = name
-        self.worker_uproot_iterators = worker_uproot_iterators
+        self.worker_iterators_df = worker_iterators_df
         self.processors = processors
+        self.filter_name = filter_name
         self.root_files_desc_dct = root_files_desc_dct
 
     def __iter__(self) -> RootLoaderIterator:
@@ -356,13 +335,13 @@ class RootIterableDataset(IterableDataset):
         else:
             worker_id = worker_info.id
 
-        uproot_iterators = self.worker_uproot_iterators[worker_id]
+        iterators_df = self.worker_iterators_df[self.worker_iterators_df["worker_id"] == worker_id].copy()
 
         return RootLoaderIterator(
             self.name,
-            uproot_iterators,
-            worker_id,
+            iterators_df,
             processors=self.processors,
+            filter_name=self.filter_name,
             root_files_desc_dct=self.root_files_desc_dct,
         )
 
@@ -376,46 +355,40 @@ def get_root_dataloader(
     processors: list[Callable[[ak.Array, dict], tuple[ak.Array, dict]]] | ProcessorsGraph | None = None,
     filter_name: Callable[[str], bool] | None = None,
     root_files_desc_dct: dict[str, dict[str, Any]] | None = None,
-    uproot_kwargs: dict[str, Any] | None = None,
     dataloader_kwargs: dict[str, Any] | None = None,
 ) -> tuple[DataLoader, int]:
 
+    if multiprocessing.get_start_method() == "fork" and num_workers > 0:
+        logging.debug("Using 'fork' start method. Consider using 'spawn' or 'forkserver'.")
+
     if num_workers == -1:
         num_workers = multiprocessing.cpu_count()
 
-    if uproot_kwargs is None:
-        uproot_kwargs = {}
-
     if dataloader_kwargs is None:
         dataloader_kwargs = {}
 
     logging.info("Making ROOT dataloader!")
 
-    uproot_iterator_maker = UprootIteratorMaker(
-        name,
-        files,
-        key,
-        step_size,
-        num_workers,
-        filter_name=filter_name,
-        **uproot_kwargs,
-    )
-    worker_uproot_iterators = uproot_iterator_maker.make()
+    df_maker = UprootIteratorsDfMaker(name, files, key, step_size, num_workers if num_workers > 0 else 1)
+    worker_iteratorts_df = df_maker.make()
 
-    total_num_entries = uproot_iterator_maker.total_num_entries
+    total_num_entries = df_maker.total_num_entries
 
     root_iterable_dataset = RootIterableDataset(
         name,
-        worker_uproot_iterators,
+        worker_iteratorts_df,
         processors=processors,
+        filter_name=filter_name,
         root_files_desc_dct=root_files_desc_dct,
     )
 
+    prefetch_factor = dataloader_kwargs.get("prefetch_factor", None)
+
     root_dataloader = DataLoader(
         root_iterable_dataset,
         batch_size=None,
         num_workers=num_workers,
-        prefetch_factor=None,
+        prefetch_factor=prefetch_factor,
         **dataloader_kwargs,
     )
 

{f9columnar-0.2.0 → f9columnar-0.2.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "f9columnar"
-version = "0.2.0"
+version = "0.2.2"
 description = "Columnar analysis utils."
 authors = ["Jan Gavranovic <jan.gavranovic@cern.ch>"]
 readme = "README.md"
@@ -28,6 +28,7 @@ requests = "^2.28.1"
 tqdm = "^4.65.0"
 networkx = "3.3"
 rich = "^13.7.0"
+psutil = "^6.1.0"
 mplhep = "^0.3.50"
 uproot = "5.3.10"
 awkward = "2.6.7"