PrismaDV: Automated Task-Aware Data Unit Test Generation

Abstract: Data is a central resource for modern enterprises, and data validation is essential for ensuring the reliability of downstream applications. However, existing automated data unit testing frameworks are largely task-agnostic: they validate datasets without considering the semantics and requirements of the code that consumes the data. We present PrismaDV, a compound AI system that analyzes downstream task code together with dataset profiles to identify data access patterns, infer implicit data assumptions, and generate task-aware executable data unit tests. To further adapt the data unit tests over time to specific datasets and downstream tasks, we propose "Selective Informative Feedback for Task Adaptation" (SIFTA), a prompt-optimization framework that leverages the scarce outcomes from the execution of data unit tests and downstream tasks. We evaluate PrismaDV on two new benchmarks spanning 60 tasks across five datasets, where it consistently outperforms both task-agnostic and task-aware baselines in generating unit tests that reflect the end-to-end impact of data errors. Furthermore, we show that with SIFTA, we can automatically learn prompts for PrismaDV's modules that outperform prompts written by hand or generated from a generic prompt optimizer. We publicly release our benchmarks and prototype implementation.

• The paper introduces task-aware test generation by integrating data profiling with downstream code analysis, yielding specialized validation constraints.
• It employs a modular architecture with LLM-based assumption inference and SIFTA prompt optimization to improve error detection and system generalization.
• Benchmark evaluations show PrismaDV achieves F1 scores of 87.8% on ICDBench and 77.4% on EIDBench, substantially outperforming task-agnostic approaches.

Automated Task-Aware Data Unit Test Generation with PrismaDV

Motivation and Problem Statement

Data unit testing is critical for robust data pipelines, but industry-standard frameworks (such as Deequ, TensorFlow Data Validation, Great Expectations) are predominantly task-agnostic, inferring constraints solely by profiling observed data. This neglects implicit assumptions and domain-specific requirements encoded in downstream task code, resulting in missed errors and false alarms. The paper introduces the problem of task-aware data unit test generation, wherein data validation constraints are specialized per downstream task, aligning with the task's actual access patterns and embedded domain logic.

Figure 1: A toy ETL scenario exemplifying failures of task-agnostic unit tests and motivating specialization to downstream task assumptions.

PrismaDV System Overview

PrismaDV is an AI-driven, modular system that automates task-aware data unit test generation by jointly analyzing dataset profiles and downstream application code. The architecture decomposes the overall process into four principal modules:

1. Profiling and Discovery: Computes descriptive statistics of input data and detects the columns accessed by downstream code.
2. Assumption Inference: Applies code analysis and LLM-based summarization to infer implicit assumptions present in task logic, generating natural language representations.
3. Constraint Code Generation: Synthesizes executable constraints targeting those assumptions, expressed in validation DSLs (e.g., Deequ).
4. Post-processing: Validates code generation syntactically and semantically, discarding inconsistent constraints.

The system builds a bipartite data-code assumption graph as an intermediate artifact, mapping data columns to inferred assumptions and relevant code locations.

Figure 2: PrismaDV constructs a data-code assumption graph, leveraging it to generate specialized validation logic for each downstream task.

Selective Informative Feedback for Task Adaptation (SIFTA)

To facilitate iterative improvement and domain adaptation, PrismaDV employs SIFTA, a prompt optimization framework that harnesses execution outcomes from both data unit tests and downstream tasks. SIFTA targets scarcity and ambiguity in supervision by quantifying "failure precision" — the fraction of constraint failures coincident with task failures. Optimization is focused on informative constraint failures and leverages backtracing to the assumptions and code lines responsible. This enables targeted prompt updates for LLM-based modules, improving system generalization across new data batches and tasks.

Benchmarking and Experimental Evaluation

Two dedicated benchmarks are proposed:

• ICDBench: Evaluates fine-grained discovery of individual constraints from data-code pairs across diverse domains and semantics.
• EIDBench: Assesses end-to-end impact detection, quantifying system performance on 60 tasks across five datasets with injected errors and ground-truth validation.

PrismaDV is compared against outlier detection, task-agnostic test generators, LLM prompting, and agentic systems. Key results include:

• On ICDBench, PrismaDV achieves an F1 score of 87.8%, exceeding baselines by over 20 points.
• On EIDBench, PrismaDV reaches 77.4% F1, strongly outperforming LLM zero-shot and few-shot baselines in error impact detection.
• SIFTA prompt optimization yields gains over manual prompts (up to 5.13% F1 increase in new data generalization scenarios), and outperforms general-purpose prompt optimizers like GEPA.

Implications and Future Directions

PrismaDV demonstrates that coupling data profiling with downstream code semantics yields robust, specialized validation artifacts, reducing false alarms and missed critical errors. The modular framework leveraging LLMs for code understanding and synthesis is extensible, enabling deployment across varied operational contexts. The explicit use of failure precision signals in SIFTA addresses supervision scarcity typical in production pipelines, making it viable for iterative adaptation.

Practically, integration with modern ML pipeline orchestration (e.g., Modyn [10.1145/3709705], CatDB [10.1145/372(2212.37250)97]), data discovery systems, and diverse enterprise workloads remains straightforward. Conceptually, PrismaDV advances the notion that data validation infrastructure should be proactive and task-aware, promoting model and application reliability.

Future work could include extending code analysis across multi-file and multi-table tasks, scaling benchmarks to industrial repositories, and exploring more advanced search strategies within SIFTA (e.g., Monte Carlo tree search). Automated synthetic error injection for prompt optimization would further enable cold-start adaptation.

Conclusion

PrismaDV provides a principled approach to task-aware automated data validation, bridging the gap between data-centric and system-centric testing. Its modular architecture, prompt optimization strategy, and strong empirical results establish it as a reference for both practical adoption and further research in AI-driven data infrastructure (2604.21765).