Coverage for packages/dqm-ml-job/src/dqm_ml_job/dataloaders/parquet.py: 90%
159 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-07-21 08:27 +0000
1"""Parquet data loader for reading Parquet files.
3This module contains the ParquetDataLoader and ParquetDataSelection classes
4for loading and iterating over Parquet file data.
5"""
7import fnmatch
8import logging
9import os
10from typing import Any
12from dqm_ml_core.utils.matching import has_pattern, resolve_include_exclude, resolve_patterns
13import pyarrow as pa
14import pyarrow.compute as pc
15import pyarrow.fs as fs
16import pyarrow.parquet as pq
18# COMPATIBILITY : from typing import Any, override # When support of 3.10 and 3.11 will be removed
19from typing_extensions import override
21from dqm_ml_job.dataloaders.proto import DataSelection
23logger = logging.getLogger(__name__)
26def _fnmatch_to_regex(pattern: str) -> str:
27 """Convert fnmatch pattern to regex pattern.
29 Args:
30 pattern: fnmatch pattern with * and ? wildcards.
32 Returns:
33 Regex pattern string.
34 """
35 return fnmatch.translate(pattern)
38def _match_wildcard_arrow(col_expr: Any, patterns: list[str]) -> Any:
39 """Return pyarrow expression for wildcard matching.
41 Args:
42 col_expr: pyarrow field expression.
43 patterns: List of fnmatch patterns.
45 Returns:
46 pyarrow compute expression for OR of all patterns.
47 """
48 regex_patterns = [_fnmatch_to_regex(p) for p in patterns]
49 # Combine with OR
50 combined_regex = "|".join(f"({p})" for p in regex_patterns)
51 return pc.match_substring_regex(col_expr, combined_regex)
54def _resolve_pyarrow_type(type_name: str) -> Any:
55 """Map TransformType string to pyarrow DataType."""
56 import pyarrow as pa
58 mapping = {
59 "int32": pa.int32(),
60 "int64": pa.int64(),
61 "float32": pa.float32(),
62 "float64": pa.float64(),
63 "bool": pa.bool_(),
64 "str": pa.utf8(),
65 "categorical": pa.dictionary(pa.int32(), pa.utf8()),
66 }
67 return mapping[type_name]
70class ParquetDataSelection(DataSelection):
71 """A specific selection of data from a Parquet dataset.
73 This class represents a filtered subset of a Parquet dataset
74 and provides an iterator over PyArrow RecordBatches.
76 Attributes:
77 name: Name identifier for this selection.
78 path: Path to the Parquet file or directory.
79 batch_size: Number of rows per batch.
80 threads: Number of threads for parallel reading.
81 filters_dict: Optional dictionary of column filters to apply.
82 filesystem: Optional PyArrow filesystem for reading.
83 sample_path: List of sample path configs describing column path prefixes.
84 transforms: List of transform configs (column cast operations).
85 """
87 def __init__(
88 self,
89 name: str,
90 path: str,
91 batch_size: int = 100_000,
92 threads: int = 4,
93 filters_dict: dict[str, Any] | None = None,
94 filesystem: Any | None = None,
95 sample_path: list[dict[str, Any]] | None = None,
96 transforms: list[dict[str, Any]] | None = None,
97 ):
98 """Initialize a Parquet data selection.
100 Args:
101 name: Name identifier for this selection.
102 path: Path to the Parquet file or directory.
103 batch_size: Number of rows per batch (default: 100000).
104 threads: Number of threads for parallel reading (default: 4).
105 filters_dict: Optional dictionary of column filters to apply.
106 filesystem: Optional PyArrow filesystem for reading.
107 sample_path: List of sample path configs describing column path prefixes.
108 transforms: List of transform configs (column cast operations).
109 """
110 self.name = name
111 self.path = path
112 self.batch_size = batch_size
113 self.threads = threads
114 self.filters_dict = filters_dict
115 self.filesystem = filesystem
116 self.sample_path = sample_path or []
117 self.transforms = transforms or []
118 self.columns_list: list[str] | None = None
119 self.dataset: pq.ParquetDataset | None = None
120 self.samples_count: int = 0
122 def _build_filter_expr(self) -> Any:
123 """Build a PyArrow filter expression from the filters dictionary.
125 Converts the filters_dict configuration into a combined pyarrow
126 compute expression using AND logic across all filter conditions.
128 Returns:
129 PyArrow compute expression for filtering, or None if no filters.
130 """
131 if self.filters_dict is None:
132 return None
133 expr = None
134 for col, val in self.filters_dict.items():
135 if self.columns_list is None: 135 ↛ 136line 135 didn't jump to line 136 because the condition on line 135 was never true
136 self.columns_list = [col]
137 elif col not in self.columns_list:
138 self.columns_list.append(col)
139 from dqm_ml_job.dataloaders.filters import build_filter_condition
141 col_expr = build_filter_condition(
142 col,
143 val,
144 wildcard_fn=lambda c, vals: _match_wildcard_arrow(pc.field(c), vals),
145 isin_fn=lambda c, vals: pc.is_in(pc.field(c), pa.array(vals)),
146 equal_fn=lambda c, v: pc.equal(pc.field(c), v),
147 )
148 expr = col_expr if expr is None else (expr & col_expr)
149 return expr
151 @override
152 def bootstrap(self, columns_list: list[str]) -> None:
153 """Initialize the parquet dataset and filter expression.
155 Args:
156 columns_list: Names of columns to load from the parquet file.
157 Empty list means read all columns.
158 """
159 self.columns_list = columns_list or None
160 logger.debug(f"[DEBUG] ParquetDataSelection.bootstrap: {self.name} received columns_list = {columns_list}")
161 self.filter_expr = self._build_filter_expr()
162 logger.debug(f"[DEBUG] ParquetDataSelection.bootstrap: filter_expr = {self.filter_expr}")
163 self.dataset = pq.ParquetDataset(self.path, filters=self.filter_expr, filesystem=self.filesystem)
164 if len(self.dataset.fragments) > 0:
165 self.samples_count = sum(p.count_rows() for p in self.dataset.fragments)
166 else:
167 self.samples_count = 0
169 def __len__(self) -> int:
170 return int(self.samples_count)
172 @override
173 def get_nb_batches(self) -> int:
174 """Return the estimated number of batches in this selection.
176 Returns:
177 Number of batches based on total samples and batch size.
178 """
179 return int(len(self) / self.batch_size) + (len(self) % self.batch_size > 0)
181 @override
182 def __iter__(self) -> Any:
183 if self.dataset is None: 183 ↛ 184line 183 didn't jump to line 184 because the condition on line 183 was never true
184 return
185 logger.debug(f"[DEBUG] ParquetDataSelection.__iter__: {self.name} using columns_list = {self.columns_list}")
186 for file in self.dataset.files:
187 parquet_file = pq.ParquetFile(file, filesystem=self.filesystem)
188 batch_iterator = parquet_file.iter_batches(
189 batch_size=self.batch_size,
190 columns=self.columns_list,
191 use_threads=self.threads,
192 )
193 for batch in batch_iterator:
194 if self.filter_expr is not None:
195 batch = batch.filter(self.filter_expr)
196 if len(batch) == 0:
197 continue
198 batch = self._apply_transforms(batch)
199 logger.debug(
200 "[DEBUG] ParquetDataSelection.__iter__: %s yielded batch with columns = %s",
201 self.name,
202 batch.schema.names,
203 )
204 yield batch
206 def _apply_transforms(self, batch: pa.RecordBatch) -> pa.RecordBatch:
207 """Apply column transforms (cast operations) to the batch.
209 For each transform entry:
210 - If ``in_place``, overwrite the column in-place.
211 - Otherwise, append a new column named ``<column>_<to_type>``.
212 """
213 for t in self.transforms:
214 col_idx = batch.schema.get_field_index(t["column"])
215 if col_idx == -1:
216 continue
217 target_type = _resolve_pyarrow_type(t["to_type"])
218 cast_col = batch.column(col_idx).cast(target_type)
219 if t.get("in_place", False):
220 batch = batch.set_column(col_idx, t["column"], cast_col)
221 else:
222 new_name = f"{t['column']}_{t['to_type']}"
223 batch = batch.append_column(pa.field(new_name, target_type), cast_col)
224 return batch
226 @override
227 def __repr__(self) -> str:
228 return f"ParquetSelection(name='{self.name}', path='{self.path}', filters={self.filters_dict})"
231class ParquetDataLoader:
232 """Data loader for Parquet files that generates one or more DataSelections.
234 This loader can read from a single Parquet file or a directory of Parquet
235 files, optionally splitting the data by a column value to create multiple
236 selections.
238 Attributes:
239 type: The loader type identifier ("parquet").
240 filesystem: Optional PyArrow filesystem for reading.
241 """
243 type: str = "parquet"
245 def __init__(self, name: str, config: dict[str, Any] | None = None):
246 """Initialize the Parquet data loader.
248 Args:
249 name: Unique name for this loader instance.
250 config: Configuration dictionary containing:
251 - path: Path to Parquet file or directory (required)
252 - batch_size: Rows per batch (default: 100000)
253 - threads: Number of threads (default: 4)
254 - split_by: Column name to split selections by
255 - split_values: Specific values to split on
256 - filter.: list of filters
257 - storage: Storage configuration (bool or dict)
259 Raises:
260 ValueError: If required config keys are missing.
261 """
262 if config is None: 262 ↛ 263line 262 didn't jump to line 263 because the condition on line 262 was never true
263 config = {}
264 self.name = name
265 self.config = config
266 self.path: str = config["path"]
267 self.batch_size = config.get("batch_size", 100_000)
268 self.threads = config.get("threads", 4)
269 # Use SplitConfig model fields instead of hardcoded keys
270 from dqm_ml_core.models.dataloaders import SplitConfig
272 split = config.get("split")
273 self.split = SplitConfig.model_validate(split) if split else None
274 self.split_by = self.split.by if self.split else None
275 self.split_values = self.split.values if self.split else None
276 filters = config.get("filters")
277 # transform the list of dict into a dict
278 self.filters_dict = {}
279 if filters is not None:
280 for item in filters:
281 column = item["column"]
282 self.filters_dict[column] = item["values"]
283 logger.debug(f"[DEBUG] ParquetDataLoader.__init__: filters_dict = {self.filters_dict}")
285 self.id_column = config.get("id_column")
286 self.sample_path = config.get("sample_path", [])
287 self.transforms = config.get("transform", [])
289 # Storage filesystem configuration - only for S3 paths, not local paths
290 self.filesystem = None
291 storage_config = None
292 storage_cfg = config.get("storage")
293 if storage_cfg: 293 ↛ 295line 293 didn't jump to line 295 because the condition on line 293 was never true
294 # Use StorageConfig model to validate and access fields
295 from dqm_ml_core.models.global_ import StorageConfig
297 storage_config = StorageConfig.model_validate(storage_cfg)
299 self.storage_config = storage_config
301 if storage_config.type == "s3":
302 from dqm_ml_job.utils.s3 import get_s3_filesystem
304 self.filesystem = get_s3_filesystem(storage_config)
306 def _resolve_selection_path(self) -> str:
307 """Resolve the full path, prepending S3 bucket if applicable."""
308 path = self.path
309 if self.filesystem is not None and isinstance(self.filesystem, fs.S3FileSystem): 309 ↛ 310line 309 didn't jump to line 310 because the condition on line 309 was never true
310 bucket_name = self.storage_config.bucket if self.storage_config else os.getenv("S3_BUCKET_NAME", "")
311 if bucket_name and not path.startswith(bucket_name + "/"):
312 path = f"{bucket_name}/{path}"
313 return path
315 def get_selections(self) -> list[DataSelection]:
316 """Create one or more ParquetDataSelection instances based on configuration.
318 Returns:
319 A list of DataSelection instances. If split_by is configured,
320 returns one selection per unique value. Otherwise, returns a
321 single selection for the entire dataset.
322 """
323 path = self._resolve_selection_path()
325 if not self.split_by:
326 # Single selection
327 return [
328 ParquetDataSelection(
329 name=self.name,
330 path=path,
331 batch_size=self.batch_size,
332 threads=self.threads,
333 filters_dict=self.filters_dict,
334 filesystem=self.filesystem,
335 sample_path=self.sample_path,
336 transforms=self.transforms,
337 )
338 ]
340 # Splitting logic
341 values = self.split_values
342 if values is None:
343 # Automatic discovery if split_values not provided
344 logger.info(f"Discovering unique values for split_by='{self.split_by}' in {path}")
345 table = pq.read_table(path, columns=[self.split_by], filesystem=self.filesystem)
346 values = [str(v) for v in pc.unique(table.column(0)).to_pylist() if v is not None]
347 else:
348 # Expand wildcard patterns in values against available data
349 if any(has_pattern(v) for v in values):
350 logger.info(f"Expanding wildcard values for split_by='{self.split_by}' in {path}")
351 table = pq.read_table(path, columns=[self.split_by], filesystem=self.filesystem)
352 available = [str(v) for v in pc.unique(table.column(0)).to_pylist() if v is not None]
353 values = resolve_patterns(values, available)
355 # Apply split.exclude (including wildcard patterns)
356 if self.split and self.split.exclude:
357 values = resolve_include_exclude(None, self.split.exclude, values)
359 selections: list[DataSelection] = []
360 for val in values:
361 selection_name = f"{self.name}_{val}"
362 # Merge existing filters with the split filter
363 merged_filters = (self.filters_dict or {}).copy()
364 # TODO: filters shouldn't be on the same column than split by
365 # raise error here ? or add this check in pydantic model ? is it possible ?
366 merged_filters[self.split_by] = val
368 selections.append(
369 ParquetDataSelection(
370 name=selection_name,
371 path=path,
372 batch_size=self.batch_size,
373 threads=self.threads,
374 filters_dict=merged_filters,
375 filesystem=self.filesystem,
376 sample_path=self.sample_path,
377 transforms=self.transforms,
378 )
379 )
380 return selections