Coverage for packages/dqm-ml-job/src/dqm_ml_job/job.py: 98%
333 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-07-21 08:27 +0000
1"""Dataset job orchestrator for end-to-end data quality assessment.
3This module contains the DatasetJob class that orchestrates the complete
4pipeline: data loading, metric computation, and result persistence.
5"""
7from collections.abc import Sequence
8import fnmatch
9import itertools
10import logging
11from typing import Any
13from dqm_ml_core.api.features_processor import FeaturesProcessor
14from dqm_ml_core.api.gap_processor import GapProcessor
15from dqm_ml_core.api.metrics_processor import MetricsProcessor
16from dqm_ml_core.api.processor import Processor
17from dqm_ml_core.utils.matching import has_pattern, resolve_include_exclude
18import numpy as np
19import pyarrow as pa
20from tqdm import tqdm
22from dqm_ml_job.dataloaders import DataLoader, DataSelection
23from dqm_ml_job.outputwriter import OutputWriter
25logger = logging.getLogger(__name__)
28class DatasetJob:
29 """
30 Orchestrates the end-to-end data quality assessment process.
32 The job handles:
33 1. Plugin discovery and component initialization.
34 2. Data selection discovery via DataLoaders.
35 3. Streaming execution: Iterating over selections and batches to
36 compute features and metrics.
37 4. Result persistence via OutputWriters.
38 5. Comparison metrics (deltas) between discovered datasets.
39 """
41 def __init__(
42 self,
43 dataloaders: dict[str, DataLoader],
44 features_processors: dict[str, FeaturesProcessor] | None = None,
45 metrics_processors: dict[str, MetricsProcessor] | None = None,
46 gap_processors: dict[str, GapProcessor] | None = None,
47 features_output: OutputWriter | None = None,
48 progress_bar: bool = True,
49 threads: int = 4,
50 errors_by_interface: dict[str, Any] | None = None,
51 compute_seed: int | None = None,
52 compute_device: str = "auto",
53 compute_max_memory: str | None = None,
54 ) -> None:
55 """
56 Initialize the pipeline components.
58 Args:
59 dataloaders: Map of initialized DataLoader instances.
60 features_processors: Map of feature extraction processors.
61 metrics_processors: Map of metric computation processors.
62 gap_processors: Map of domain gap processors.
63 features_output: Optional writer for persisting per-sample features.
64 progress_bar: Whether to display execution progress in the terminal.
65 threads: Number of threads for parallel processing.
66 errors_by_interface: Per-interface error configuration.
67 compute_seed: Seed for reproducible RNG in processors.
68 compute_device: Device hint ("auto", "cpu", "cuda") for processors.
69 compute_max_memory: Optional max memory string (e.g. "2GB") for features flushing.
70 """
71 # We initialize loaded pluging elements
72 self.dataloaders = dataloaders
73 self.features_processors = features_processors or {}
74 self.metrics_processors = metrics_processors or {}
75 self.gap_processors = gap_processors or {}
76 self.features_output = features_output
77 self.progress_bar = progress_bar
78 self.threads = threads
79 self.errors_by_interface = errors_by_interface or {}
80 self.compute_max_memory = compute_max_memory
82 self._resolve_output_columns()
83 self._analyze_processor_columns()
85 # Inject per-interface errors into processors
86 self._inject_per_interface_errors()
88 # Inject compute config into processors
89 self._inject_per_interface_compute(compute_seed, compute_device)
91 logger.info(
92 f"DQM job pipeline initialized will process "
93 f"{len(self.dataloaders)} dataloaders, "
94 f"{len(self.features_processors) + len(self.metrics_processors) + len(self.gap_processors)} processors, "
95 f"{1 if self.features_output else 0} output writers"
96 )
98 @property
99 def all_processors(self) -> dict[str, Processor]:
100 """Return all processors across all interfaces."""
101 return {
102 **self.features_processors,
103 **self.metrics_processors,
104 **self.gap_processors,
105 }
107 def _resolve_output_columns(self) -> None:
108 """Resolve features_output include/exclude columns from the output writer config."""
109 self.features_output_include = None
110 self.features_output_exclude = None
111 if not self.features_output:
112 return
113 self.features_output_include = self.features_output.columns or None
114 self.features_output_exclude = getattr(self.features_output, "exclude", None)
116 def _analyze_processor_columns(self) -> None:
117 """Collect needed input columns, generated features, and generated metrics from all processors."""
118 self.needed_input_columns = []
119 self.generated_features = []
120 self.generated_metrics = []
121 self._has_wildcard_columns = False
122 for proc in self.all_processors.values():
123 cols = proc.needed_columns()
124 self.needed_input_columns.extend(cols)
125 if not self._has_wildcard_columns: 125 ↛ 122line 125 didn't jump to line 122 because the condition on line 125 was always true
126 self._has_wildcard_columns = any(has_pattern(c) for c in cols)
128 for proc in self.features_processors.values():
129 self.generated_features.extend(proc.generated_features())
131 for proc in self.metrics_processors.values():
132 self.generated_metrics.extend(proc.generated_metrics())
134 self.needed_input_columns = list(dict.fromkeys(self.needed_input_columns))
135 self.generated_features = list(dict.fromkeys(self.generated_features))
136 self.generated_metrics = list(dict.fromkeys(self.generated_metrics))
138 if self._has_wildcard_columns:
139 self.needed_input_columns = []
141 if not self.features_output_include:
142 return
143 for col in self.features_output_include:
144 if has_pattern(col):
145 continue
146 if col not in self.generated_features: 146 ↛ 143line 146 didn't jump to line 143 because the condition on line 146 was always true
147 logger.info(f"Adding required output column '{col}' to input columns")
148 self.needed_input_columns.insert(0, col)
150 @staticmethod
151 def _get_interface_for_processor(processor: Processor) -> str | None:
152 """Determine which interface a processor belongs to.
154 Args:
155 processor: The processor instance.
157 Returns:
158 Interface name ("features", "metrics", "gap") or None if unknown.
159 """
160 if isinstance(processor, FeaturesProcessor):
161 return "features"
162 elif isinstance(processor, MetricsProcessor):
163 return "metrics"
164 elif isinstance(processor, GapProcessor):
165 return "gap"
166 return None
168 def _inject_per_interface_errors(self) -> None:
169 """Inject per-interface errors into processors based on their interface."""
170 for proc in self.all_processors.values():
171 interface = self._get_interface_for_processor(proc)
172 if interface and interface in self.errors_by_interface:
173 proc.errors_config = self.errors_by_interface[interface]
175 def _inject_per_interface_compute(self, compute_seed: int | None, compute_device: str) -> None:
176 """Inject compute config into processors for device and seed."""
177 for proc in self.all_processors.values():
178 proc.compute_device = compute_device
179 if compute_seed is not None:
180 proc.compute_seed = compute_seed
182 def get_ordered_processors(self) -> list[Processor]:
183 """
184 Return the list of all processors in dependency order.
186 Processors that generate columns (via ``generated_features()`` or
187 ``generated_columns()``) are placed before processors that depend on
188 those columns (via ``needed_columns()``). This ensures, for example,
189 that an ``image_embedding`` processor that produces the ``embedding``
190 column runs before a ``domain_gap`` processor that consumes it,
191 regardless of the order in which they appear in the YAML config.
192 """
193 procs = list(self.all_processors.values())
194 if len(procs) <= 1:
195 return procs
197 dep_on = self._build_dependency_graph(procs)
198 return self._topological_sort(procs, dep_on)
200 @staticmethod
201 def _register_generated_columns(
202 procs: list[Processor],
203 ) -> dict[str, set[int]]:
204 """Build a mapping from column names to the processor indices that generate them.
206 Args:
207 procs: List of metric processors.
209 Returns:
210 Dict mapping column names to sets of processor indices.
211 """
212 generated_by: dict[str, set[int]] = {}
213 for i, p in enumerate(procs):
214 if hasattr(p, "generated_features"):
215 for col in p.generated_features():
216 generated_by.setdefault(col, set()).add(i)
217 return generated_by
219 @staticmethod
220 def _resolve_dependency_col(
221 col: str,
222 generated_names: list[str],
223 generated_by: dict[str, set[int]],
224 exclude_idx: int,
225 ) -> set[int]:
226 """Resolve processor dependencies for a required column.
228 Matches the column pattern against generated column names and returns
229 indices of processors that produce matching columns (excluding self).
231 Args:
232 col: Required column name (may contain fnmatch patterns).
233 generated_names: List of all column names generated by any processor.
234 generated_by: Mapping from column name to set of processor indices.
235 exclude_idx: Index of the processor requesting the dependency (excluded).
237 Returns:
238 Set of processor indices that generate matching columns.
239 """
240 matching_cols = fnmatch.filter(generated_names, col) if has_pattern(col) else [col]
241 deps: set[int] = set()
242 for gen_col in matching_cols:
243 for gen_idx in generated_by.get(gen_col, ()):
244 if gen_idx != exclude_idx: 244 ↛ 243line 244 didn't jump to line 243 because the condition on line 244 was always true
245 deps.add(gen_idx)
246 return deps
248 @staticmethod
249 def _build_dependency_graph(procs: list[Processor]) -> list[set[int]]:
250 """Build a dependency graph from a list of processors.
252 Args:
253 procs: List of processors.
255 Returns:
256 List of sets where dep_on[i] contains indices of processors
257 that processor i depends on.
258 """
259 generated_by = DatasetJob._register_generated_columns(procs)
260 generated_names = list(generated_by.keys())
261 dep_on: list[set[int]] = [set() for _ in procs]
262 for i, p in enumerate(procs):
263 for col in p.needed_columns():
264 dep_on[i] |= DatasetJob._resolve_dependency_col(col, generated_names, generated_by, i)
266 return dep_on
268 @staticmethod
269 def _topological_sort(procs: list[Processor], dep_on: list[set[int]]) -> list[Processor]:
270 """Topological sort of processors using Kahn's algorithm.
272 Args:
273 procs: List of processors.
274 dep_on: Dependency graph as produced by _build_dependency_graph.
276 Returns:
277 Processors in dependency order.
278 """
279 ordered: list[Processor] = []
280 remaining = set(range(len(procs)))
281 while remaining:
282 ready = {i for i in remaining if not (dep_on[i] & remaining)}
283 if not ready:
284 ready = {min(remaining)}
285 for i in sorted(ready):
286 ordered.append(procs[i])
287 remaining.remove(i)
288 return ordered
290 def describe(self, selections: list[DataSelection]) -> None:
291 """Log a summary of the execution plan, including selections and metrics."""
292 total = len(self.all_processors)
293 logger.info(f"Executing dqm-ml-job on {len(selections)} selections, using {total} processors ")
295 for selection in selections:
296 logger.info(f" Selection: {selection.name} -> {selection}")
298 for proc_name, proc in self.all_processors.items():
299 logger.info(f" Processor: {proc_name} -> {proc}")
300 logger.info(f" Needed columns: {proc.needed_columns()}")
301 if isinstance(proc, FeaturesProcessor):
302 logger.info(f" Generated features: {proc.generated_features()}")
303 elif isinstance(proc, MetricsProcessor):
304 logger.info(f" Generated metrics: {proc.generated_metrics()}")
306 def _discover_selections(self) -> list[DataSelection]:
307 """Discover all data selections from all configured dataloaders.
309 Returns:
310 List of DataSelection instances.
311 """
312 all_selections: list[DataSelection] = []
313 for loader in self.dataloaders.values():
314 all_selections.extend(loader.get_selections())
315 return all_selections
317 def _compute_selection_metrics(
318 self,
319 selection_name: str,
320 batches_metrics_array: dict[str, Any],
321 metrics_processors: Sequence[MetricsProcessor | GapProcessor],
322 ) -> dict[str, Any]:
323 """Compute dataset-level metrics for a single selection.
325 Args:
326 selection_name: Name of the selection.
327 batches_metrics_array: Accumulated batch metrics.
328 metrics_processors: List of processors.
330 Returns:
331 Dictionary of computed dataset metrics.
332 """
333 dataset_metrics: dict[str, Any] = {}
334 metrics_iter = (
335 tqdm(metrics_processors, desc="metrics", position=1, leave=False)
336 if self.progress_bar
337 else metrics_processors
338 )
339 for metric in metrics_iter:
340 if logging.getLogger().level == logging.DEBUG:
341 logger.debug(f"Metric computation {metric.__class__.__name__} for dataselection {selection_name}")
342 dataset_metrics.update(metric.compute(batch_metrics=batches_metrics_array))
343 if logging.getLogger().level == logging.DEBUG:
344 logger.debug(f"Available metrics {list(dataset_metrics.keys())}")
345 return dataset_metrics
347 def run(self) -> tuple[dict[Any, dict[str, Any]], pa.Table | None]:
348 """
349 Execute the job on all discovered data selections.
351 This is the main entry point for execution. It iterates through every
352 selection found by the loaders, computes statistics, and finally
353 calculates deltas between datasets.
355 Returns:
356 A tuple containing:
357 - Mapping of selection names to their final metric dictionaries.
358 - pyarrow Table containing all computed deltas.
359 """
360 ordered_processors = self.get_ordered_processors()
361 all_selections = self._discover_selections()
363 self.describe(all_selections)
365 dataselection_metrics_list: dict[Any, dict[str, Any]] = {}
366 job_iter = tqdm(all_selections, desc="selection", position=0) if self.progress_bar else all_selections
368 for selection in job_iter:
369 selection_name = selection.name
370 logger.info(f"Processing selection '{selection_name}'")
372 selection.bootstrap(self.needed_input_columns)
373 batches_metrics_array = self._compute_batches_metrics(selection_name, selection, ordered_processors)
375 metrics_and_gap = list(self.metrics_processors.values()) + list(self.gap_processors.values())
376 dataset_metrics = self._compute_selection_metrics(selection_name, batches_metrics_array, metrics_and_gap)
377 dataselection_metrics_list[selection_name] = dataset_metrics
379 # Reset processor state between selections — processors like
380 # RepresentativenessProcessor cache per-selection state (e.g.
381 # quantile bin edges). Without a reset those cached values
382 # leak across selections and produce NaN/incorrect results
383 # when the next selection's distribution differs from the first
384 # one that was processed. See AGENTS.md for background.
385 for proc in ordered_processors:
386 proc.reset()
388 gap_list = list(self.gap_processors.values())
389 delta_metrics_table = self._compute_delta_metrics(gap_list, dataselection_metrics_list)
391 if self.features_output and hasattr(self.features_output, "flush"):
392 self.features_output.flush()
394 return dataselection_metrics_list, delta_metrics_table
396 @staticmethod
397 def _to_pa_array(value: Any, key: str) -> pa.Array:
398 """Convert a delta metric value to PyArrow array.
400 Args:
401 value: The value to convert (float, int, str, np.ndarray, or pa.Array).
402 key: The metric name for error logging.
404 Returns:
405 PyArrow array containing the value.
407 Raises:
408 TypeError: If the value type is not supported.
409 """
410 if isinstance(value, pa.Array):
411 return value
412 elif isinstance(value, (int, float, np.number)):
413 return pa.array([float(value)])
414 elif isinstance(value, str):
415 return pa.array([value])
416 elif isinstance(value, np.ndarray):
417 return pa.array([value.tolist()])
418 else:
419 logger.error(f"Cannot convert delta metric '{key}' to pa.Array: type={type(value)}")
420 raise TypeError(f"Unsupported delta metric type: {type(value)} for key '{key}'")
422 def _compute_delta_metrics(
423 self,
424 metrics_processors: Sequence[GapProcessor],
425 dataselection_metrics_list: dict[str, dict[str, Any]],
426 ) -> pa.Table | None:
427 """Compute comparison metrics between every unique pair of data selections.
429 Builds a single table with one row per (pair, metric) combination.
430 Different metric processors may produce different columns; missing
431 values are padded with nulls via ``pa.concat_tables``.
433 Args:
434 metrics_processors: List of processors capable of computing deltas.
435 dataselection_metrics_list: Map of selection names to their metrics.
437 Returns:
438 A pyarrow Table with one row per (pair, metric) combination.
439 """
441 selection_combinations = itertools.combinations(dataselection_metrics_list, 2)
443 tables: list[pa.Table] = []
444 for combination in selection_combinations:
445 src_metrics = dataselection_metrics_list[combination[0]]
446 target_metrics = dataselection_metrics_list[combination[1]]
448 for metric in metrics_processors:
449 delta_metrics = metric.compute_delta(src_metrics, target_metrics)
451 if len(delta_metrics) == 0: 451 ↛ 452line 451 didn't jump to line 452 because the condition on line 451 was never true
452 continue
454 row = {key: self._to_pa_array(value, key) for key, value in delta_metrics.items()}
455 row["selection_source"] = pa.array([combination[0]])
456 row["selection_target"] = pa.array([combination[1]])
457 tables.append(pa.table(row))
459 if not tables:
460 return None
462 return pa.concat_tables(tables, promote_options="default")
464 @staticmethod
465 def _process_batch(
466 batch: Any,
467 ordered_processors: list[Processor],
468 metrics_processors: list[MetricsProcessor],
469 gap_processors: list[GapProcessor],
470 ) -> tuple[dict[str, Any], dict[str, Any]]:
471 """Compute features and batch-level metrics for a single batch.
473 Two-phase dispatch:
474 1. Extract data from the batch using the interface-specific method.
475 2. Compute batch-level metrics only for metrics and gap processors.
477 Args:
478 batch: Input data batch.
479 ordered_processors: All processors in dependency order.
480 metrics_processors: List of metric processors.
481 gap_processors: List of gap processors.
483 Returns:
484 Tuple of (batch_features, batch_metrics).
485 """
486 batch_features: dict[str, Any] = {}
487 batch_metrics: dict[str, Any] = {}
489 # Phase 1: Extract data from batch (interface-specific method)
490 for proc in ordered_processors:
491 if isinstance(proc, FeaturesProcessor):
492 batch_features.update(proc.compute_features(batch, prev_features=batch_features))
493 elif isinstance(proc, MetricsProcessor):
494 batch_features.update(proc.select_columns(batch, prev_features=batch_features))
495 elif isinstance(proc, GapProcessor): 495 ↛ 490line 495 didn't jump to line 490 because the condition on line 495 was always true
496 batch_features.update(proc.select_features(batch, prev_features=batch_features))
498 # Phase 2: Batch-level metric aggregation (only metrics and gap)
499 for proc in metrics_processors:
500 batch_metrics.update(proc.compute_batch_metric(batch_features))
501 for proc in gap_processors:
502 batch_metrics.update(proc.compute_batch_metric(batch_features))
504 if logging.getLogger().level == logging.DEBUG:
505 m_keys, m_features = (
506 list(batch_metrics.keys()),
507 list(batch_features.keys()),
508 )
509 logger.debug(f"Available batch_metrics {m_keys} - features {m_features}")
511 return batch_features, batch_metrics
513 def _accumulate_source_features(
514 self,
515 batch: Any,
516 features_accumulator: dict[str, list[Any]],
517 feature_array_size: int,
518 ) -> int:
519 """Accumulate source dataset columns into the features accumulator.
521 Args:
522 batch: Input data batch.
523 features_accumulator: Dict accumulating feature lists.
524 feature_array_size: Current memory usage estimate.
526 Returns:
527 Updated feature_array_size.
528 """
529 if self.features_output is None:
530 return feature_array_size
532 available = batch.column_names
533 keep = resolve_include_exclude(
534 self.features_output_include,
535 self.features_output_exclude,
536 available,
537 )
539 for col_name in keep:
540 col_data = batch.column(col_name)
541 if col_name not in features_accumulator:
542 features_accumulator[col_name] = []
543 features_accumulator[col_name].append(col_data)
544 feature_array_size += col_data.get_total_buffer_size()
545 return feature_array_size
547 def _accumulate_generated_features(
548 self,
549 batch: Any,
550 batch_features: dict[str, Any],
551 batch_metrics: dict[str, Any],
552 features_accumulator: dict[str, list[Any]],
553 feature_array_size: int,
554 ) -> int:
555 """Accumulate generated features into the features accumulator.
557 Args:
558 batch: Input data batch (used to identify source columns).
559 batch_features: Features generated by processors.
560 batch_metrics: Metrics generated by processors.
561 features_accumulator: Dict accumulating feature lists.
562 feature_array_size: Current memory usage estimate.
564 Returns:
565 Updated feature_array_size.
566 """
567 if self.features_output is None:
568 return feature_array_size
570 # Generated features are always included in the output.
571 # The include/exclude filter applies only to source columns
572 # (handled in _accumulate_source_features).
573 source_cols = set(batch.schema.names)
575 for k, v in batch_features.items():
576 if k in batch_metrics or k in source_cols:
577 continue
578 if k not in features_accumulator:
579 features_accumulator[k] = []
580 features_accumulator[k].append(v)
581 feature_array_size += v.get_total_buffer_size()
582 return feature_array_size
584 def _maybe_flush_features(
585 self,
586 selection_name: str,
587 features_accumulator: dict[str, list[Any]],
588 feature_array_size: int,
589 part_index: int,
590 memory_threshold: int,
591 ) -> int:
592 """Flush features to disk if memory threshold is exceeded.
594 Args:
595 selection_name: Name of the current data selection.
596 features_accumulator: Dict accumulating feature lists (mutated in place on flush).
597 feature_array_size: Current memory usage estimate.
598 part_index: Current chunk index.
599 memory_threshold: Memory threshold in bytes.
601 Returns:
602 Updated part_index (incremented if flush occurred).
603 """
604 if feature_array_size <= memory_threshold or not self.features_output:
605 return part_index
607 logger.info(f"Memory threshold reached ({feature_array_size / 1024**2:.1f}MB). Flushing chunk {part_index}")
608 features_chunk: dict[str, Any] = {}
609 for k, v_list in features_accumulator.items():
610 features_chunk[k] = pa.concat_arrays(v_list)
612 self._inject_dataloader_column(selection_name, features_chunk)
613 self.features_output.write_table(selection_name, features_chunk, part_index)
614 features_accumulator.clear()
615 return part_index + 1
617 def _write_remaining_features(
618 self,
619 selection_name: str,
620 features_accumulator: dict[str, list[Any]],
621 part_index: int,
622 ) -> None:
623 """Concatenate and write remaining features that were never flushed.
625 Args:
626 selection_name: Name of the current data selection.
627 features_accumulator: Dict accumulating feature lists.
628 part_index: Current chunk index.
629 """
630 if not self.features_output or not features_accumulator:
631 return
633 features_array: dict[str, Any] = {}
634 for k, v_list in features_accumulator.items():
635 features_array[k] = pa.concat_arrays(v_list)
637 self._inject_dataloader_column(selection_name, features_array)
638 self.features_output.write_table(selection_name, features_array, part_index)
640 @staticmethod
641 def _concatenate_accumulator(
642 accumulator: dict[str, list[Any]],
643 ) -> dict[str, Any]:
644 """Concatenate lists of arrays into a single dict of arrays."""
645 return {k: pa.concat_arrays(v) for k, v in accumulator.items()}
647 @staticmethod
648 def _inject_path_prefixes(selection: DataSelection, processors: list[Processor]) -> None:
649 """Build per-column path prefix map from selection's sample_path config and inject into processors."""
650 prefix_map: dict[str, str] = {}
651 for entry in getattr(selection, "sample_path", []):
652 col = entry.get("column")
653 if col and entry.get("prefix"):
654 prefix_map[col] = entry["prefix"]
655 for proc in processors:
656 proc.current_path_prefix = prefix_map
658 @staticmethod
659 def _clear_path_prefixes(processors: list[Processor]) -> None:
660 """Clear per-selection path prefix state from processors."""
661 for proc in processors:
662 if hasattr(proc, "current_path_prefix"): 662 ↛ 661line 662 didn't jump to line 661 because the condition on line 662 was always true
663 del proc.current_path_prefix
665 def _compute_batches_metrics(
666 self,
667 selection_name: str,
668 selection: DataSelection,
669 ordered_processors: list[Processor],
670 ) -> dict[str, Any]:
671 """Process all batches to compute intermediate statistics and features.
673 Memory Management:
674 - Batch-level statistics (`batch_metrics`) are accumulated in lists
675 and concatenated once the selection is complete.
676 - Per-sample features are also accumulated in memory before being
677 passed to the OutputWriter.
678 - NOTE: For large datasets, accumulation can lead to high memory
679 usage. Future versions will implement disk-flushing (chunking).
681 Args:
682 selection_name: Name of the current data selection.
683 selection: The selection iterator.
684 ordered_processors: All processors in dependency order.
686 Returns:
687 Dictionary of concatenated intermediate statistics arrays.
688 """
689 self._inject_path_prefixes(selection, ordered_processors)
691 batch_metrics_accumulator: dict[str, list[Any]] = {}
692 features_accumulator: dict[str, list[Any]] = {}
693 feature_array_size = 0
694 part_index = 0
696 compute_max_memory = getattr(self, "compute_max_memory", None)
697 memory_threshold = self._parse_memory_string(compute_max_memory) if compute_max_memory else 512 * 1024 * 1024
699 dataloader_iter = (
700 tqdm(
701 selection,
702 desc="batches",
703 position=1,
704 leave=False,
705 total=selection.get_nb_batches(),
706 )
707 if self.progress_bar
708 else selection
709 )
711 for batch in dataloader_iter:
712 logger.debug(f"[DEBUG] _compute_batches_metrics: {selection_name} batch columns = {batch.schema.names}")
713 metrics_list = list(self.metrics_processors.values())
714 gap_list = list(self.gap_processors.values())
715 batch_features, batch_metrics = self._process_batch(batch, ordered_processors, metrics_list, gap_list)
717 for k, v in batch_metrics.items():
718 if k not in batch_metrics_accumulator:
719 batch_metrics_accumulator[k] = []
720 batch_metrics_accumulator[k].append(v)
722 feature_array_size = self._accumulate_source_features(batch, features_accumulator, feature_array_size)
723 feature_array_size = self._accumulate_generated_features(
724 batch,
725 batch_features,
726 batch_metrics,
727 features_accumulator,
728 feature_array_size,
729 )
730 part_index = self._maybe_flush_features(
731 selection_name,
732 features_accumulator,
733 feature_array_size,
734 part_index,
735 memory_threshold,
736 )
737 if part_index > 0: 737 ↛ 738line 737 didn't jump to line 738 because the condition on line 737 was never true
738 feature_array_size = 0
740 batches_metrics_array = self._concatenate_accumulator(batch_metrics_accumulator)
741 self._write_remaining_features(selection_name, features_accumulator, part_index)
742 self._clear_path_prefixes(ordered_processors)
744 return batches_metrics_array
746 def _parse_memory_string(self, memory_str: str) -> int:
747 """Parse memory string (e.g., "2GB", "500MB") to bytes.
749 Args:
750 memory_str: Memory string to parse.
752 Returns:
753 Memory in bytes.
754 """
755 memory_str = memory_str.strip().upper()
756 if memory_str.endswith("GB"):
757 return int(float(memory_str[:-2]) * 1024 * 1024 * 1024)
758 elif memory_str.endswith("MB"):
759 return int(float(memory_str[:-2]) * 1024 * 1024)
760 elif memory_str.endswith("KB"):
761 return int(float(memory_str[:-2]) * 1024)
762 elif memory_str.endswith("B"):
763 return int(float(memory_str[:-1]))
764 else:
765 # Assume it's in bytes
766 return int(memory_str)
768 def _inject_dataloader_column(self, selection_name: str, features: dict[str, Any]) -> None:
769 """Inject the dataloader column into a features dict.
771 Adds the selection name as a column so the output parquet contains a
772 ``dataloader`` column identifying which dataset each row originates from.
774 Args:
775 selection_name: Name of the current data selection (dataloader name).
776 features: Mutable dict of column_name -> pa.Array to inject into.
777 """
778 if not self.features_output:
779 return
780 if not features:
781 return
783 sample = next(iter(features.values()))
784 features["dataloader"] = pa.array([selection_name] * len(sample))