Feature engineering with Spark: best practices explained

Feature engineering with Spark: best practices explained

Getting raw data into your ML pipeline is only half the battle. The other half—crafting the right signals that models can actually learn from—often determines whether your predictions land or miss. Feature engineering with Spark: best practices explained covers the techniques that let us build and manage features at scale, from simple aggregations to time-aware transformations that capture business logic without leaking information from the future.

In this article, we'll walk through production patterns we've refined across e-commerce, trading, and SaaS workloads. We're not aiming for exhaustive theory; instead, we're focusing on concrete approaches that solve real problems when you're handling billions of events and need features ready for model training or inference.

Why feature engineering matters in Spark

Raw data rarely feeds directly into a model. A user ID, a timestamp, and a purchase amount don't tell the model much until we ask: What is this user's average order value over the last 30 days? or How does this purchase compare to their typical behavior? These questions become features—and answering them at scale requires both Spark's parallelism and careful thought about what we're computing and when.

Feature engineering in Spark differs from doing it on a single machine. We can't just loop through rows; we think in partitions, joins, and window functions. The payoff is processing terabytes of data in minutes, but the cost is that mistakes—like including future information in historical features—are amplified and harder to catch.

Aggregations per user or product

The foundation of most ML features is the aggregate. User-level aggregations answer questions like: How many orders has this customer placed?, What's their total spending?, What's their average review rating? Product-level aggregations ask: How many units sold?, What's the trend in views?

These are typically computed from event-level data (orders, clicks, reviews) grouped by the entity of interest. In Spark, a groupBy followed by an aggregation is straightforward, but scale and correctness are where care matters.

from pyspark.sql import functions as F
from pyspark.sql.window import Window

# Read raw events: user_id, event_time, amount, product_id
events = spark.read.parquet("s3://data/events")

# User-level aggregations
user_features = events.groupBy("user_id").agg(
    F.count("*").alias("total_orders"),
    F.sum("amount").alias("total_spent"),
    F.avg("amount").alias("avg_order_value"),
    F.stddev("amount").alias("stddev_order_value"),
    F.max("event_time").alias("last_order_time")
)

# Product-level aggregations
product_features = events.groupBy("product_id").agg(
    F.count("*").alias("total_sold"),
    F.avg("amount").alias("avg_price"),
    F.countDistinct("user_id").alias("unique_buyers")
)

user_features.write.mode("overwrite").parquet("s3://features/user_level")
product_features.write.mode("overwrite").parquet("s3://features/product_level")

In production, we partition these outputs by a snapshot date so we can rebuild features incrementally. Also, pay attention to nulls and edge cases: a user with one order has zero standard deviation, which can break downstream models if not handled.

We've found that computing aggregations once and reusing them across multiple models saves compute and keeps feature definitions consistent. Store them early, use them often.

Window features over time

Aggregations give us lifetime snapshots, but models often need signals that capture recent behavior. A user who spent $100 yesterday is different from one who spent $100 once five years ago. This is where window functions shine: they let us compute rolling sums, moving averages, and time-decay features without expensive self-joins.

In Spark, the Window API partitions data by an entity and orders it by time, then applies an aggregate over a frame—typically the last N days or N preceding rows.

# Compute 7-day and 30-day rolling sum of spending per user
events_sorted = events.withColumn(
    "event_date", F.to_date("event_time")
).select("user_id", "event_date", "amount")

# Define windows for different lookbacks
window_7d = Window.partitionBy("user_id") \
    .orderBy(F.col("event_date").cast("long")) \
    .rangeBetween(-7*86400, 0)  # seconds

window_30d = Window.partitionBy("user_id") \
    .orderBy(F.col("event_date").cast("long")) \
    .rangeBetween(-30*86400, 0)

rolling_features = events_sorted.withColumn(
    "spending_7d", F.sum("amount").over(window_7d)
).withColumn(
    "spending_30d", F.sum("amount").over(window_30d)
).withColumn(
    "order_count_7d", F.count("*").over(window_7d)
)

rolling_features.write.mode("overwrite").parquet("s3://features/rolling")

A subtlety here: when you use rangeBetween with time windows, Spark compares the ordering column directly. For timestamps, cast to long (seconds or milliseconds) so the range math is correct.

Time windows also expose leakage risks. If you're predicting churn for a user on Monday and you include their Monday spending in the 7-day window, you've leaked information. We'll return to this in the point-in-time section.

Write to a feature store

Once you've computed features, where do they live? A feature store is a central repository—often built on top of a database or object store—that serves features to both training and inference pipelines. It keeps feature logic versioned, enables reuse, and makes it possible to serve consistent features at prediction time.

On a practical level, we typically write Spark outputs to a feature store by:

1. Writing features to partitioned Parquet or Delta in S3/ADLS, keyed by the entity (user_id, product_id) and a snapshot date.
2. Registering metadata in a catalog (e.g., a custom table in a metadata database) that tracks feature name, version, computation logic, and freshness.
3. Reading from the feature store during training to construct training examples, and during inference to fetch the latest features.

Using Delta Lake is a strong choice if your infra supports it: ACID guarantees mean we can safely overwrite, and time travel lets us grab features from a specific date for reproducibility.

# Write to Delta with partitioning by snapshot date
snapshot_date = "2024-01-15"

user_features.withColumn("snapshot_date", F.lit(snapshot_date)) \
    .write \
    .mode("overwrite") \
    .partitionBy("snapshot_date") \
    .format("delta") \
    .save("s3://feature_store/user_features")

# Read for training: all users, specific snapshot
training_features = spark.read.format("delta") \
    .load("s3://feature_store/user_features") \
    .filter(F.col("snapshot_date") == "2024-01-15")

In a real system, you'd also track which features are used by which models, maintain SLAs on freshness (e.g., "this feature is refreshed daily"), and log feature values for debugging. A lightweight metadata layer—even just a few Spark DataFrame operations—pays dividends.

Avoid leakage with point-in-time joins

This is where many teams stumble, and it's critical: point-in-time correctness means that for a training example with a label date, you only include features computed from data available before that date. Include future information, and your model looks brilliant in testing but fails in production.

Suppose we're predicting whether a user will churn on 2024-01-15. We should only see features computed from events up to 2024-01-14. Accidentally including 2024-01-15 events—even if they're just the label itself—introduces leakage.

The standard approach: compute features for multiple snapshots (daily, weekly) and join them to labels using a point-in-time lookup. For a label with date D, find the most recent feature snapshot with date < D.

# Labels: user_id, label_date, churned
labels = spark.read.parquet("s3://labels/churn")

# Features with snapshots: user_id, snapshot_date, spending_7d, etc.
features = spark.read.parquet("s3://feature_store/user_features")

# Point-in-time join: for each label, get features from the most recent snapshot before label_date
from pyspark.sql import functions as F

# Add a snapshot_date to labels (the date we want features from, e.g., day before label)
labels_with_feature_date = labels.withColumn(
    "feature_lookup_date", F.date_sub(F.col("label_date"), 1)
)

# Join features to labels: user_id matches, and feature snapshot_date <= feature_lookup_date
pit_join = labels_with_feature_date.join(
    features,
    on=[
        F.col("labels_with_feature_date.user_id") == F.col("features.user_id"),
        F.col("features.snapshot_date") <= F.col("labels_with_feature_date.feature_lookup_date")
    ],
    how="left"
).groupBy("labels_with_feature_date.user_id", "label_date") \
    .agg(
        F.max(F.col("snapshot_date")).alias("actual_snapshot_date"),
        F.first(F.col("spending_7d")).alias("spending_7d"),
        F.first(F.col("order_count_7d")).alias("order_count_7d")
    )

pit_join.write.mode("overwrite").parquet("s3://training_data/churn_examples")

The groupBy and max ensures we're taking the most recent snapshot before the label date. This is the "point-in-time" join—and it's non-negotiable for honest model evaluation.

Common mistakes

• Forgetting to handle nulls in aggregations: A user with no orders returns NULL for sum and avg. Either fill with 0 (for counts) or a sentinel value, or be explicit about missing data in your model.
• Mixing snapshot dates in joins: Accidentally joining today's features to yesterday's labels. Always explicitly filter features to the snapshot date you intend to use.
• Not partitioning the feature store: Writing all features to a single partition makes reads slow and rebuilds expensive. Partition by snapshot date and entity type.
• Computing features at inference time from training data: If your training features use a 30-day rolling sum, you can't compute that at inference without access to 30 days of history. Pre-compute and cache it.
• Ignoring data drift: Features computed from old data may not represent current behavior. Monitor feature distributions and refresh on a schedule that matches your domain.

Key takeaways

• Start with simple user and product aggregations, then layer in time-aware features using Spark's Window functions.
• Write computed features to a centralized feature store (Parquet, Delta, or a dedicated platform) partitioned by snapshot date for easy replay.
• Always perform point-in-time joins when building training data: features must come from a snapshot strictly before the label date.
• Partition and version your features; a catalog tracking what's computed, when, and which models use it saves hours of debugging.
• Test for leakage by verifying that no feature includes data from the label or prediction date—this is the most common production failure mode.
• Compute features once, use them many times; reusing a feature definition across models keeps pipelines consistent and cuts compute costs.