Coverage for packages/dqm-ml-job/src/dqm_ml_job/dataloaders/pandas.py: 92%
105 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-07-21 08:27 +0000
1"""Pandas data loader for reading CSV files.
3This module contains the PandasDataLoader and PandasDataSelection classes
4for loading and iterating over CSV file data using Pandas.
5"""
7import fnmatch
8import logging
9from typing import Any
11from dqm_ml_core.utils.matching import has_pattern, resolve_include_exclude, resolve_patterns
12import pandas as pd
13import pyarrow as pa
15# COMPATIBILITY : from typing import Any, override # When support of 3.10 and 3.11 will be removed
16from typing_extensions import override
18from dqm_ml_job.dataloaders.proto import DataSelection
20logger = logging.getLogger(__name__)
23def _match_wildcard(series: pd.Series, patterns: list[str]) -> pd.Series:
24 """Return boolean mask for rows matching any wildcard pattern.
26 Args:
27 series: Pandas Series to filter.
28 patterns: List of patterns with fnmatch wildcards (* and ?).
30 Returns:
31 Boolean Series indicating matches.
32 """
33 mask = pd.Series(False, index=series.index)
34 for pattern in patterns:
35 mask = mask | series.astype(str).apply(lambda x, p=pattern: fnmatch.fnmatch(x, p))
36 return mask
39_PANDAS_TYPE_MAP: dict[str, Any] = {
40 "int32": "int32",
41 "int64": "int64",
42 "float32": "float32",
43 "float64": "float64",
44 "bool": "bool",
45 "str": "string",
46 "categorical": "category",
47}
50def _apply_pandas_transforms(df: pd.DataFrame, transforms: list[dict[str, Any]]) -> None:
51 """Apply column transforms to a pandas DataFrame in-place."""
52 for t in transforms:
53 col = t["column"]
54 if col not in df.columns:
55 continue
56 pandas_type = _PANDAS_TYPE_MAP[t["to_type"]]
57 if t.get("in_place", False):
58 df[col] = df[col].astype(pandas_type)
59 else:
60 new_name = f"{col}_{t['to_type']}"
61 df[new_name] = df[col].astype(pandas_type)
64class PandasDataSelection(DataSelection):
65 """A selection of data from a CSV file loaded via Pandas.
67 This class represents data loaded from a CSV file and provides
68 an iterator over PyArrow RecordBatches.
70 Attributes:
71 name: Name identifier for this selection.
72 path: Path to the CSV file.
73 data: The loaded pandas DataFrame.
74 sample_path: List of sample path configs describing column path prefixes.
75 transforms: List of transform configs (column cast operations).
76 """
78 def __init__(
79 self,
80 name: str,
81 path: str,
82 sample_path: list[dict[str, Any]] | None = None,
83 transforms: list[dict[str, Any]] | None = None,
84 filters_dict: dict[str, Any] | None = None,
85 ):
86 """Initialize a Pandas data selection.
88 Args:
89 name: Name identifier for this selection.
90 path: Path to the CSV file.
91 sample_path: List of sample path configs describing column path prefixes.
92 transforms: List of transform configs (column cast operations).
93 filters_dict: Column-value pairs to filter rows by.
94 """
95 self.name = name
96 self.path = path
97 self.sample_path = sample_path or []
98 self.transforms = transforms or []
99 self.filters_dict = filters_dict or {}
100 self.data: pd.DataFrame | None = None
102 @override
103 def bootstrap(self, columns_list: list[str] | None = None) -> None:
104 """Load the CSV file into memory as a pandas DataFrame.
106 Args:
107 columns_list: Unused, kept for API compatibility.
108 """
109 from dqm_ml_job.dataloaders.filters import build_filter_condition
111 data = pd.read_csv(self.path, sep=",")
112 assert isinstance(data, pd.DataFrame)
113 self.data = data
114 for col, val in self.filters_dict.items():
115 condition = build_filter_condition(
116 col,
117 val,
118 wildcard_fn=lambda c, vals: _match_wildcard(data[c], vals),
119 isin_fn=lambda c, vals: data[c].isin(vals),
120 equal_fn=lambda c, v: data[c] == v,
121 )
122 self.data = self.data[condition.reindex(self.data.index, fill_value=False)]
124 def __len__(self) -> int:
125 return len(self.data) if self.data is not None else 0
127 @override
128 def get_nb_batches(self) -> int:
129 """Return the estimated number of batches (always 1 for CSV).
131 Returns:
132 1 if data is loaded, 0 otherwise.
133 """
134 return 1 if self.data is not None else 0
136 @override
137 def __iter__(self) -> Any:
138 if self.data is not None: 138 ↛ exitline 138 didn't return from function '__iter__' because the condition on line 138 was always true
139 df = self.data.copy() if self.transforms else self.data
140 _apply_pandas_transforms(df, self.transforms)
141 yield pa.RecordBatch.from_pandas(df)
143 @override
144 def __repr__(self) -> str:
145 return f"PandasSelection(name='{self.name}', path='{self.path}')"
148class PandasDataLoader:
149 """Data loader for CSV files using Pandas.
151 This loader reads CSV files and provides DataSelections for
152 processing by the DQM pipeline.
154 Attributes:
155 type: The loader type identifier ("csv").
156 """
158 type: str = "csv"
160 def __init__(self, name: str, config: dict[str, Any] | None = None):
161 """Initialize the Pandas data loader.
163 Args:
164 name: Unique name for this loader instance.
165 config: Configuration dictionary containing:
166 - path: Path to CSV file (required)
168 Raises:
169 ValueError: If required config keys are missing.
170 """
171 if config is None: 171 ↛ 172line 171 didn't jump to line 172 because the condition on line 171 was never true
172 config = {}
173 self.name = name
174 self.path = config["path"]
175 # Use SplitConfig model fields instead of hardcoded keys
176 from dqm_ml_core.models.dataloaders import SplitConfig
178 split = config.get("split")
179 self.split = SplitConfig.model_validate(split) if split else None
180 self.split_by = self.split.by if self.split else None
181 self.split_values = self.split.values if self.split else None
182 filters = config.get("filters")
183 # transform the list of dict into a dict
184 self.filters_dict = {}
185 if filters is not None:
186 for item in filters:
187 column = item["column"]
188 self.filters_dict[column] = item["values"]
189 self.id_column = config.get("id_column")
190 self.sample_path = config.get("sample_path", [])
191 self.transforms = config.get("transform", [])
193 # Storage filesystem configuration - only for S3 paths, not local paths
194 self.filesystem = None
195 storage_cfg = config.get("storage")
196 if storage_cfg: 196 ↛ 198line 196 didn't jump to line 198 because the condition on line 196 was never true
197 # Use StorageConfig model to validate and access fields
198 from dqm_ml_core.models.global_ import StorageConfig
200 storage_config = StorageConfig.model_validate(storage_cfg)
202 if storage_config.type == "s3":
203 from dqm_ml_job.utils.s3 import get_s3_filesystem
205 self.filesystem = get_s3_filesystem(storage_config)
207 def get_selections(self) -> list[DataSelection]:
208 """Create one or more PandasDataSelection instances based on split config.
210 If split is configured, returns one selection per split value.
211 Otherwise returns a single selection for the entire CSV file.
213 Returns:
214 A list of DataSelection instances.
215 """
216 if not self.split_by:
217 return [
218 PandasDataSelection(
219 name=self.name,
220 path=self.path,
221 sample_path=self.sample_path,
222 transforms=self.transforms,
223 filters_dict=self.filters_dict,
224 )
225 ]
227 # Determine split values
228 values = self.split_values
229 if values is None:
230 # Auto-discover unique values from the CSV
231 df = pd.read_csv(self.path, sep=",", usecols=[self.split_by])
232 values = [str(v) for v in df[self.split_by].unique() if v is not None]
233 else:
234 # Expand wildcard patterns in values against available data
235 if any(has_pattern(v) for v in values):
236 df = pd.read_csv(self.path, sep=",", usecols=[self.split_by])
237 available = [str(v) for v in df[self.split_by].unique() if v is not None]
238 values = resolve_patterns(values, available)
240 # Apply split.exclude (including wildcard patterns)
241 if self.split and self.split.exclude:
242 values = resolve_include_exclude(None, self.split.exclude, values)
244 # Create one selection per value
245 selections: list[DataSelection] = []
246 for val in values:
247 selection_name = f"{self.name}_{val}"
248 merged_filters = (self.filters_dict or {}).copy()
249 merged_filters[self.split_by] = val
250 selections.append(
251 PandasDataSelection(
252 name=selection_name,
253 path=self.path,
254 sample_path=self.sample_path,
255 transforms=self.transforms,
256 filters_dict=merged_filters,
257 )
258 )
259 return selections