Machine Learning in Lyft's Business Context Questions

import os
import pyarrow.parquet as pq
import pyarrow.compute as pc
from datetime import datetime
from torch.utils.data import Dataset
import torch

class ParquetTripDataset(Dataset):
    def __init__(self, root_dir, date_col='trip_date',
                 start_date=None, end_date=None,
                 numeric_cols=None, stats=None, transform=None):
        """
        root_dir: base dir with partitioned parquet files (subdirs or files)
        date_col: column name storing date (ISO or YYYY-MM-DD)
        start_date,end_date: filter inclusive (strings or datetime)
        numeric_cols: list of numeric feature names to normalize
        stats: dict {col: (mean, std)}. If None and numeric_cols provided, compute streamingly.
        transform: optional callable to apply to each sample
        """
        self.root_dir = root_dir
        self.date_col = date_col
        self.transform = transform
        self.numeric_cols = numeric_cols or []
        self.files = []  # list of file paths that match date filter
        self.file_row_counts = []  # per-file total rows
        self.cum_offsets = []  # prefix sums for indexing
        self.stats = stats

        # normalize date inputs
        def to_dt(x):
            if x is None: return None
            return x if isinstance(x, datetime) else datetime.fromisoformat(x)
        self.start_date = to_dt(start_date)
        self.end_date = to_dt(end_date)

        # discover parquet files (non-recursive could be adapted)
        for dirpath, _, filenames in os.walk(root_dir):
            for fn in filenames:
                if fn.endswith('.parquet'):
                    self.files.append(os.path.join(dirpath, fn))

        # lazily gather metadata: row counts and optional date-based filter using metadata columns/row-groups
        filtered_files = []
        for fp in self.files:
            pqf = pq.ParquetFile(fp)
            # sum rows across row groups
            nrows = pqf.metadata.num_rows
            # try to read min/max of date_col from row-group metadata if present
            accept = True
            if self.start_date or self.end_date:
                # scan row-group statistics
                accept = False
                for rg in range(pqf.metadata.num_row_groups):
                    rg_meta = pqf.metadata.row_group(rg)
                    try:
                        col_index = pqf.schema_arrow.get_field_index(self.date_col)
                        col_meta = rg_meta.column(col_index)
                        stats = col_meta.statistics
                        if stats is None:
                            accept = True; break  # can't decide, keep file
                        min_dt = datetime.fromisoformat(stats.min.decode() if isinstance(stats.min, (bytes,bytearray)) else stats.min)
                        max_dt = datetime.fromisoformat(stats.max.decode() if isinstance(stats.max, (bytes,bytearray)) else stats.max)
                        if (self.start_date is None or max_dt >= self.start_date) and (self.end_date is None or min_dt <= self.end_date):
                            accept = True; break
                    except Exception:
                        accept = True; break
            if accept:
                filtered_files.append((fp, nrows))
        self.files = [f for f,_ in filtered_files]
        self.file_row_counts = [r for _,r in filtered_files]
        # build cumulative offsets
        s = 0
        for cnt in self.file_row_counts:
            self.cum_offsets.append(s)
            s += cnt
        self.total_rows = s

        # compute stats if needed (single streaming pass, memory-efficient)
        if self.numeric_cols and self.stats is None:
            # compute mean and std via two-pass streaming over files (or Welford online single pass)
            sums = {c:0.0 for c in self.numeric_cols}
            sumsq = {c:0.0 for c in self.numeric_cols}
            count = 0
            for fp in self.files:
                table = pq.read_table(fp, columns=self.numeric_cols)
                arrs = {c: table[c].to_pandas().astype(float) for c in self.numeric_cols}
                n = len(arrs[self.numeric_cols[0]])
                count += n
                for c in self.numeric_cols:
                    v = arrs[c].values
                    sums[c] += v.sum()
                    sumsq[c] += (v**2).sum()
            self.stats = {}
            for c in self.numeric_cols:
                mean = sums[c]/count
                var = max(0.0, sumsq[c]/count - mean*mean)
                std = var**0.5 if var>0 else 1.0
                self.stats[c] = (mean, std)

    def __len__(self):
        return int(self.total_rows)

    def _locate(self, idx):
        # binary search over cum_offsets
        lo, hi = 0, len(self.cum_offsets)-1
        while lo <= hi:
            mid = (lo+hi)//2
            start = self.cum_offsets[mid]
            end = start + self.file_row_counts[mid]
            if idx < start: hi = mid-1
            elif idx >= end: lo = mid+1
            else: return mid, idx - start
        raise IndexError

    def __getitem__(self, idx):
        if idx < 0: idx += len(self)
        if idx < 0 or idx >= len(self): raise IndexError
        file_idx, local_row = self._locate(idx)
        fp = self.files[file_idx]
        # read only needed row using row_group/row selection if supported
        # simplest: read with row_groups where that row resides
        pqf = pq.ParquetFile(fp)
        # find row-group containing local_row
        acc = 0
        for rg in range(pqf.metadata.num_row_groups):
            rg_rows = pqf.metadata.row_group(rg).num_rows
            if local_row < acc + rg_rows:
                row_in_rg = local_row - acc
                table = pqf.read_row_group(rg, columns=None, use_pandas=False)
                row = {name: table.column(i)[row_in_rg].as_py() for i,name in enumerate(table.column_names)}
                break
            acc += rg_rows
        # optional date filter: verify row matches date range (could skip earlier)
        if self.date_col:
            dt = row.get(self.date_col)
            if dt is not None:
                dt_parsed = dt if isinstance(dt, datetime) else datetime.fromisoformat(dt)
                if (self.start_date and dt_parsed < self.start_date) or (self.end_date and dt_parsed > self.end_date):
                    raise IndexError  # this can happen if metadata was coarse; caller can handle or DataLoader shuffle reduces collisions

        # normalize numeric features on-the-fly
        features = {}
        for c in self.numeric_cols:
            val = row.get(c)
            if val is None:
                norm = 0.0
            else:
                mean, std = self.stats.get(c, (0.0, 1.0))
                norm = (float(val) - mean) / std if std != 0 else 0.0
            features[c] = torch.tensor(norm, dtype=torch.float32)
        sample = {**row, **features}
        if self.transform:
            sample = self.transform(sample)
        return sample