Python Data Manipulation with Pandas & PySpark Questions

Techniques for cleaning, transforming, and analyzing data in Python using Pandas and PySpark. Covers working with DataFrames, data wrangling, missing-value handling, filtering, aggregations, joins, grouping, and typical patterns for data preparation and exploratory analysis, including both in-memory Pandas workflows and distributed PySpark processing.

HardTechnical

59 practiced

Design a robust A/B testing metric computation pipeline in PySpark to compute experiment metrics nightly across billions of rows. Requirements: deduplicate exposures/events, attribute events to the correct exposure and time window, compute per-cohort aggregated metrics and confidence intervals, and support incremental re-computation. Sketch schema, joins, and job orchestration details.

Sample Answer

Requirements & constraints:- Nightly batch on billions of rows- Deduplicate exposures/events, attribute events to correct exposure & time window- Per-cohort aggregates + confidence intervals- Support incremental re-compute (late arrives, backfills)

High-level design:- Sources: exposures (impression_id, user_id, treatment, timestamp, attrs), events (event_id, user_id, event_type, timestamp, attrs)- Output: metrics_by_cohort(partitioned by date, treatment, cohort_key) with counts, rates, CI, version metadata

Schema sketch:- exposures(exposure_id STRING, user_id STRING, treatment STRING, exp_ts TIMESTAMP, attrs MAP, dedupe_ts TIMESTAMP)- events(event_id STRING, user_id STRING, event_type STRING, evt_ts TIMESTAMP, attrs MAP, dedupe_ts TIMESTAMP)

Deduplication & canonicalization:1. Define canonical keys: exposure by exposure_id; if missing, use (user_id, treatment, windowed exp_ts floor) and min(ts).2. Use Spark window dedupe: rank by dedupe_ts DESC, keep first.3. For events, dedupe by event_id; if absent, use (user_id, event_type, evt_ts rounded to minute) and take earliest.

Attribution logic (PySpark sketch):- Join events to exposures by user_id where evt_ts in [exp_ts, exp_ts + attribution_window] and pick nearest preceding exposure per user/event (use range join + row_number partition by event_id order by abs(evt_ts - exp_ts))- Efficient implementation: broadcast small treatment assignment table; for big joins use interval join via Spark 3 range join (using condition) or explode exposures into attribution windows bucketed by date.

PySpark pseudocode:

python

# dedupe exposures
from pyspark.sql import Window, functions as F
w = Window.partitionBy('exposure_id').orderBy(F.col('dedupe_ts').desc())
expos_canonical = exposures.withColumn('rn', F.row_number().over(w)).filter("rn=1").drop('rn')

# dedupe events
we = Window.partitionBy('event_id').orderBy(F.col('dedupe_ts').desc())
events_canonical = events.withColumn('rn', F.row_number().over(we)).filter("rn=1").drop('rn')

# attribution: join on user_id and time window, pick closest prior exposure
joined = events_canonical.alias('e').join(
  expos_canonical.alias('x'),
  (F.col('e.user_id')==F.col('x.user_id')) &
  (F.col('e.evt_ts') >= F.col('x.exp_ts')) &
  (F.col('e.evt_ts') <= F.col('x.exp_ts') + F.expr('interval 7 days')),
  how='left'
)
w2 = Window.partitionBy('e.event_id').orderBy(F.abs(F.unix_timestamp('e.evt_ts')-F.unix_timestamp('x.exp_ts')))
attr = joined.withColumn('rn', F.row_number().over(w2)).filter('rn=1').drop('rn')

Cohorting & metrics:- Define cohorts (e.g., signup_date, geography) from exposure attrs.- Aggregate per partition (date = exposure_date), treatment, cohort: counts, conversions, sums.- Confidence intervals: - For binomial rates use Wilson or Clopper-Pearson for accuracy. - For mean metrics use bootstrap if non-normal or delta method for efficiency. - Implement analytic variance: var(p̂)=p(1-p)/n; CI = p̂ ± z*sqrt(var). For A/B diff, compute pooled variance and effect CI.

Incremental recompute & partitioning:- Partition outputs by exposure_date, treatment, cohort_key.- Nightly job processes only changed partitions: new exposures/events and late-arriving windows. Maintain high-water mark per source and include grace/backfill window (e.g., 7 days).- For recompute, read partitions affected and overwrite them atomically (use Hive ACID or write to staging then atomic swap).

Orchestration & reliability:- Use Airflow (or Dagster) DAG: 1. Ingest validation & schema check 2. Dedupe exposures/events 3. Attribution job (Spark) 4. Aggregation + CI computation 5. Publish metrics to warehouse / BI / experiment service 6. Backfill/reconciliation task weekly to detect drift- Use job checkpoints, write lineage metadata (source offsets, row counts, hashes) for reproducibility.- Use Spark optimizations: partition by user_id hash for joins, use broadcast for small tables, enable Adaptive Query Exec, cache intermediate if reused.

Edge cases & monitoring:- Missing exposure -> attribute to "no-exposure" cohort- Multiple overlapping exposures -> tie-break by latest exposure prior to event- Late-arriving events beyond grace window -> flag for manual backfill- Monitor sample ratio, variance anomalies, monotonicity checks

Trade-offs:- Exact nearest-exposure attribution is more accurate but heavier joins; bucketing exposures into windows reduces cost.- Bootstrap gives robust CIs but is costlier; use analytic where possible and sample-based bootstrap for heavy scale.

This design scales to billions via partitioned Spark jobs, incremental processing, and careful join strategies while ensuring correctness, reproducibility, and observability.

HardTechnical

48 practiced

Describe and implement optimizations to reduce the memory and CPU overhead when converting a large Spark DataFrame to Pandas via toPandas(). Include code snippets using Arrow, column selection, sampling, repartitioning, and driver memory settings. Explain trade-offs.

Sample Answer

To safely convert a large Spark DataFrame to Pandas with reduced memory/CPU overhead, combine selective projection, sampling, Arrow serialization, controlled partitioning, and driver memory tuning. Trade-offs: lower memory vs potential sampling bias, more shuffle vs parallelism, and Arrow requiring compatible dtypes.

Approach:1. Select only required columns and cast to Arrow-friendly dtypes.2. Optionally sample or filter to reduce rows.3. Repartition to a controlled number of partitions (not too many small tasks).4. Enable Arrow and use toPandas().5. Adjust driver memory and use chunked conversions if still large.

Code:

python

from pyspark.sql import SparkSession
spark = SparkSession.builder \
    .appName("toPandas-optim") \
    .config("spark.sql.execution.arrow.pyspark.enabled", "true") \
    .config("spark.driver.memory", "16g") \
    .getOrCreate()

# 1. Column selection and casting
needed_cols = ["id", "feature1", "feature2"]
df_sel = df.select(*[
    df[c].cast("double") if c.startswith("feature") else df[c] for c in needed_cols
])

# 2. Optional sampling / filtering
df_small = df_sel.sample(withReplacement=False, fraction=0.1, seed=42)

# 3. Repartition to match driver cores/memory (fewer partitions -> fewer pandas concatenations)
df_part = df_small.repartition(8)  # tune: ~driver_cores * 2

# 4. Convert using Arrow
pdf = df_part.toPandas()

If still too big, stream in batches:

python

# iterate partitions and process incrementally to avoid full materialization
rdd = df_part.rdd
for i, part in enumerate(rdd.toLocalIterator()):
    pdf_part = spark.createDataFrame(part).toPandas()
    # process or write out incrementally

Key points & trade-offs:- Arrow: much faster, lower serialization overhead, but requires compatible dtypes (use .cast()).- Column selection & casting: biggest win — reduce data transferred.- Sampling: reduces cost but can bias results unless used for exploratory work.- Repartitioning: balance between parallelism and driver overhead; too many partitions causes many small pandas frames and higher merge cost.- Driver memory: increase spark.driver.memory and spark.driver.maxResultSize; too high may be infeasible on shared clusters.- Alternative: use toPandas on a downsampled dataset or write partitions to parquet and load with pandas in parallel on a sufficiently sized VM.

MediumTechnical

45 practiced

You're preparing to write a large dataset (~1TB) to Parquet partitioned by event_date using PySpark. Provide code that repartitions appropriately, writes compressed Parquet to s3://bucket/path/ with snappy compression, and explain best practices around partition column choice and avoiding small files.

Sample Answer

Approach: compute a reasonable number of output files (based on target file size, e.g. 128 MB), repartition the DataFrame by event_date while controlling total parallelism so partitions produce reasonably large files, then write Parquet with snappy compression to S3. Key trade-off: partitioning by a high-cardinality column creates many small folders/files; choose partition columns that align with query patterns (date/time, low-to-medium cardinality) and aim for files ~64–256MB.

python

# python
from pyspark.sql import SparkSession
import math

spark = SparkSession.builder.appName("write_parquet").getOrCreate()

# df: your 1TB DataFrame with a column 'event_date' as 'yyyy-MM-dd' or date type
# Example assumptions:
total_size_bytes = 1 * 1024**4  # 1 TB in bytes
target_file_size = 128 * 1024**2  # 128 MB target
# compute approximate number of output files
num_output_files = max(1, int(total_size_bytes / target_file_size))

# Decide how many shuffle partitions to use; keep a practical upper bound
max_tasks = min(num_output_files, 2000)

# Repartition by event_date but also control total partitions so files are sizable.
# This creates up to `max_tasks` partitions, grouped by event_date.
df_repart = df.repartition(max_tasks, "event_date")

# Write partitioned Parquet to S3 with snappy compression
(
    df_repart
    .write
    .mode("overwrite")  # or "append" depending on use-case
    .option("compression", "snappy")
    .partitionBy("event_date")
    .parquet("s3a://bucket/path/")
)

Best practices and notes:- Choose partition columns that are commonly used in filters (e.g., event_date) and have moderate cardinality. Avoid high-cardinality features (user_id) as partition keys.- Aim for file sizes 64–256 MB to balance read throughput and task overhead. Calculate num files = total_data_size / target_file_size.- Use repartition(n, partition_col) to both group by partition column and control total files. repartitionByRange can help produce balanced ranges for date-like columns.- For S3, use the s3a:// connector and set committer/consistency configs for reliable writes (e.g., EMRFS or Hadoop S3A optimized committers) to avoid eventual-consistency issues.- If source data skew leads to many small files for some partitions, consider coalescing small partitions, combining files with an extra job, or writing per-date batches.- Monitor output with AWS Glue/Athena or list objects; perform occasional compaction jobs to merge small files.

EasyTechnical

51 practiced

As a PySpark starter task: write Python code to create a SparkSession, read a CSV file at s3://my-bucket/data.csv with header and inferSchema turned on, show the inferred schema and 10 sample rows. Add config to set shuffle partitions to 200 for the session.

Sample Answer

To do this in PySpark we create a SparkSession with the desired config (shuffle partitions = 200), read the CSV with header and inferSchema enabled, print the inferred schema, and show 10 sample rows.

python

from pyspark.sql import SparkSession

# Create SparkSession with shuffle partitions set to 200
spark = (
    SparkSession.builder
    .appName("ReadCSVFromS3")
    .config("spark.sql.shuffle.partitions", "200")
    .getOrCreate()
)

# Path to CSV on S3
s3_path = "s3://my-bucket/data.csv"

# Read CSV with header and infer schema
df = (
    spark.read
    .option("header", "true")
    .option("inferSchema", "true")
    .csv(s3_path)
)

# Show inferred schema
df.printSchema()

# Show 10 sample rows
df.show(10, truncate=False)

# Stop session when done
spark.stop()

Key points:- spark.sql.shuffle.partitions controls the default number of partitions used for shuffle operations; 200 is good default for moderate clusters.- inferSchema triggers an extra pass to detect types; for large files prefer providing schema for efficiency.- Ensure IAM credentials / Hadoop/S3 configs are set (AWS keys, role, or instance profile) and correct S3 connector is available. Edge cases: very large CSV (schema inference cost), inconsistent headers, malformed rows — consider .option("mode","DROPMALFORMED") or providing explicit schema.

MediumTechnical

51 practiced

Explain when to use cache() / persist() in Spark. Compare storage levels (MEMORY_ONLY, MEMORY_AND_DISK, MEMORY_ONLY_SER), describe how to choose them based on workload (iterative ML training vs one-off ETL), and show how to unpersist and monitor cache usage.

Sample Answer

Use cache()/persist() when you will reuse an RDD/DataFrame multiple times and want to avoid recomputing its lineage or re-reading source data. cache() is shorthand for persist(StorageLevel.MEMORY_ONLY).

When to cache/persist- Iterative ML training (e.g., repeated epochs, cross-validation, feature engineering reused across folds): persist the preprocessed training DataFrame so transforms and shuffles aren’t recomputed each epoch.- One-off ETL: don’t persist unless you read the same dataset multiple times in the job; otherwise streaming/pipe stages are fine and caching wastes memory.

Storage levels — trade-offs- MEMORY_ONLY: keeps deserialized JVM objects in RAM. Fastest CPU access, lowest IO, but if not enough memory, partitions are recomputed on demand (no disk spill).- MEMORY_AND_DISK: keeps as much as fits in RAM; spills remaining partitions to disk. Good for large datasets where you prefer correctness over recompute cost.- MEMORY_ONLY_SER: stores serialized binary objects in memory (smaller footprint, less GC pressure) but CPU overhead to (de)serialize. Good when memory is tight and you accept CPU cost.

How to choose- Iterative ML (many reuses, moderate dataset size): MEMORY_AND_DISK or MEMORY_ONLY_SER if memory constrained; if fits comfortably, MEMORY_ONLY for max speed.- One-off ETL (single pass): avoid persisting; if you need temporary reuse within a small scope, persist with MEMORY_AND_DISK to be safe.

Examples (PySpark)

python

from pyspark import StorageLevel

# cache shorthand
df = spark.read.parquet("/data/train")
df.cache()  # MEMORY_ONLY

# explicit persist with storage level
df.persist(StorageLevel.MEMORY_AND_DISK)

# serialized option
df.persist(StorageLevel.MEMORY_ONLY_SER)

Unpersist and monitoring- Unpersist when done to free memory: df.unpersist(blocking=True)- Monitor via Spark UI -> Storage tab (shows RDD/DataFrame blocks, used memory, disk size, number of partitions). Programmatic checks:

python

spark.catalog.isCached("train_table")  # if registered as table
spark.catalog.clearCache()  # clears all cached tables

- Use SparkListener or REST API / metrics to alert on cache pressure (usedMemory, maxMemory) and tune executor memory or storage level accordingly.

Key rules- Persist only when reuse outweighs memory cost.- Prefer MEMORY_AND_DISK for large iterative jobs; use serialized storage to reduce footprint at CPU cost.- Always unpersist after the workload and monitor via the Spark UI to avoid eviction and GC issues.

Unlock Full Question Bank

Get access to hundreds of Python Data Manipulation with Pandas & PySpark interview questions and detailed answers.

Join thousands of developers preparing for their dream job.