# Spark Policy Toolkit: Semantic Contracts and Scalable Execution for Policy Learning in Spark

arXiv: https://arxiv.org/abs/2604.25061 · HTML: https://mlsanity.com/paper/2604.25061/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 8 of its 19 references that are in the map.
- First posted on arXiv 27 Apr 2026.

## Builds on (most cited first, 8 of 8)

- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [Estimation and Inference of Heterogeneous Treatment Effects using Random Forests](https://mlsanity.com/paper/1510.04342/) (2015)
- [Recursive Partitioning for Heterogeneous Causal Effects](https://mlsanity.com/paper/1504.01132/) (2015)
- [CausalML: Python Package for Causal Machine Learning](https://mlsanity.com/paper/2002.11631/) (2020)
- [MLlib: Machine Learning in Apache Spark](https://mlsanity.com/paper/1505.06807/) (2015)
- [Distributed Double Machine Learning with a Serverless Architecture](https://mlsanity.com/paper/2101.04025/) (2021)
- [mlf-core: a framework for deterministic machine learning](https://mlsanity.com/paper/2104.07651/) (2021)
- [MMLSpark: Unifying Machine Learning Ecosystems at Massive Scales](https://mlsanity.com/paper/1810.08744/) (2018)

## Abstract (arXiv)

Custom policy-learning pipelines in Spark fail for two coupled systems reasons: rowwise Python execution makes inference impractical, and driver-side candidate materialization makes split search fragile at feature scale. We present Spark Policy Toolkit, a semantics-governed systems toolkit for scalable policy learning in Spark. The toolkit provides two Spark-native primitives: partition-initialized vectorized inference through mapInPandas and mapInArrow, and collect-less split search that scores candidates on executors. Both primitives are governed by one fixed-input semantic contract: the same rows, feature order, treatment vocabulary, preprocessing manifest, and split boundaries must preserve per-row score vectors, best-split decisions, and end-to-end learned policy outputs. The evaluation combines practical baseline ladders, backend parity checks, measured split-search scale results, synthetic and Hillstrom end-to-end policy preservation, missingness stress, partition and order perturbation tests, quantile-boundary sensitivity, and a concrete adversarial failure catalog. On a 40-worker Databricks cluster, mapInArrow reaches 4.72M rows/s at 10M matched rows and 7.23M rows/s at 50M rows, while collect-less split search remains valid from F = 10 through F = 1000 with 124000 candidate rows, where the driver-collect baseline is intentionally skipped. Across 24 backend-ablation settings, mapInArrow wins 18 while mapInPandas wins 6, so the paper treats backend choice as workload-dependent rather than universal. Once the fixed-input lock is enforced, all six tested repartition/coalesce/shuffle perturbations preserve identical signatures; before lock, all six drift. The central result is not speed alone: throughput and collect-less execution are the mechanisms that let policy semantics survive at Spark scale.