Coverage for ml_workbench/feature.py: 14%
155 statements
« prev ^ index » next coverage.py v7.11.0, created at 2026-01-06 16:09 +0200
1from __future__ import annotations
3from collections.abc import Iterable, Mapping, Sequence
4from dataclasses import dataclass
5from typing import TYPE_CHECKING, Any
7import pandas as pd
9from .dataset import Dataset
11if TYPE_CHECKING: 11 ↛ 12line 11 didn't jump to line 12 because the condition on line 11 was never true
12 from .config import YamlConfig
15@dataclass(frozen=True)
16class FeatureSpec:
17 name: str
18 dataset: str
19 numerical: list[str]
20 categorical: list[str]
21 column: list[str]
22 description: str | None = None
25class Feature:
26 """Represents a feature set (group) defined in a features YAML.
28 Expected structure under ``features``:
30 features:
31 group_name:
32 description: "..." # optional
33 dataset: dataset_name # required
34 numerical: [col_a, col_b] # optional list of columns
35 categorical: [col_c, col_d] # optional list of columns
37 Also supports "columns" or "column" forms for auto-inferred types (they are equivalent):
39 features:
40 feature_name:
41 dataset: dataset_name
42 columns: [col_a, col_b, col_c]
43 # OR
44 column: [col_a, col_b, col_c]
46 The "columns"/"column" fields are stored as a separate attribute "column" (as a list of column names).
47 Their types (numerical/categorical) will be inferred automatically during experiment setup,
48 rather than being assigned here.
49 """
51 def __init__(self, name: str, config: YamlConfig) -> None:
52 self.name = name
53 self.config = config
55 features_section = self._get_features_section(config)
56 if name not in features_section:
57 raise KeyError(f"Feature '{name}' not found in configuration") # noqa: TRY003
59 spec_raw = features_section[name]
60 if not isinstance(spec_raw, Mapping):
61 raise TypeError(f"Feature '{name}' specification must be a mapping") # noqa: TRY003
63 dataset_name = spec_raw.get("dataset")
64 description = spec_raw.get("description")
66 if not isinstance(dataset_name, str) or not dataset_name:
67 raise ValueError( # noqa: TRY003
68 f"Feature '{name}' is missing required 'dataset' string field"
69 )
71 # Accept both grouped (numerical/categorical), columns, and column schema
72 numerical_cols: list[str] = []
73 categorical_cols: list[str] = []
74 column_cols: list[str] = []
76 if "numerical" in spec_raw or "categorical" in spec_raw:
77 raw_num = spec_raw.get("numerical", [])
78 raw_cat = spec_raw.get("categorical", [])
79 if not isinstance(raw_num, list) or not all(
80 isinstance(x, str) for x in raw_num
81 ):
82 raise ValueError( # noqa: TRY003
83 f"Feature '{name}'.numerical must be a list of strings"
84 )
85 if not isinstance(raw_cat, list) or not all(
86 isinstance(x, str) for x in raw_cat
87 ):
88 raise ValueError( # noqa: TRY003
89 f"Feature '{name}'.categorical must be a list of strings"
90 )
91 numerical_cols = list(raw_num)
92 categorical_cols = list(raw_cat)
93 elif "columns" in spec_raw or "column" in spec_raw:
94 # Treat "columns" and "column" as equivalent
95 column_names = spec_raw.get("columns") or spec_raw.get("column", [])
96 if not isinstance(column_names, list) or not all(
97 isinstance(x, str) for x in column_names
98 ):
99 field_name = "columns" if "columns" in spec_raw else "column"
100 raise ValueError( # noqa: TRY003
101 f"Feature '{name}'.{field_name} must be a list of strings"
102 )
103 # Store in column_cols for type inference later
104 # These will be re-evaluated based on the dataset dtypes
105 column_cols = list(column_names)
106 else:
107 raise ValueError( # noqa: TRY003
108 f"Feature '{name}' must define either 'numerical'/'categorical', 'columns', or 'column'"
109 )
111 self.spec = FeatureSpec(
112 name=name,
113 dataset=dataset_name,
114 numerical=numerical_cols,
115 categorical=categorical_cols,
116 column=column_cols,
117 description=description if isinstance(description, str) else None,
118 )
120 # Validate referenced dataset exists
121 datasets = config.get_data().get("datasets")
122 if not (isinstance(datasets, Mapping) and self.spec.dataset in datasets):
123 raise ValueError( # noqa: TRY003
124 f"Feature '{name}' references unknown dataset '{self.spec.dataset}'"
125 )
127 @staticmethod
128 def _get_features_section(config: YamlConfig) -> Mapping[str, Any]:
129 features = config.get_data().get("features")
130 if not isinstance(features, Mapping):
131 raise KeyError("No 'features' section found in configuration") # noqa: TRY003
132 return features
134 def get_series(self, column: str, *, index: str | None = None) -> pd.Series:
135 """Materialize a single column from this feature set as a pandas Series.
137 The ``column`` must be listed under this feature set's numerical,
138 categorical, or column lists.
139 """
140 if (
141 column not in self.spec.numerical
142 and column not in self.spec.categorical
143 and column not in self.spec.column
144 ):
145 raise KeyError( # noqa: TRY003
146 f"Column '{column}' is not declared in feature set '{self.name}'"
147 )
149 dataset = Dataset(self.spec.dataset, self.config)
150 df = dataset.read_pandas()
151 if index is not None and index in df.columns:
152 df = df.set_index(index)
153 if column not in df.columns:
154 raise KeyError( # noqa: TRY003
155 f"Column '{column}' not found in dataset '{self.spec.dataset}'"
156 )
157 series = df[column].copy()
158 series.name = f"{self.name}.{column}"
159 return series
161 def to_dict(self) -> dict[str, Any]:
162 return {
163 "name": self.spec.name,
164 "dataset": self.spec.dataset,
165 "numerical": list(self.spec.numerical),
166 "categorical": list(self.spec.categorical),
167 "column": list(self.spec.column),
168 "description": self.spec.description,
169 }
171 def get_columns_by_type(self) -> dict[str, list[str]]:
172 return {
173 "numerical": list(self.spec.numerical),
174 "categorical": list(self.spec.categorical),
175 "column": list(self.spec.column),
176 }
178 @classmethod
179 def list_feature_names(cls, config: YamlConfig) -> list[str]:
180 features = cls._get_features_section(config)
181 return list(features.keys())
183 @classmethod
184 def load_all(cls, config: YamlConfig) -> dict[str, Feature]:
185 names = cls.list_feature_names(config)
186 return {name: cls(name, config) for name in names}
188 @classmethod
189 def to_dataframe(
190 cls,
191 config: YamlConfig,
192 *,
193 feature_sets: Iterable[str] | None = None,
194 include_types: Sequence[str] | None = None,
195 index: str | None = None,
196 ) -> pd.DataFrame:
197 """Materialize multiple feature sets into a pandas DataFrame.
199 Parameters
200 ----------
201 feature_sets: optional iterable of feature set names to include. If None, all are used.
202 include_types: optional list among ["numerical", "categorical", "column"] to filter columns.
203 index: optional column name used as index before joining.
204 """
205 selected_sets = list(feature_sets or cls.list_feature_names(config))
206 type_filter = set(include_types or ["numerical", "categorical"])
208 series_list: list[pd.Series] = []
209 for fs_name in selected_sets:
210 fs = cls(fs_name, config)
211 columns: list[str] = []
212 if "numerical" in type_filter:
213 columns.extend(fs.spec.numerical)
214 if "categorical" in type_filter:
215 columns.extend(fs.spec.categorical)
216 if "column" in type_filter:
217 columns.extend(fs.spec.column)
218 for col in columns:
219 series_list.append(fs.get_series(col, index=index))
221 if not series_list:
222 return pd.DataFrame()
224 df = series_list[0].to_frame()
225 for s in series_list[1:]:
226 df = df.join(s, how="outer")
227 return df
229 @classmethod
230 def verify_config(cls, config: YamlConfig) -> None:
231 """Validate that all features are well-formed and reference known datasets.
233 Checks
234 ------
235 - features section exists if referenced and is a mapping
236 - each feature has a string "dataset" and the dataset exists in config.datasets
237 - either numerical/categorical lists (lists of strings) are present, or a legacy
238 single "column" (string) is provided
239 - no column appears in both numerical and categorical within the same feature
240 - at least one column is declared per feature
241 """
243 features = config.get_data().get("features")
244 if features is None:
245 return # Nothing to verify
246 if not isinstance(features, Mapping):
247 raise TypeError("'features' section must be a mapping") # noqa: TRY003
249 datasets = config.get_data().get("datasets")
250 if not isinstance(datasets, Mapping):
251 raise TypeError("No 'datasets' section found while validating features") # noqa: TRY003
253 for feature_name, raw in features.items():
254 if not isinstance(raw, Mapping):
255 raise TypeError( # noqa: TRY003
256 f"Feature '{feature_name}' specification must be a mapping"
257 )
259 dataset_name = raw.get("dataset")
260 if not isinstance(dataset_name, str) or not dataset_name:
261 raise ValueError( # noqa: TRY003
262 f"Feature '{feature_name}' is missing required 'dataset' string field"
263 )
264 if dataset_name not in datasets:
265 raise ValueError( # noqa: TRY003
266 f"Feature '{feature_name}' references unknown dataset '{dataset_name}'"
267 )
269 has_grouped = ("numerical" in raw) or ("categorical" in raw)
270 has_columns_or_column = "columns" in raw or "column" in raw
271 declared_columns: list[str] = []
272 if has_grouped:
273 for key in ("numerical", "categorical"):
274 val = raw.get(key, [])
275 if val is None:
276 val = []
277 if not isinstance(val, list) or not all(
278 isinstance(x, str) for x in val
279 ):
280 raise ValueError( # noqa: TRY003
281 f"Feature '{feature_name}'.{key} must be a list of strings"
282 )
283 num_list = list(raw.get("numerical", []) or [])
284 cat_list = list(raw.get("categorical", []) or [])
285 # Check duplicates across types
286 overlap = set(num_list).intersection(cat_list)
287 if overlap:
288 dup = ", ".join(sorted(overlap))
289 raise ValueError( # noqa: TRY003
290 f"Feature '{feature_name}' declares columns in both numerical and categorical: {dup}"
291 )
292 declared_columns.extend(num_list)
293 declared_columns.extend(cat_list)
294 # '__all__' is only allowed in the 'columns'/'column' form and must be alone
295 if "__all__" in num_list or "__all__" in cat_list:
296 raise ValueError( # noqa: TRY003
297 f"Feature '{feature_name}' cannot use '__all__' inside 'numerical' or 'categorical'; use 'columns: [__all__]' or 'column: [__all__]'"
298 )
299 elif has_columns_or_column:
300 # Treat "columns" and "column" as equivalent
301 column_names = raw.get("columns") or raw.get("column", [])
302 field_name = "columns" if "columns" in raw else "column"
303 if not isinstance(column_names, list) or not all(
304 isinstance(x, str) for x in column_names
305 ):
306 raise ValueError( # noqa: TRY003
307 f"Feature '{feature_name}'.{field_name} must be a list of strings"
308 )
309 # If '__all__' is used, it must be the only entry
310 if "__all__" in column_names and len(column_names) != 1:
311 raise ValueError( # noqa: TRY003
312 f"Feature '{feature_name}' uses '__all__' alongside other columns; '__all__' must be alone"
313 )
314 declared_columns.extend(column_names)
315 else:
316 raise ValueError( # noqa: TRY003
317 f"Feature '{feature_name}' must define either 'numerical'/'categorical', 'columns', or 'column'"
318 )
319 if not declared_columns:
320 raise ValueError(f"Feature '{feature_name}' declares no columns") # noqa: TRY003
321 # '__all__' is allowed only when it is the single entry in the 'columns' or 'column' list.
322 # The grouped form is already rejected above.
325__all__ = ["Feature", "FeatureSpec"]