One ML pipeline, three execution layers: PySpark, Polars, and Rust

The expensive part may be the boundary

It is easy to draw a pipeline containing PySpark, Polars, Rust, and SageMaker. It is harder to explain why every boundary belongs there. Moving a dataset between execution engines can introduce serialization, materialization, schema drift, and another environment to reproduce. A faster transformation is not necessarily a faster experiment if preparing its inputs consumes the time it saves.

I would evaluate this architecture as a sequence of falsifiable choices. Does the preparation workload need distribution? Does a local columnar engine shorten iteration without changing feature semantics? Is there a measured hot path that justifies custom native code? Can training and evaluation consume immutable artifacts rather than relying on hidden notebook state? This article defines the experiment. It does not report a benchmark or claim that a three-language pipeline is the best default.

Give each execution layer a narrow job

PySpark is a candidate for large joins, broad historical backfills, and workloads that need distributed execution. Before changing engines, inspect its query plan, partition sizes, shuffle volume, and skew. Apache Spark documents how partitioning, statistics, join selection, and adaptive execution affect performance. An expensive join caused by a bad key distribution will not necessarily disappear when the surrounding application is rewritten.

Polars is a candidate for local feature exploration and reproducible transformations over a bounded dataset or suitable streaming plan. Its lazy API gives the optimizer visibility over the query before execution. I would begin with a scan, select the necessary columns, and push filters into the plan where supported. Calling a lazy engine after eagerly reading everything into memory forfeits part of that opportunity.

Rust enters only after profiling identifies a custom operation that the existing engines cannot express efficiently enough. Polars already has a Rust implementation; replacing a vectorized expression with an external Rust service can add a network trip without improving the underlying work. A native parser, specialized validator, or carefully batched extension is a more plausible candidate than a wholesale rewrite of orchestration code.

Freeze the dataset before comparing tools

The experimental unit should be an immutable input manifest. It identifies object versions or content hashes, schema, extraction logic, eligibility rules, and the event-time range. It also identifies the feature contract. If one implementation silently drops malformed records while another retains them as nulls, their runtime comparison is not comparing the same computation.

I would prepare three workload families: a join-heavy historical build, a narrow local iteration, and a custom parsing workload. Each family should include an ordinary case and a stress case, such as skewed keys, many small files, or unexpectedly long strings. The shapes are part of the benchmark record. A single conveniently sized dataset is not enough to characterize an engine.

The following manifest is a synthetic example of the record I want. Its paths and hash are placeholders, not published artifacts. Production manifests should capture actual immutable versions, avoid sensitive identifiers, and include a reference output or defined numerical tolerance. The distinction between a data checksum and a semantic equality check matters when engines serialize the same table differently.

1{
2  "study": "feature-build-boundaries",
3  "input_manifest": "s3://example-data/manifests/frozen-input.json",
4  "schema_version": "orders-v3",
5  "feature_spec": "features-v2",
6  "split_spec": "event-time-with-label-delay",
7  "implementations": [
8    "spark-sql",
9    "polars-lazy",
10    "native-parser-ablation"
11  ],
12  "required_checks": [
13    "key-set-equality",
14    "null-policy",
15    "decimal-precision",
16    "feature-equivalence"
17  ],
18  "results": "not-yet-measured"
19}

Test equivalent features across engines

Feature parity needs its own fixture suite. I would include missing values, duplicate events, timezone transitions, decimal boundaries, late-arriving updates, and an entity with no historical activity. Aggregation windows need explicit endpoint rules. A transaction at exactly the prediction timestamp should not enter a historical feature merely because one implementation uses an inclusive comparison.

A useful initial transformation selects only the fields required for the experiment and preserves the event timestamp. The small Polars example below validates a simple positive-amount slice; it is not a full training pipeline. Real feature construction would add point-in-time joins, label maturity rules, schema assertions, and an output manifest. Those rules should be implemented once as a contract and checked against every execution engine.

For integer and categorical features, exact comparison is often appropriate. Floating-point reductions may require a justified tolerance because execution order can differ. I would never use a loose tolerance as a blanket escape from investigating a mismatch. Record the feature, its expected numerical range, and why the allowed error cannot change an important downstream decision.

1import polars as pl
2 
3features = (
4    pl.scan_parquet("fixtures/orders/*.parquet")
5    .select("order_id", "customer_id", "event_at", "amount_minor")
6    .filter(pl.col("amount_minor").is_not_null())
7    .filter(pl.col("amount_minor") > 0)
8    .with_columns(
9        pl.col("amount_minor").cast(pl.Int64),
10        pl.col("customer_id").cast(pl.String),
11    )
12)
13 
14# Inspect before materializing; use a bounded fixture first.
15print(features.explain())
16result = features.collect()
17assert result["order_id"].null_count() == 0

Orchestrate artifacts, not invisible state

SageMaker Pipelines represents dependencies between processing, training, and evaluation steps. I would use those dependencies to make artifact flow explicit: a preparation output becomes a training input, a model artifact becomes an evaluation input, and an evaluation report becomes evidence for a release decision. The pipeline definition should expose what changed between two experiments rather than merely proving that several jobs ran.

The preparation environment and training environment need separate identities and permissions. A feature job may read a larger raw dataset than the training job should ever access. Pin container images and dependencies, preserve configuration, and record the exact artifact selected for evaluation. A successful job is not evidence that it used the intended data, especially when mutable prefixes are reused between runs.

I would keep the first orchestration deliberately small. A processing step writes a validated feature artifact and manifest; training writes a model and training record; evaluation writes both machine-readable metrics and failure examples. Any promotion gate should check the identity of those artifacts as well as its metric thresholds. Reusing a stale evaluation report is a provenance failure even when its score looks excellent.

Partial output is another boundary worth testing. I would interrupt a feature job after it writes some files but before it completes validation. Downstream training should consume only a committed manifest of complete outputs, never a wildcard that happens to include both finished and abandoned files. A retry should write to a distinct run location and publish its manifest only after checks pass. This small failure injection tests whether the pipeline really operates on artifacts or still depends on timing. It also makes cleanup safer: abandoned work can be identified without deleting files referenced by a successful run.

Measure the whole experiment and keep the losing result

The benchmark should report elapsed time from input availability to validated output, compute consumption, peak memory, bytes transferred, and engineering complexity. Break down startup, reading, transformation, writing, and verification. Repeat runs under declared warm and cold conditions. Fix the resource envelope or report its differences honestly. A larger machine is a valid engineering choice, but it is not evidence that a programming language alone caused the improvement.

For a Rust optimization, retain the original implementation as an ablation. Compare the full pipeline with and without the native component, not only the inner loop. Include packaging, cross-platform builds, observability, and failure recovery in the decision record. If the native step saves little overall time but introduces a fragile deployment dependency, the experiment has produced a useful reason to remove it.

The result I want is a boundary decision supported by measurements. Perhaps Spark remains necessary for the historical join while Polars makes a bounded study simpler. Perhaps the existing vectorized operation eliminates the need for custom Rust entirely. Perhaps SageMaker orchestration is useful while a local runner remains better for fixture tests. A credible research pipeline can explain those choices, reproduce their evidence, and change them when the workload changes.