Coverage for ml_workbench/dataset.py: 8%

1"""

2Module for working with datasets.

4Dataset class is responsible to read dataset in accordance with it type and format.

5It also provides basic statistics about the dataset.

7For combined datasets, it is responsible to read the datasets and merge them in accordance with the merge specification.

9Dataset class is initialized with a dataset name and a dataset specification.

10Primitive dataset specification is a dictionary with the following keys:

11- name: str

12- description: str

13- path: str

14- format: str

15- type: str

17Combined dataset specification is a dictionary with the following keys:

18- name: str

19- description: str

20- merge_specs: dict

21 - dataset_name: dict

22 - right_on: str

23 - left_on: str

24 - how: str

26Dataset class has the following methods:

27- read: read the dataset in accordance with it type and format

28- get_statistics: get basic statistics about the dataset

29- get_schema: get the schema of the dataset

30- get_columns: get the columns of the dataset

31- get_rows: get the rows of the dataset

32- get_head: get the head of the dataset

33"""

35from __future__ import annotations

37from pathlib import Path

38from typing import TYPE_CHECKING, Any

40import pandas as pd

42# Constants

43DELTA_TABLE_PARTS_COUNT = 2 # catalog.schema.table format has 2 dots

45if TYPE_CHECKING: 45 ↛ 46line 45 didn't jump to line 46 because the condition on line 45 was never true

46 from .config import YamlConfig

49def _is_databricks_table_name(path: str) -> bool:

50 """Return True if ``path`` looks like a Databricks table identifier.

52 Heuristic: exactly two dots separating catalog.schema.table and no path separators

53 or URL schemes. This keeps the check permissive and practical.

54 """

56 if "/" in path or "\\" in path:

57 return False

58 # Exclude common URL schemes

59 if "://" in path:

60 return False

61 if path.count(".") != DELTA_TABLE_PARTS_COUNT:

62 return False

63 parts = path.split(".")

64 return all(part.strip() for part in parts)

67def _infer_dataset_format_from_path(path: str) -> str | None:

68 """Infer dataset type from a path-like string.

70 - If the value looks like a Databricks table name (catalog.schema.table), return "delta".

71 - Otherwise infer by file extension (case-insensitive): csv, txt, parquet, json.

72 - Returns None if no inference is possible.

73 """

75 if _is_databricks_table_name(path):

76 return "delta"

78 # Extract filename segment and extension without being confused by schemes

79 last_segment = path.rsplit("/", 1)[-1]

80 if "." not in last_segment:

81 return None

82 ext = last_segment.rsplit(".", 1)[-1].lower()

83 if ext in {"csv", "txt", "parquet", "json"}:

84 return ext

85 return None

88class Dataset:

89 """Dataset abstraction for reading data from various sources.

91 Supports:

92 - Local files (CSV, TXT, Parquet, JSON)

93 - S3 paths (CSV, Parquet, JSON)

94 - Databricks Delta tables and /Volumes/ paths

95 - Combined datasets via merge specifications

97 Always returns pandas DataFrame regardless of source.

98 """

100 def __init__(self, name: str, config: YamlConfig) -> None:

101 """Initialize a Dataset.

102

103 Parameters

104 ----------

105 name : str

106 Dataset name

107 config : YamlConfig

108 Configuration object containing dataset specifications

109 """

110 self.name = name

111 self.config = config

112

113 # Get dataset spec from config

114 spec = config.get_dataset_config(name)

115

116 self.path = spec.get("path")

117 self.format = spec.get("format")

118 self.type = spec.get("type", "local")

119 self.description = spec.get("description")

120 self.is_combined = "merge_specs" in spec

121 self._merge_specs = spec.get("merge_specs") if self.is_combined else None

122 self._df: pd.DataFrame | None = None

123

124 def read_pandas(self) -> pd.DataFrame:

125 """Read the dataset and return a pandas DataFrame.

126

127 For combined datasets, reads all participating datasets and merges them

128 according to merge specifications.

129

130 Returns

131 -------

132 pd.DataFrame

133 The loaded dataset

134

135 Raises

136 ------

137 ValueError

138 If path or format is missing or unsupported, or if config is missing

139 for combined datasets

140 RuntimeError

141 If reading fails

142 """

143 if self._df is not None:

144 return self._df

145

146 # Handle combined datasets

147 if self.is_combined:

148 self._df = self._read_combined()

149 return self._df

150

151 # Handle primitive datasets

152 if not self.path:

153 raise ValueError(f"Dataset '{self.name}' has no path specified") # noqa: TRY003

154

155 if self.type == "databricks":

156 self._df = self._read_databricks()

157 elif self.type == "s3":

158 self._df = self._read_s3()

159 elif self.type == "local":

160 self._df = self._read_local()

161 else:

162 raise ValueError(f"Unsupported dataset type: {self.type}") # noqa: TRY003

163

164 return self._df

165

166 def _read_combined(self) -> pd.DataFrame:

167 """Read and merge multiple datasets according to merge specifications.

168

169 Returns

170 -------

171 pd.DataFrame

172 Merged dataset

173

174 Raises

175 ------

176 ValueError

177 If merge specs are invalid

178 """

179 if not self._merge_specs:

180 raise ValueError(f"Combined dataset '{self.name}' has no merge_specs") # noqa: TRY003

181

182 # Start with None; first dataset becomes the base

183 result_df: pd.DataFrame | None = None

184

185 # Convert to list to track position

186 merge_items = list(self._merge_specs.items())

187

188 # Iterate through merge specs in order

189 for idx, (dataset_name, merge_spec) in enumerate(merge_items):

190 # Get the dataset configuration

191 try:

192 self.config.get_dataset_config(dataset_name)

193 except KeyError as exc:

194 raise ValueError( # noqa: TRY003

195 f"Dataset '{dataset_name}' referenced in merge_specs not found"

196 ) from exc

197

198 # Create and read the dataset

199 ds = Dataset(dataset_name, self.config)

200 df = ds.read_pandas()

201

202 # First dataset becomes the base (no merge needed)

203 if result_df is None:

204 result_df = df

205 continue

206

207 # Merge with the accumulated result

208 # right_on comes from current dataset's merge_spec

209 # left_on and how come from previous dataset's merge_spec

210 prev_dataset_name, prev_merge_spec = merge_items[idx - 1]

211 right_on = merge_spec.get("right_on")

212 left_on = prev_merge_spec.get("left_on")

213 how = prev_merge_spec.get("how", "inner")

214

215 # Perform the merge

216 if right_on and left_on:

217 # Merge on specified columns

218 result_df = result_df.merge(

219 df,

220 left_on=left_on,

221 right_on=right_on,

222 how=how,

223 )

224 elif left_on:

225 # Merge on index (right) and column (left)

226 result_df = result_df.merge(

227 df,

228 right_index=True,

229 left_on=left_on,

230 how=how,

231 )

232 elif right_on:

233 # Last dataset: merge on column (right) and index (left)

234 result_df = result_df.merge(

235 df,

236 right_on=right_on,

237 left_index=True,

238 how=how,

239 )

240 else:

241 # Both use index

242 result_df = result_df.merge(

243 df,

244 right_index=True,

245 left_index=True,

246 how=how,

247 )

248

249 if result_df is None:

250 raise ValueError(f"Combined dataset '{self.name}' produced no data") # noqa: TRY003

251

252 return result_df

253

254 def _read_databricks(self) -> pd.DataFrame:

255 """Read from Databricks (Delta tables or /Volumes/ paths)."""

256 try:

257 from pyspark.sql import ( # noqa: PLC0415

258 SparkSession, # type: ignore[import-not-found] # noqa: PLC0415

259 )

260

261 spark = SparkSession.builder.getOrCreate()

262

263 # Check if it's a table name (catalog.schema.table)

264 if self.path.count(".") == DELTA_TABLE_PARTS_COUNT and "/" not in self.path: # type: ignore[union-attr]

265 # Delta table

266 spark_df = spark.table(self.path)

267 else:

268 # File path under /Volumes/ or similar

269 fmt = self.format or "delta"

270 if fmt == "delta":

271 spark_df = spark.read.format("delta").load(self.path)

272 elif fmt == "parquet":

273 # special handling for parquet to fix INT64 (TIMESTAMP_MICROS) error in Databricks

274 pdf = pd.read_parquet(self.path)

275 return pdf

276 elif fmt == "csv":

277 spark_df = spark.read.csv(self.path, header=True, inferSchema=True)

278 elif fmt == "json":

279 spark_df = spark.read.json(self.path)

280 elif fmt == "txt":

281 spark_df = spark.read.text(self.path)

282 else:

283 raise ValueError(f"Unsupported format for Databricks: {fmt}") # noqa: TRY003

284

285 return spark_df.toPandas()

286 except ImportError as exc:

287 raise RuntimeError( # noqa: TRY003

288 "PySpark is required for Databricks datasets. "

289 "Install with: pip install pyspark or databricks-connect"

290 ) from exc

291

292 def _read_s3(self) -> pd.DataFrame:

293 """Read from S3 paths by fetching into an IO object via boto3."""

294 import io # noqa: PLC0415

295 import re # noqa: PLC0415

296

297 import boto3 # noqa: PLC0415

298

299 fmt = (self.format or "").lower()

300

301 # Extract bucket and key from the S3 path

302 # supports s3://bucket-name/key/to/object

303 match = re.match(r"s3://([^/]+)/(.+)", self.path)

304 if not match:

305 raise ValueError(f"Invalid S3 path: {self.path}") # noqa: TRY003

306 bucket, key = match.groups()

307

308 s3 = boto3.client("s3")

309 obj = s3.get_object(Bucket=bucket, Key=key)

310 file_obj = io.BytesIO(obj["Body"].read())

311

312 if fmt == "csv":

313 return pd.read_csv(file_obj)

314 if fmt == "parquet":

315 return pd.read_parquet(file_obj)

316 if fmt == "json":

317 return pd.read_json(file_obj)

318 raise ValueError(f"Unsupported format for S3: {fmt}") # noqa: TRY003

319

320 def _read_local(self) -> pd.DataFrame:

321 """Read from local filesystem."""

322 path = Path(self.path) # type: ignore[arg-type]

323

324 if not path.exists():

325 raise FileNotFoundError(f"Local file not found: {path}") # noqa: TRY003

326

327 fmt = (self.format or "").lower()

328

329 if fmt == "csv":

330 return pd.read_csv(path)

331 if fmt == "txt":

332 # TXT files often are CSV-like with different delimiters

333 return pd.read_csv(path, sep="\t")

334 if fmt == "parquet":

335 return pd.read_parquet(path)

336 if fmt == "json":

337 return pd.read_json(path)

338 raise ValueError(f"Unsupported format for local files: {fmt}") # noqa: TRY003

339

340 def get_statistics(self) -> dict[str, Any]:

341 """Get basic statistics about the dataset.

342

343 Returns

344 -------

345 Dict[str, Any]

346 Statistics including num_rows, num_columns, column_names, dtypes

347 """

348 df = self.read_pandas()

349 return {

350 "num_rows": len(df),

351 "num_columns": len(df.columns),

352 "column_names": list(df.columns),

353 "dtypes": {col: str(dtype) for col, dtype in df.dtypes.items()},

354 "memory_usage_bytes": int(df.memory_usage(deep=True).sum()),

355 }

356

357 def get_schema(self) -> dict[str, str]:

358 """Get the schema (column names and types).

359

360 Returns

361 -------

362 Dict[str, str]

363 Mapping of column name to dtype string

364 """

365 df = self.read_pandas()

366 return {col: str(dtype) for col, dtype in df.dtypes.items()}

367

368 def get_columns(self) -> list[str]:

369 """Get column names.

370

371 Returns

372 -------

373 list[str]

374 List of column names

375 """

376 df = self.read_pandas()

377 return list(df.columns)

378

379 def get_rows(self) -> int:

380 """Get number of rows.

381

382 Returns

383 -------

384 int

385 Number of rows

386 """

387 df = self.read_pandas()

388 return len(df)

389

390 def get_head(self, n: int = 5) -> pd.DataFrame:

391 """Get first n rows.

392

393 Parameters

394 ----------

395 n : int, default 5

396 Number of rows to return

397

398 Returns

399 -------

400 pd.DataFrame

401 First n rows

402 """

403 df = self.read_pandas()

404 return df.head(n)

405

406 def __repr__(self) -> str:

407 return f"Dataset(name={self.name!r}, type={self.type}, format={self.format}, path={self.path!r})"

408

409 @classmethod

410 def _impute_dataset_types(cls, config: YamlConfig) -> None:

411 """Impute dataset format and source type for entries in ``datasets``.

412

413 - "format": csv, txt, parquet, json, delta (inferred from path or table name)

414 - "type": databricks, s3, local (inferred from path and/or format)

415

416 Datasets that represent composite/merged inputs (e.g. have ``merge_specs``)

417 are left unchanged.

418 """

419

420 datasets = config.get_data().get("datasets")

421 if not isinstance(datasets, dict):

422 return

423

424 for _dataset_name, dataset_spec in datasets.items():

425 if not isinstance(dataset_spec, dict):

426 continue

427

428 # Skip combined datasets defined via merge specs

429 if "merge_specs" in dataset_spec:

430 continue

431

432 dataset_path = dataset_spec.get("path")

433 if not isinstance(dataset_path, str):

434 continue

435

436 # 1) Impute format if missing

437 has_format = "format" in dataset_spec and dataset_spec.get(

438 "format"

439 ) not in (None, "")

440 if not has_format:

441 inferred_fmt = _infer_dataset_format_from_path(dataset_path)

442 if inferred_fmt:

443 dataset_spec["format"] = inferred_fmt

444

445 # 2) Impute source type if missing

446 has_type = "type" in dataset_spec and dataset_spec.get("type") not in (

447 None,

448 "",

449 )

450 if not has_type:

451 fmt = str(dataset_spec.get("format") or "").lower()

452 if _is_databricks_table_name(dataset_path) or dataset_path.startswith(

453 "/Volumes/"

454 ):

455 src_type = "databricks"

456 elif dataset_path.startswith("s3://"):

457 src_type = "s3"

458 elif fmt == "delta":

459 src_type = "databricks"

460 else:

461 src_type = "local"

462 dataset_spec["type"] = src_type

463

464 @classmethod

465 def verify_config(cls, config: YamlConfig) -> None:

466 """Verify dataset configuration integrity.

467

468 Steps:

469 1. Impute dataset formats and types where missing.

470 2. Validate that combined datasets reference only datasets defined in this config.

471

472 Raises

473 ------

474 ValueError

475 If a combined dataset references a dataset not present in the configuration.

476 """

477

478 # Step 1: impute

479 cls._impute_dataset_types(config)

480

481 # Step 2: validate combined dataset references

482 datasets = config.get_data().get("datasets")

483 if not isinstance(datasets, dict):

484 return

485

486 for combined_name, combined_spec in datasets.items():

487 if not isinstance(combined_spec, dict):

488 continue

489 merge_specs = combined_spec.get("merge_specs")

490 if not isinstance(merge_specs, dict):

491 continue

492

493 for referenced_name in merge_specs:

494 if referenced_name not in datasets:

495 raise ValueError( # noqa: TRY003

496 f"Combined dataset '{combined_name}' references unknown dataset '{referenced_name}'"

497 )