Working with Large-Scale Datasets Questions
Analyzing data that does not fit comfortably in memory or a single query. Covers distributed processing concepts, partitioning and sampling strategies, query and pipeline performance, and trade-offs when scaling an analysis. Focuses on getting correct answers efficiently at scale.
You need to compute PCA on a numeric dataset stored as a 200GB binary matrix on disk. Explain how to use numpy.memmap for memory-efficient access and design an approach that computes the leading k principal components without loading the whole dataset into RAM. Provide code sketches and discuss trade-offs between accuracy and memory usage.
Sample Answer
First, use numpy.memmap to avoid loading the 200GB file into RAM — memmap gives array-like access that fetches pages on demand.
High-level approach:
- One-pass compute column means (needed to center data) by reading in chunks.
- Use a memory-efficient algorithm for leading-k PCs:
- Option A: randomized SVD on a streamed matrix (block-wise multiply: Y = X @ Omega computed by summing X_chunk @ Omega).
- Option B: Incremental PCA (online covariance accumulation or sklearn.decomposition.IncrementalPCA) processing chunks sequentially.
Code sketch (randomized SVD using memmap and one-pass multiplications):
import numpy as np
from numpy.linalg import svd
# open memmap (dtype, shape known)
X = np.memmap('data.bin', dtype=np.float32, mode='r', shape=(n_rows, n_cols))
def compute_col_mean(memmap, chunk_rows=100000):
mean = np.zeros(memmap.shape[1], dtype=np.float64)
total = 0
for i in range(0, memmap.shape[0], chunk_rows):
chunk = memmap[i:i+chunk_rows].astype(np.float64)
mean += chunk.sum(axis=0)
total += chunk.shape[0]
return mean / total
mean = compute_col_mean(X)
def randomized_pca(memmap, k, oversample=10, n_iter=2, chunk_rows=100000):
ncols = memmap.shape[1]
omega = np.random.randn(ncols, k+oversample)
# Y = (X - mean) @ omega computed in chunks
Y = np.zeros((memmap.shape[0], omega.shape[1]), dtype=np.float64)
# Instead of storing full Y (could be large), compute Q via tall-skinny QR in streaming fashion:
# Here we accumulate Y^T Y to compute orthonormal basis via eigen-decomp (alternative streaming)
B = np.zeros((omega.shape[1], omega.shape[1]), dtype=np.float64)
for i in range(0, memmap.shape[0], chunk_rows):
chunk = memmap[i:i+chunk_rows].astype(np.float64)
chunk -= mean
Z = chunk @ omega # small matrix
B += Z.T @ Z
# eigen-decompose B to get approximate subspace
eigvals, eigvecs = np.linalg.eigh(B)
idx = np.argsort(eigvals)[::-1][:k]
Qhat = eigvecs[:, idx]
# Form small matrix B2 = Qhat^T * (X^T X) * Qhat computed by streaming if needed, then SVD
# For brevity: construct projected matrix via another pass and do final SVD
return Qhat # map back to principal components by additional passes
Practical simpler alternative: sklearn.decomposition.IncrementalPCA with memmap:
from sklearn.decomposition import IncrementalPCA
ipca = IncrementalPCA(n_components=k)
for i in range(0, X.shape[0], chunk_rows):
chunk = X[i:i+chunk_rows].astype(np.float64)
chunk -= mean
ipca.partial_fit(chunk)
# then transform or get components_
Trade-offs and considerations:
- Centering requires at least one full pass (compute mean) — unavoidable unless using approximate centering.
- Chunk size trades IO overhead vs memory: larger chunks reduce loop overhead but need more RAM.
- Randomized methods: fewer passes (1–3) and low memory but approximate; more power iterations improve accuracy at cost of extra passes and IO.
- Incremental PCA: single-pass partial_fit, stable and simple, but may be less accurate than full SVD.
- Numerical stability: use float64 where possible for accumulation (means/covariance) even if stored data is float32.
- I/O dominates runtime for 200GB; parallel reads or SSDs help.
- If exact principal components are required, you must form covariance (or perform SVD) which may need more passes or more memory.
Summary: use memmap + chunked mean, then either IncrementalPCA for simplicity or streamed randomized SVD for better accuracy/control; pick chunk size and number of power iterations to balance accuracy vs memory/IO.
A long-running pandas job takes hours on a single machine. Outline a step-by-step migration plan to Dask (or Ray DataFrame) including minimal code changes to get started, how to test correctness, common unsupported pandas methods and workarounds, and how to measure and validate performance improvements and scaling behavior.
Sample Answer
Requirements & constraints:
- Keep code changes minimal; correctness identical to pandas; scale from single machine to cluster; support GPU later; measure speedup and scaling.
Migration plan (step-by-step):
- Inventory & prioritize: profile job to find hot paths (cProfile, pandas .info(), memory). Identify dataframe sizes, IO, groupbys, joins, custom apply.
- Proof-of-concept (PoC) locally:
- Replace imports and create Dask/Ray DataFrame wrappers.
- Run a small sample (10% of data) to validate logic and outputs.
- Expand to full dataset on single machine with multiple workers/threads.
- Move to cluster: configure scheduler (Dask Scheduler / Ray cluster), tune partitions and worker resources.
- Optimize: persist/cache, repartition, avoid wide shuffles, vectorize custom functions (numba/cython), use GPUs when available.
- Monitor & iterate: instrument metrics, compare results and perf.
Minimal code changes (examples):
Dask:
import dask.dataframe as dd
# read
ddf = dd.read_parquet("s3://bucket/data/*.parquet")
# previously: df = pd.read_parquet(...)
# then keep same pandas-like chain, compute at end
result = ddf.groupby("key")["val"].sum().compute()
Ray DataFrame (Modin on Ray or ray.data):
import modin.pandas as pd # minimal drop-in for many pandas calls
df = pd.read_parquet("data.parquet")
res = df.groupby("key")["val"].sum()
Testing correctness:
- Unit tests on small deterministic samples comparing pandas vs Dask/Ray: use pandas .equals() or numeric tolerances.
- Integration tests on full dataset sample; compare checksums, row counts, key aggregates.
- Use deterministic partitions and sort keys to compare ordered outputs.
- Add property tests for commutativity/associativity where applicable.
Common unsupported pandas methods & workarounds:
- .apply on rows: slow or unsupported — replace with vectorized ops, map_partitions, or use .map_partitions with pandas functions.
- Iteration (.itertuples/.iterrows): avoid; use vectorized transforms.
- Complex multi-index operations: may require compute(), convert to pandas, or redesign to single-index.
- inplace=True semantics: avoid relying on side-effects; prefer functional style.
- .pivot_table with complex agg: use groupby+unstack or compute then pivot in pandas on smaller subsets.
Measuring & validating performance & scaling:
- Baseline: record wall-clock, CPU, memory, I/O and peak memory on pandas run.
- Microbenchmarks: measure key operators (read, groupby, join, apply) across frameworks.
- Scaling tests:
- Strong scaling: fixed problem size, increase workers; expect reduced runtime until overhead dominates.
- Weak scaling: increase data proportional to workers; aim for constant runtime.
- Use Dask dashboard / Ray dashboard, Prometheus/Grafana for cluster metrics.
- Report speedup = T_pandas / T_dask and efficiency = speedup / N_workers.
- Verify no correctness regressions after scaling: re-run checksums/aggregates.
Trade-offs & risks:
- Some pandas semantics differ (lazy vs eager, partitioning). Expect initial correctness checks and possible algorithm refactoring.
- Shuffles (joins/groupby) dominate cost — minimize and tune partitioning.
- Choose Dask for rich pandas API compatibility; Ray/Modin offers easier drop-in but fewer advanced APIs.
This plan gets you from minimal code change PoC to validated, scalable distributed runs while ensuring correctness and measurable gains.
Compare using pandas vs executing SQL queries via DuckDB for complex multi-join analytical queries. Provide an example where DuckDB can outperform pandas on a multi-join aggregation, show how to convert a pandas DataFrame to DuckDB for processing and back, and explain I/O and memory trade-offs.
Sample Answer
Brief comparison (high-level):
- pandas: in-memory, excellent API ergonomics, vectorized ops; best for datasets that fit memory and when you need Python-native row/column manipulation and custom functions.
- DuckDB: vectorized analytical SQL engine, columnar, can operate on disk-backed files (Parquet/CSV) and stream from Python objects; excels at large multi-join, group-by, windowed analytics thanks to optimized query planning, late materialization, and single-node parallelism.
When DuckDB outperforms pandas
- Scenario: three large tables (events ~50M rows, users ~5M, products ~1M). Query: join all three, filter, then group-by product category with complex aggregations (distinct counts, percentiles). pandas will need to materialize joins in memory (possibly requiring repeated copies), while DuckDB can pipeline joins and aggregations, use hash join spilling to disk, and push filters early.
Example: convert pandas -> DuckDB, run multi-join aggregation, convert result back
import duckdb
import pandas as pd
# assume df_events, df_users, df_products are large pandas DataFrames
con = duckdb.connect()
# register pandas DataFrames as DuckDB views (zero-copy where possible)
con.register('events', df_events)
con.register('users', df_users)
con.register('products', df_products)
query = """
SELECT p.category,
COUNT(*) AS events,
COUNT(DISTINCT u.user_id) AS unique_users,
approx_quantile(e.value, 0.5) AS median_value
FROM events e
JOIN users u ON e.user_id = u.user_id
JOIN products p ON e.product_id = p.product_id
WHERE e.timestamp >= '2025-01-01'
GROUP BY p.category
ORDER BY events DESC
"""
result_df = con.execute(query).df() # returns a pandas DataFrame
Key reasons DuckDB wins here:
- Query planner pushes WHERE/filter predicates before joins.
- Joins and aggregations are vectorized and can spill to disk, avoiding OOM.
- Built-in approximate aggregations (approx_quantile) are optimized C++ implementations.
- Avoids repeated Python-level row iteration or creating many intermediate DataFrames.
I/O and memory trade-offs
- pandas: everything in-memory; fast for small-to-medium data, but OOM risk when joins create intermediate copies. Serialization cost if reading/writing to disk (Parquet/CSV).
- DuckDB: can operate in-memory or spill to disk, reads columnar Parquet efficiently (predicate/column pruning). When data already in Parquet on disk, DuckDB reads only required columns and can be much faster and more memory efficient than loading full DataFrames into pandas.
- Conversion cost: registering DataFrame in DuckDB is cheap (often zero-copy); calling .df() materializes results back to pandas — cost proportional to result size. For very large aggregated results (still small relative to raw data), this is acceptable.
Practical advice for AI Engineers
- Use DuckDB for feature engineering over large analytic tables, join-heavy ETL, or when working with Parquet data lakes. Materialize final feature set into a compact parquet or pandas DataFrame for model training.
- For custom per-row Python transforms require pandas or apply in DuckDB via UDFs (newer DuckDB Python UDF support), but prefer pushing as SQL or vectorized ops to retain performance.
You want to port a pandas pipeline to GPU using RAPIDS cuDF to accelerate groupby and merge operations. Discuss what pandas APIs are supported by cuDF, what code changes are typically required, how to handle unsupported operations (fallback to CPU), and provide an example replacing pd.read_csv with cudf.read_csv then performing a groupby. Also discuss GPU memory limitations and strategies to handle them.
Sample Answer
High-level approach
- cuDF implements a large subset of pandas' DataFrame/Series APIs (arithmetic, indexing, groupby/agg, join/merge, many datetime and string ops, read_csv, categorical). But not everything (custom Python-pandas apply, many complex .pivot_table edge cases, some multi-index ops, third‑party extension dtypes). Typical porting is "drop‑in" for many workflows but requires validation and a few code changes and fallbacks.
Typical code changes
- Replace imports: pandas -> cudf (or use both and convert).
- Read/write: pd.read_csv -> cudf.read_csv; df.to_parquet/from_parquet have cuDF equivalents.
- dtype tuning: cast to optimized GPU dtypes (int32, float32, categoricals).
- Avoid Python-level rowwise operations (df.apply with Python funcs): replace with vectorized cuDF ops, numba/cupy kernels, or run on CPU.
- Conversions: df.to_pandas() and cudf.from_pandas() where needed.
- Use dask-cudf for datasets larger than one GPU.
Handling unsupported operations (fallback patterns)
- Detect and fallback selectively:
- Bring only the small, unsupported part to CPU: df_cpu = df.to_pandas(); df_cpu['col'] = df_cpu['col'].apply(py_func); df = cudf.from_pandas(df_cpu)
- Or run entire pipeline on CPU if many ops unsupported.
- Use try/except and feature-detection:
import cudf, pandas as pd
def read_and_group(path):
try:
gdf = cudf.read_csv(path) # GPU read
res = gdf.groupby('key')['value'].sum()
return res
except Exception as e:
# fallback to pandas for compatibility
pdf = pd.read_csv(path)
return pdf.groupby('key')['value'].sum()
- Consider rewriting custom Python functions as numba.cuda, cupy, or use apply_rows/apply_grouped (cuDF UDF APIs) where available.
Example: replace pd.read_csv then groupby
import cudf
# read on GPU
gdf = cudf.read_csv('data.csv') # similar args as pd.read_csv
# cast to efficient dtypes
gdf['user_id'] = gdf['user_id'].astype('int32')
gdf['event'] = gdf['event'].astype('category')
# groupby and aggregate on GPU
agg = gdf.groupby('user_id').agg({'event': 'count', 'value': 'sum'})
# if needed back to pandas for a CPU-only step:
pdf = agg.to_pandas()
GPU memory limitations & strategies
- Limitation: GPU memory is typically 16–80+ GB but still far smaller than host RAM; cuDF operates in GPU memory — OOMs are common on large datasets.
- Strategies:
- Use dask-cudf to shard data across GPUs and process in parallel with out-of-core support.
- Chunked processing: read in smaller batches (cudf.read_csv has chunksize via iterator patterns or use pandas chunking + to_gpu per chunk).
- Reduce memory footprint: drop unused columns early, downcast dtypes, convert strings to categoricals, compress with parquet.
- External shuffle/aggregation: perform partial aggregations per chunk and reduce (map-reduce pattern) to avoid materializing full joins/groupbys.
- Spill to host: use libraries that support host spill (RAPIDS + UCX + Dask can spill).
- Monitor and tune memory pool: set RMM (RAPIDS Memory Manager) with pool sizing and enable managed memory or pinning for better allocation.
- Multi-GPU: horizontal scale with dask-cudf; use NVLink/UCX for inter-GPU transfers.
- Use nvprof/nsight and RAPIDS profiler to find hotspots and memory peaks.
Key tradeoffs and tips
- Start by porting IO, simple vectorized ops, joins and groupbys — these give biggest wins.
- Keep heavy Python UDFs on CPU or rewrite them for GPU with numba/cupy to maximize benefits.
- Validate numeric equivalence and null handling — semantics sometimes differ slightly.
- Use dask-cudf when dataset > single GPU or to pipeline chunked reads + map-reduce aggregations.
This approach minimizes code churn while leveraging GPU speed for groupby/merge-heavy pipelines and provides clear fallbacks and memory strategies when datasets or operations exceed GPU capabilities.
Implement (in Python) a streaming-friendly function that reads inference logs with fields (request_id, model_version, start_ts, end_ts nullable, status) and computes approximate p95 and p99 latencies per model_version over the last 24 hours using a single pass and limited memory. Provide code sketch, explain choice of streaming quantile algorithm (e.g., t-digest), and analyze time and memory complexity.
Sample Answer
Approach: single-pass streaming with limited memory by keeping, per model_version, a ring of fixed-duration time buckets (e.g., 1-minute) each containing a compact streaming quantile sketch (t-digest). For each completed inference (end_ts present), compute latency = end_ts - start_ts, map it to the bucket for end_ts, and update that bucket's t-digest. To report p95/p99 for the last 24 hours, merge the t-digests of the buckets within the 24h window for that model_version and query quantiles. This bounds memory (24h worth of buckets) and is merge-friendly.
Code sketch (requires tdigest package: pip install tdigest):
from tdigest import TDigest
from collections import defaultdict, deque
import time
import math
BUCKET_SECONDS = 60 # 1-minute buckets
WINDOW_SECONDS = 24 * 3600 # 24 hours
BUCKETS_PER_WINDOW = WINDOW_SECONDS // BUCKET_SECONDS
class ModelWindow:
def __init__(self):
# deque of (bucket_start_ts, TDigest)
self.buckets = deque()
def _ensure_bucket(self, bucket_start):
if self.buckets and self.buckets[-1][0] == bucket_start:
return self.buckets[-1][1]
# append new bucket
td = TDigest()
self.buckets.append((bucket_start, td))
# drop old buckets
while self.buckets and bucket_start - self.buckets[0][0] >= WINDOW_SECONDS:
self.buckets.popleft()
return td
def add(self, ts, value):
bucket_start = ts - (ts % BUCKET_SECONDS)
td = self._ensure_bucket(bucket_start)
td.update(value)
def get_quantiles(self, now_ts=None):
if not self.buckets:
return None
now_ts = now_ts or int(time.time())
# merge relevant buckets into a temp TDigest
merged = TDigest()
cutoff = now_ts - WINDOW_SECONDS
for bstart, td in self.buckets:
if bstart >= cutoff:
merged = TDigest.merge(merged, td)
if merged.n == 0:
return None
return {'p95': merged.quantile(0.95), 'p99': merged.quantile(0.99)}
class LatencyAggregator:
def __init__(self):
self.per_model = defaultdict(ModelWindow)
def process_log(self, request_id, model_version, start_ts, end_ts, status):
# Expect timestamps as epoch seconds (or convert)
if end_ts is None or status != 'completed':
return
latency = end_ts - start_ts
if latency < 0:
return
self.per_model[model_version].add(end_ts, latency)
def get_model_metrics(self, model_version, now_ts=None):
return self.per_model[model_version].get_quantiles(now_ts)
Why t-digest:
- Accurate for extreme quantiles (p95/p99) and handles skewed latency distributions well.
- Compact (small set of centroids), supports online updates and fast merge operations — ideal for time-bucketed, distributed streams.
- Low memory per sketch and fast to query quantiles.
Complexity:
- Per record time: O(log k) amortized to update a t-digest (k = centroid count). Bucket lookup O(1). Merging when querying is O(b * k) where b = number of buckets in window (<= BUCKETS_PER_WINDOW).
- Memory: O(M * b * k) where M = number of model_versions being tracked, b ≈ 1440 for 1-minute buckets, k typically 50-200 centroids. This is bounded and tunable (increase BUCKET_SECONDS or reduce k if memory constrained).
Edge cases & notes:
- Handle null end_ts (incomplete requests) by ignoring or tracking separately.
- Choose bucket size to trade latency of window accuracy vs memory.
- If model count is very large, evict inactive models or use LRU expiry.
- Use epoch ms if higher precision needed; ensure consistent units.
That is every published Working with Large-Scale Datasets question for AI Engineer so far. Browse the other topics in this category, or practice this one interactively.