The Feature Engineering Crossroads
There's a point in every ML project where you're processing more data than fits in memory on a single machine. That's when the conversation about distributed feature engineering starts, and invariably someone proposes either Spark MLlib or distributed PyTorch. They solve overlapping problems but come from fundamentally different design philosophies, and picking the wrong one costs you months.
I've run both in production for feature engineering at scale, and the short version is: Spark MLlib wins for tabular data transformations on datasets measured in terabytes. Distributed PyTorch wins when your features involve learned representations or when the boundary between feature engineering and model training is blurry.
Spark MLlib Pipeline Model
MLlib thinks about feature engineering as a series of deterministic transformations applied to DataFrames. You chain together Transformers and Estimators into a Pipeline object, and the whole thing can be serialized, versioned, and replayed. This is incredibly useful for production systems where reproducibility matters.
from pyspark.ml import Pipeline
from pyspark.ml.feature import (
StringIndexer, VectorAssembler,
StandardScaler, Bucketizer
)
indexer = StringIndexer(
inputCol="device_type", outputCol="device_idx"
)
bucketizer = Bucketizer(
splits=[0, 1, 5, 20, 100, float("inf")],
inputCol="session_count",
outputCol="session_bucket"
)
assembler = VectorAssembler(
inputCols=["device_idx", "session_bucket",
"avg_duration", "page_views", "bounce_rate"],
outputCol="raw_features"
)
scaler = StandardScaler(
inputCol="raw_features", outputCol="features",
withStd=True, withMean=True
)
pipeline = Pipeline(stages=[indexer, bucketizer, assembler, scaler])
model = pipeline.fit(train_df)
features_df = model.transform(train_df)The strength here is that Spark handles the distributed execution transparently. You write the transformation logic once, and it runs across 500 executors without you thinking about data partitioning or communication patterns. A well-tuned Spark job processes about 2TB of tabular data per hour on a 20-node cluster.
Where MLlib Falls Short
MLlib transformation library is extensive but shallow. You get standard scalers, encoders, tokenizers, and basic statistical transformations. What you don't get is anything that requires gradient-based optimization as part of the feature extraction. If you need embeddings from a pretrained model, sequence features from an RNN, or any feature that involves a neural network forward pass, MLlib can't help you directly.
The workaround is to use Spark UDFs that wrap PyTorch inference, but this is clunky. Each executor loads a copy of the model, the serialization overhead is significant for large models, and error handling across the boundary between Spark and PyTorch is painful.
Related reading: Distributed Training with DeepSpeed ZeRO: Practical Configur.
Distributed PyTorch for Learned Features
When your features involve neural network inference, distributed PyTorch is the natural choice. You're already in the ecosystem, the model loading is straightforward, and you can use the same infrastructure for both feature generation and model training.
import torch
import torch.distributed as dist
from torch.utils.data import DataLoader, DistributedSampler
from transformers import AutoModel, AutoTokenizer
def generate_embeddings(rank, world_size, data_paths, output_dir):
dist.init_process_group("nccl", rank=rank, world_size=world_size)
torch.cuda.set_device(rank)
model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
model = model.to(rank)
model.eval()
dataset = TextDataset(data_paths)
sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
loader = DataLoader(dataset, batch_size=256, sampler=sampler,
num_workers=4, pin_memory=True)
embeddings = []
with torch.no_grad():
for batch in loader:
inputs = {k: v.to(rank) for k, v in batch.items()}
outputs = model(**inputs)
emb = outputs.last_hidden_state.mean(dim=1)
embeddings.append(emb.cpu())
all_embeddings = torch.cat(embeddings, dim=0)
torch.save(all_embeddings, f"{output_dir}/embeddings_rank{rank}.pt")
dist.destroy_process_group()The Data Loading Challenge
The biggest practical problem with distributed PyTorch for feature engineering isn't the compute, it's getting the data to the GPUs fast enough. Spark was built for this; it reads from HDFS, S3, or Delta Lake with optimized I/O and columnar access. PyTorch data loading was designed for image files on local disk, and adapting it for terabyte-scale tabular data takes real engineering work.
We've had good results with Petastorm, which bridges the gap by providing a PyTorch-compatible DataLoader that reads directly from Parquet files on distributed storage.
Hybrid Architectures That Work
In practice, the best feature engineering pipelines use both. Spark handles the heavy tabular transformations, joins, aggregations, window functions, encoding, and writes intermediate features to Parquet. Then PyTorch reads those features, generates learned representations, and writes the final feature vectors back to storage.
The handoff point between Spark and PyTorch is the critical design decision. Put it too early and you're doing simple transformations in PyTorch that Spark would handle more efficiently. Put it too late and you're trying to run neural network inference inside Spark UDFs, which is the worst of both worlds.
Related reading: Real-Time Video Analytics Architecture on Edge Devices.
My rule of thumb: if the operation is a SQL-expressible transformation (group by, join, filter, aggregate), do it in Spark. If it requires a model forward pass or gradient computation, do it in PyTorch. If it's a simple mathematical operation, either tool works fine, pick whichever is upstream of the next step.
One last thing: whatever you build, make sure the feature engineering pipeline is idempotent. Running it twice on the same input data should produce identical output. This sounds obvious but it's surprisingly easy to violate when you have got random sampling, non-deterministic GPU operations, or time-dependent features in the mix.
Performance Benchmarking Methodology
When comparing Spark MLlib and distributed PyTorch for your specific workload, the benchmarking methodology matters as much as the results. Here is the framework we use.
For Spark, we measure end-to-end pipeline execution time including data reading, transformation, and writing results. We test at three data scales: 100GB, 1TB, and 10TB. The important metric is throughput (GB processed per hour), not latency for individual records. We also track the peak memory usage per executor and the shuffle data volume, because these determine your cluster sizing.
For PyTorch, we measure GPU utilization alongside throughput. If your GPUs are sitting idle waiting for data, you have a data loading bottleneck, and throwing more GPUs at the problem won't help. We aim for at least 75 percent GPU utilization during feature generation. Below that, the data pipeline needs optimization before the compute pipeline.
For a related perspective, see Camera Calibration and Geometric Transforms for 3D Vision Ap.
The cost comparison needs to account for infrastructure differences. Spark runs on CPU instances (typically r-series on AWS for memory-intensive feature engineering). PyTorch needs GPU instances, which cost 3-10x more per hour. But if the PyTorch pipeline finishes 5x faster because of GPU acceleration, the total cost might be lower. Always compare total job cost, not per-hour instance cost.
Error Handling and Recovery
Both systems handle failures differently, and this affects your production reliability.
Spark has built-in stage-level retry logic. If an executor fails during a transformation, Spark re-schedules that partition on another executor automatically. For idempotent transformations, this is transparent. The catch is that speculative execution can cause duplicate writes if your output format doesn't handle concurrent writers gracefully.
PyTorch distributed training is less forgiving. If one worker fails, the entire training group typically needs to restart from the last checkpoint. The elastic training features in PyTorch 2.0 (torchelastic) improve this by allowing workers to rejoin, but it requires additional setup and doesn't work with all distributed strategies.
For production pipelines, we add our own checkpointing layer on top of both systems. Every N batches, we write intermediate results to durable storage. If the job fails, it restarts from the last checkpoint rather than from scratch. This is especially important for multi-hour jobs where a failure at 90 percent completion would otherwise waste all the preceding compute.
When to Consider Ray
I should mention Ray as a third option that bridges the gap between Spark and PyTorch. Ray Data provides a DataFrame-like API that can distribute both standard transformations and GPU-accelerated operations. It doesn't have Spark MLlib built-in pipeline abstractions, but it avoids the JVM-Python boundary issues that plague Spark UDFs.
We've started using Ray for new projects where the feature engineering involves a mix of tabular and embedding operations. The programming model feels more natural than the Spark-to-PyTorch handoff, and the performance is competitive with both. The ecosystem is still maturing though, so expect to hit rough edges that Spark and PyTorch solved years ago.