From 8d920650895d12709b4bf747d003721fccb62628 Mon Sep 17 00:00:00 2001 From: Miguel Twahirwa Date: Thu, 23 Oct 2025 23:30:50 -0400 Subject: [PATCH] Add pandas data loader with schema inference --- .gitignore | 2 + macrobase_py/__init__.py | 5 + macrobase_py/data/__init__.py | 5 + macrobase_py/data/loader.py | 299 ++++++++++++++++++++++++++++++ macrobase_py/tests/conftest.py | 6 + macrobase_py/tests/test_loader.py | 105 +++++++++++ 6 files changed, 422 insertions(+) create mode 100644 macrobase_py/__init__.py create mode 100644 macrobase_py/data/__init__.py create mode 100644 macrobase_py/data/loader.py create mode 100644 macrobase_py/tests/conftest.py create mode 100644 macrobase_py/tests/test_loader.py diff --git a/.gitignore b/.gitignore index 0939a5f86..054af1bcb 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,6 @@ data/ +!macrobase_py/data/ +!macrobase_py/data/** conf/ *.class .classpath diff --git a/macrobase_py/__init__.py b/macrobase_py/__init__.py new file mode 100644 index 000000000..4d88e3181 --- /dev/null +++ b/macrobase_py/__init__.py @@ -0,0 +1,5 @@ +"""MacroBase Python utilities.""" + +from .data.loader import ColumnMetadata, DataLoader + +__all__ = ["ColumnMetadata", "DataLoader"] diff --git a/macrobase_py/data/__init__.py b/macrobase_py/data/__init__.py new file mode 100644 index 000000000..ce9f90f13 --- /dev/null +++ b/macrobase_py/data/__init__.py @@ -0,0 +1,5 @@ +"""Data utilities for MacroBase Python.""" + +from .loader import ColumnMetadata, DataLoader + +__all__ = ["ColumnMetadata", "DataLoader"] diff --git a/macrobase_py/data/loader.py b/macrobase_py/data/loader.py new file mode 100644 index 000000000..4e177883d --- /dev/null +++ b/macrobase_py/data/loader.py @@ -0,0 +1,299 @@ +"""Data loading utilities for MacroBase Python integration.""" +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, Iterable, Iterator, Mapping, Optional + +import pandas as pd +from pandas.api import types as pdt + + +@dataclass(frozen=True) +class ColumnMetadata: + """Metadata describing a single column in a dataset.""" + + name: str + pandas_dtype: str + logical_type: str + nullable: bool + null_count: int + null_fraction: float + unique_values: Optional[int] = None + categorical: bool = False + + +class DataLoader: + """Utility class for reading tabular data with schema inference. + + Parameters + ---------- + source : str + Path to the CSV or Excel file to load. + file_type : Optional[str] + Explicit file type (``"csv"`` or ``"excel"``). If not supplied the + loader will infer the type from the file extension. + chunk_size : Optional[int] + Number of rows per chunk when iterating over data. Chunked loading is + only supported for CSV inputs. + read_options : Optional[Mapping[str, object]] + Additional keyword arguments forwarded to ``pandas.read_csv`` or + ``pandas.read_excel``. + schema_overrides : Optional[Mapping[str, Mapping[str, object]]] + Per-column overrides for the inferred schema. Keys may include + ``logical_type``, ``pandas_dtype``, ``nullable`` and ``categorical``. + sample_rows : int + Number of rows to sample during schema inference. + max_categorical_cardinality : int + Upper bound on the number of unique values for an object column to be + considered categorical. + """ + + SUPPORTED_FILE_TYPES = {"csv", "excel"} + + def __init__( + self, + source: str, + file_type: Optional[str] = None, + *, + chunk_size: Optional[int] = None, + read_options: Optional[Mapping[str, object]] = None, + schema_overrides: Optional[Mapping[str, Mapping[str, object]]] = None, + sample_rows: int = 500, + max_categorical_cardinality: int = 50, + ) -> None: + self.source = Path(source) + self.file_type = self._normalize_file_type(file_type) + self.chunk_size = chunk_size + self._base_read_options = dict(read_options or {}) + self.schema_overrides: Mapping[str, Mapping[str, object]] = ( + schema_overrides or {} + ) + self.sample_rows = sample_rows + self.max_categorical_cardinality = max_categorical_cardinality + + if self.chunk_size is not None and self.file_type != "csv": + raise ValueError("Chunked loading is only supported for CSV files") + + self._schema: Dict[str, ColumnMetadata] = {} + self._infer_schema() + + @property + def schema(self) -> Mapping[str, ColumnMetadata]: + """Return the inferred schema metadata.""" + + return dict(self._schema) + + # ------------------------------------------------------------------ + # Public API + def load(self, limit: Optional[int] = None) -> pd.DataFrame: + """Load data into a typed :class:`pandas.DataFrame`. + + Parameters + ---------- + limit : Optional[int] + Maximum number of rows to read from the source. ``None`` reads the + entire dataset. + """ + + df = self._read_full(limit=limit) + return self._apply_schema(df) + + def iter_batches(self, chunk_size: Optional[int] = None) -> Iterator[pd.DataFrame]: + """Yield dataframes in chunks using the inferred schema.""" + + actual_chunk_size = chunk_size or self.chunk_size + if actual_chunk_size is None: + yield self.load() + return + + for chunk in self._read_in_chunks(actual_chunk_size): + yield self._apply_schema(chunk) + + # ------------------------------------------------------------------ + # Schema inference + def _infer_schema(self) -> None: + sample = self._load_sample() + if sample is None or sample.empty: + self._schema = {} + return + + inferred: Dict[str, ColumnMetadata] = {} + for column in sample.columns: + series = sample[column] + inferred[column] = self._infer_column(series) + + self._schema = inferred + + def _infer_column(self, series: pd.Series) -> ColumnMetadata: + logical_type = self._detect_logical_type(series) + pandas_dtype = str(series.dtype) + null_count = int(series.isna().sum()) + total = len(series) + null_fraction = float(null_count / total) if total > 0 else 0.0 + unique_values: Optional[int] = None + categorical = logical_type == "categorical" + + try: + unique_values = int(series.nunique(dropna=True)) + except TypeError: + unique_values = None + + overrides = self.schema_overrides.get(series.name, {}) + logical_type = overrides.get("logical_type", logical_type) + if logical_type == "boolean" and "pandas_dtype" not in overrides: + pandas_dtype = "boolean" + pandas_dtype = overrides.get("pandas_dtype", pandas_dtype) + categorical = overrides.get("categorical", categorical) + nullable = overrides.get("nullable", null_count > 0) + + return ColumnMetadata( + name=series.name, + pandas_dtype=pandas_dtype, + logical_type=logical_type, + nullable=nullable, + null_count=null_count, + null_fraction=null_fraction, + unique_values=unique_values, + categorical=categorical, + ) + + def _detect_logical_type(self, series: pd.Series) -> str: + if pdt.is_bool_dtype(series): + return "boolean" + if pdt.is_datetime64_any_dtype(series): + return "datetime" + if pdt.is_numeric_dtype(series): + return "numeric" + if isinstance(series.dtype, pdt.CategoricalDtype): + return "categorical" + if pdt.is_object_dtype(series): + non_null = series.dropna() + if not non_null.empty: + normalized = non_null.map(self._normalize_boolean_like) + if normalized.notna().all(): + return "boolean" + + datetime_coerced = pd.to_datetime(non_null, errors="coerce") + valid_datetime_ratio = float( + datetime_coerced.notna().sum() / len(non_null) + ) + if valid_datetime_ratio > 0.9: + return "datetime" + + numeric_coerced = pd.to_numeric(non_null, errors="coerce") + valid_numeric_ratio = float( + numeric_coerced.notna().sum() / len(non_null) + ) + if valid_numeric_ratio > 0.9: + return "numeric" + unique_count = series.nunique(dropna=True) + if unique_count <= self.max_categorical_cardinality: + return "categorical" + return "text" + return "text" + + # ------------------------------------------------------------------ + # IO utilities + def _normalize_file_type(self, file_type: Optional[str]) -> str: + if file_type is not None: + normalized = file_type.lower() + if normalized not in self.SUPPORTED_FILE_TYPES: + raise ValueError(f"Unsupported file type: {file_type}") + return normalized + + suffix = self.source.suffix.lower() + if suffix == ".csv": + return "csv" + if suffix in {".xls", ".xlsx"}: + return "excel" + raise ValueError( + "Could not infer file type from extension; please specify file_type" + ) + + def _load_sample(self) -> Optional[pd.DataFrame]: + options = self._reader_options() + options.pop("chunksize", None) + options.pop("iterator", None) + if self.file_type == "csv": + return pd.read_csv(self.source, nrows=self.sample_rows, **options) + return pd.read_excel(self.source, nrows=self.sample_rows, **options) + + def _reader_options(self) -> Dict[str, object]: + return dict(self._base_read_options) + + def _read_full(self, limit: Optional[int]) -> pd.DataFrame: + options = self._reader_options() + if self.file_type == "csv": + if limit is not None: + options["nrows"] = limit + return pd.read_csv(self.source, **options) + + if limit is not None: + options["nrows"] = limit + return pd.read_excel(self.source, **options) + + def _read_in_chunks(self, chunk_size: int) -> Iterable[pd.DataFrame]: + options = self._reader_options() + options["chunksize"] = chunk_size + if self.file_type == "csv": + reader = pd.read_csv(self.source, **options) + return reader + raise ValueError("Chunked iteration is only implemented for CSV files") + + def _apply_schema(self, df: pd.DataFrame) -> pd.DataFrame: + if not self._schema: + return df + + typed_df = df.copy() + for column, metadata in self._schema.items(): + if column not in typed_df.columns: + continue + + series = typed_df[column] + desired_logical = metadata.logical_type + desired_dtype = metadata.pandas_dtype + + override = self.schema_overrides.get(column, {}) + if "pandas_dtype" in override: + desired_dtype = override["pandas_dtype"] + + try: + if desired_logical == "numeric": + typed_df[column] = pd.to_numeric(series, errors="coerce") + elif desired_logical == "datetime": + typed_df[column] = pd.to_datetime(series, errors="coerce") + elif desired_logical == "boolean": + typed_df[column] = self._coerce_boolean(series) + elif metadata.categorical: + typed_df[column] = series.astype("category") + elif desired_dtype not in (str(series.dtype), "object"): + typed_df[column] = series.astype(desired_dtype) + except (ValueError, TypeError) as exc: + raise ValueError( + f"Failed to cast column '{column}' to {desired_logical}/{desired_dtype}" + ) from exc + + return typed_df + + def _normalize_boolean_like(self, value) -> Optional[bool]: + if pd.isna(value): + return None + if isinstance(value, bool): + return value + if isinstance(value, (int, float)) and not pd.isna(value): + if value == 1: + return True + if value == 0: + return False + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"true", "t", "yes", "y", "1"}: + return True + if normalized in {"false", "f", "no", "n", "0"}: + return False + return None + + def _coerce_boolean(self, series: pd.Series) -> pd.Series: + coerced = [self._normalize_boolean_like(value) for value in series] + return pd.Series(coerced, index=series.index, dtype="boolean") diff --git a/macrobase_py/tests/conftest.py b/macrobase_py/tests/conftest.py new file mode 100644 index 000000000..662d3b3fd --- /dev/null +++ b/macrobase_py/tests/conftest.py @@ -0,0 +1,6 @@ +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) diff --git a/macrobase_py/tests/test_loader.py b/macrobase_py/tests/test_loader.py new file mode 100644 index 000000000..059ee14f8 --- /dev/null +++ b/macrobase_py/tests/test_loader.py @@ -0,0 +1,105 @@ +import pandas as pd +import pytest + +from macrobase_py.data.loader import ColumnMetadata, DataLoader + + +def test_csv_ingestion_and_schema_detection(tmp_path): + df = pd.DataFrame( + { + "id": [1, 2, 3, 4], + "value": ["10.5", "20.1", None, "42.0"], + "category": ["A", "A", "B", "C"], + "timestamp": pd.date_range("2024-01-01", periods=4, freq="D"), + "flag": [True, False, True, None], + } + ) + csv_path = tmp_path / "sample.csv" + df.to_csv(csv_path, index=False) + + loader = DataLoader(csv_path) + + loaded = loader.load() + assert loaded["value"].dtype.kind in {"f", "i"} + assert loaded["timestamp"].dtype.kind == "M" + assert loaded["category"].dtype.name == "category" + + schema = loader.schema + assert isinstance(schema["value"], ColumnMetadata) + assert schema["value"].logical_type == "numeric" + assert schema["category"].categorical is True + assert schema["flag"].logical_type == "boolean" + assert schema["value"].null_count == 1 + + +def test_chunked_iteration_applies_schema(tmp_path): + df = pd.DataFrame( + { + "id": [1, 2, 3, 4], + "value": ["1", "2", "3", "4"], + } + ) + csv_path = tmp_path / "chunked.csv" + df.to_csv(csv_path, index=False) + + loader = DataLoader(csv_path, chunk_size=2) + + batches = list(loader.iter_batches()) + assert len(batches) == 2 + assert all(batch["value"].dtype.kind in {"i", "f"} for batch in batches) + + +def test_excel_ingestion(tmp_path): + pytest.importorskip("openpyxl") + + df = pd.DataFrame( + { + "name": ["alice", "bob"], + "score": [10, 20], + } + ) + excel_path = tmp_path / "sample.xlsx" + try: + df.to_excel(excel_path, index=False) + except ModuleNotFoundError: + pytest.skip("openpyxl is required for Excel tests") + + loader = DataLoader(excel_path) + loaded = loader.load() + + assert loaded.equals(df) + assert loader.schema["name"].logical_type == "categorical" + assert loader.schema["score"].logical_type == "numeric" + + +def test_schema_overrides(tmp_path): + df = pd.DataFrame({"text": ["1", "2"]}) + csv_path = tmp_path / "override.csv" + df.to_csv(csv_path, index=False) + + loader = DataLoader( + csv_path, + schema_overrides={"text": {"logical_type": "text", "categorical": False}}, + ) + + assert loader.schema["text"].logical_type == "text" + assert loader.schema["text"].categorical is False + + +def test_invalid_chunking_for_excel(tmp_path): + excel_path = tmp_path / "sample.xlsx" + try: + pd.DataFrame({"a": [1]}).to_excel(excel_path, index=False) + except ModuleNotFoundError: + pytest.skip("openpyxl is required for Excel tests") + + with pytest.raises(ValueError): + DataLoader(excel_path, chunk_size=1) + + +def test_unknown_file_type(tmp_path): + path = tmp_path / "data.unknown" + path.write_text("a,b\n1,2\n", encoding="utf-8") + + with pytest.raises(ValueError): + DataLoader(path)