arXiv:2607.00254v1 [cs.DB] 30 Jun 2026
Query-Centric Optimization of AI Workflows via Approximate Query Processing and Proxy Models Huayi Wang
Jun Xu
Gromit Yeuk-Yin Chan
Georgia Institute of Technology Atlanta, USA [email protected]
Georgia Institute of Technology Atlanta, USA [email protected]
Adobe Research San Jose, USA [email protected]
ABSTRACT
1
Many modern AI workflows—ranging from LLM post-training pipelines to agentic reasoning tasks—can be expressed as declarative queries whose expensive predicate is evaluated by a large model or reward function. We propose a query-centric formulation of these workflows and show that classical database techniques, namely approximate query processing (AQP) and proxy-model (PM) based filtering, can substantially reduce the number of expensive model invocations without requiring changes to the underlying models or pipelines. Our first strategy treats the workflow as an online aggregation problem: it progressively samples records, maintains a running aggregate estimate with a confidence interval, and terminates early once the interval stabilizes, accepting the estimate when it falls within a user-specified error bound. Our second strategy trains a lightweight, CPU-resident decision tree on a small set of oracle-labeled examples and uses it to pre-filter records whose outcome can be predicted with high confidence, routing only uncertain records to the expensive model. We evaluate both strategies on TPC-DS aggregate queries and on real LLM post-training pipelines including math reasoning, general instruction following, and code generation. On TPC-DS, Strategy AQP keeps aggregate error under 10% while reaching its adaptive stopping point at 10–15% of oracle calls under balanced distributions—an 85–90% reduction—and Strategy PM reduces oracle calls by 60–70% on natural-label workloads. On LLM pipelines, Strategy AQP reaches its adaptive stopping point at 20–50% of oracle calls with less than 5% accuracy loss on the structured math and code tasks (open-ended instruction following, scored by a reward model, shows a larger but bounded reduction), and Strategy PM reduces reward-model scoring time by up to 19× on structured tasks with less than 10% accuracy loss.
The rapid growth of expensive Large Language Model (LLM) computation in AI workflows [4, 9]—from training foundation models to invoking them hundreds of times in agentic task pipelines [3, 20]— has created new opportunities to optimize latency and cost. More importantly, many of the current AI-centric workflows can be expressed as database queries. Consider the following two examples:
PVLDB Reference Format: Huayi Wang, Jun Xu, and Gromit Yeuk-Yin Chan. Query-Centric Optimization of AI Workflows. PVLDB, 14(1): XXX-XXX, 2020. doi:XX.XX/XXX.XX PVLDB Artifact Availability: The source code, data, and/or other artifacts have been made available at https://github.com/Why1221/Query-Centric-AI-optimization.
This work is licensed under the Creative Commons BY-NC-ND 4.0 International License. Visit https://creativecommons.org/licenses/by-nc-nd/4.0/ to view a copy of this license. For any use beyond those covered by this license, obtain permission by emailing [email protected]. Copyright is held by the owner/author(s). Publication rights licensed to the VLDB Endowment. Proceedings of the VLDB Endowment, Vol. 14, No. 1 ISSN 2150-8097. doi:XX.XX/XXX.XX
INTRODUCTION
Example 1.1. A recurring pattern in modern LLM post-training pipelines — explicitly documented in both the Llama 2 and Llama 3 technical reports [7, 26] — is what can be understood as the following SQL query operation: SELECT * FROM model_outputs ORDER BY reward_model_score LIMIT K
Demonstrated in Figure 1, for each training prompt, K=8–32 candidate responses are sampled from the current model, scored by a reward model, and only the top-scoring response is retained as supervised fine-tuning data for the next training round. This rejection sampling step is applied iteratively over six rounds, with the reward model also serving as a retrospective quality filter over the full training corpus. This generate-rank-select loop is a central mechanism by which the model bootstraps its own training data quality without additional human annotation. Example 1.2. Agentic reasoning tasks—workflows in which an AI model must gather context, apply multi-step inference, and produce a decision or action—are typically treated as orchestration problems, implemented through bespoke application logic spread across pipelines, queues, and LLM API calls. Yet consider a canonical example: “computing the average order value among customers whose recent orders suggest churn risk”, a task that today is commonly delegated to an agent framework with hand-written retrieval, prompt assembly, and filtering logic [3, 22]. As illustrated in Figure 2, this workflow can be expressed as the following single declarative query: SELECT customer_id, AVG(order_total) FROM orders WHERE llm_classify(order_description, 'churn_risk') = 'HIGH' GROUP BY customer_id
where llm_classify itself may encapsulate a multi-step agentic pipeline, yet appears in the query simply as a semantic predicate that filters records before a standard aggregation is applied. This observation suggests that agentic reasoning tasks can be naturally viewed as declarative queries with semantic predicates. These two examples, drawn from opposite ends of the LLM lifecycle, illustrate a recurring pattern we refer to as query-centric AI workflows. To increase the throughput of these workflows, we observe an opportunity that lies between the query results and the AI pipeline: query-level properties—such as the distribution of aggregated columns or the selectivity of ranking clauses—can guide
Figure 1: Illustration of rejection sampling as a SQL query in the LLM post-training recipe.
the selective substitution or early termination of expensive model calls, increasing throughput while bounding or empirically controlling the impact on the final result. For example, in the churn query, if order_total is approximately uniformly distributed across churn risk levels, a cheaper surrogate model can replace llm_classify without materially affecting the AVG result; in the rejection sampling query, the LIMIT K clause implies that only the top-𝐾 candidates need to reach the reward model, allowing the remainder to be pruned early. Thus, while naively executing the AI pipeline over all records is prohibitively expensive, query processing systems can exploit these query-result relationships to avoid redundant model invocations and reduce end-to-end cost. Our Approach. To exploit these opportunities, we propose two complementary approaches. First, we cast query-centric AI workflows as an approximate query processing (AQP) [8] problem: rather than invoking the AI pipeline on every record, we draw a progressive sample of rows, compute an estimate of the aggregate result along with a confidence interval, and terminate early once the estimate converges to within a user-specified error bound. Second, we introduce a proxy-based filtering strategy: a lightweight surrogate model is trained to identify records whose AI pipeline output can be predicted with high confidence—routing only the remaining “hard” records to the expensive model. Together, these two approaches reduce the number of expensive model invocations: the AQP-based approach provides statistical confidence-interval control on aggregate results, together with an empirically validated adaptive stopping criterion for ranking (LIMIT) workloads, while the filtering-based approach routes only low-confidence records to the expensive model based on the proxy’s confidence. A key advantage of the query-centric formulation is that it makes explicit a rich set of structural predicates — such as response length, presence of special symbols, or formatting patterns in mathematical outputs — that can be expressed as lightweight WHERE clauses over the generated records. For the proxy-model based approach, these predicates are invisible to inference-based proxies such as smaller reward models [19, 27], which operate purely on semantic content.
Figure 2: Illustration of an analytical query involving AI agents as a User Defined Function (UDF). By treating them as input features, we show that a straightforward decision tree trained on a small set of reward-model-labeled examples can exploit these query-derived signals to filter records early with a favorable cost–quality trade-off. This result underscores a broader point: the query-centric view does not merely reframe existing workflows — it actively surfaces optimization opportunities that are inaccessible from the pipeline view alone. We conduct a comprehensive experimental evaluation on two classes of query-centric AI workloads. The first is a set of TPC-DS queries whose expensive predicate is evaluated by an ML classifier; these experiments demonstrate the effectiveness of the AQP-based approach, showing that progressive sampling achieves <10% aggregate error while invoking the oracle on only around 10% of tuples. The second is a suite of real LLM post-training pipelines spanning math reasoning (GSM8K), general instruction following (UltraFeedback), and code generation (HumanEval+, MBPP), which demonstrate both strategies: the proxy-based filter reduces rewardmodel scoring time by up to 19× with less than 10% downstream quality loss on the structured tasks. Contributions. Our main contributions are as follows: • We formalize query-centric AI workflows as ML queries and propose two complementary strategies to reduce expensive model invocations: an AQP-based approach with statistical confidence-interval control (and an empirically validated adaptive stopping criterion for ranking workloads) and a proxy-based filtering approach that routes only lowconfidence records to the oracle. • We discuss a lightweight strategy-selection rule that chooses between AQP and proxy-based filtering based on query structure, confidence-interval behavior, and observed data characteristics. 2
• We evaluate both strategies on TPC-DS and real LLM pipelines, reducing reward-model scoring time by up to 19× while staying within 10% of full-RM accuracy on structured LLM tasks, and oracle calls by 85–90% with <10% aggregate error on TPC-DS under balanced distributions.
2
model architecture. We refer to these inference steps as targeted AI pipelines. During its usage, the pipeline queries the data that needs to be predicted from sources and acquires the data in a tabular data format. Then, it transforms the input data into a target schema. For example, records are transformed into labels by a classification pipeline, numbers by a regression pipeline, and clusters by a clustering pipeline. Yet, as inference pipelines grow more sophisticated, they become increasingly expensive to execute. We summarize these problems into two goals. Given a database, a query, and an AI pipeline, we want to (1) derive an estimator of the query result based on samples to reduce the number of records sent through the AI pipeline (G1); and (2) build a simpler proxy model that estimates the AI workflow’s result on the query efficiently and accurately (G2). While many query optimization techniques for ML inference have been proposed to speed up the process [9, 12, 19, 27], one of our distinctions is to do it without changing the database architecture. In the industrial setting, it is preferred to have minimal changes on the query processor and the layers underneath due to extra costs incurred on database migration and maintenance. Besides, the database already contains several optimization features like caching and indexing that we can leverage.
BACKGROUND AND PROBLEM
In this section, we describe the background, as well as the inputs, outputs, and goals of our query-centric AI workflow. To begin with, we explain the characteristics of common AI workflows that can be expressed as queries. Consider the following query examples: -- Q1 SELECT AVG(churn_rate), predicted_segment FROM customer_analytics_db GROUP BY predicted_segment -- Q2 SELECT genre FROM movie_db WHERE predicted_rating > 4 -- Q3 SELECT * FROM loan_application_db WHERE predicted_default = True -- Q4 SELECT t1.user, t2.movie FROM user_db as t1 LEFT JOIN movie_db as t2 ON t1.predicted_preferred_genre = t2.genre --Q5 SELECT * FROM census_db ORDER BY predicted_income ASC LIMIT 50000
2.1
Inputs, Outputs, and Goals
Our strategies take a query, a dataset, and an AI workflow as input, and output an approximated query result. Our goals on the query optimization are to reduce the size of the query data to the AI workflow and speed up the remaining queries. To be specific, for the targeted ML pipeline, we expect the output to be either a classification label or a regression value. Our scope of the queries is as follows: (1) Aggregate Queries on a numerical column with the group-by clause based on ML pipeline. (2) Selection Queries on particular rows with the filter clause based on ML pipeline. (3) Limit Queries that control the size of output with sorting criteria based on the ML pipeline.
The queries Q1-5 demonstrate various needs of predictive AI workflows to output designated tables. In a straightforward data science pipeline, we first fetch all the rows into memory and input them into the workflows. Then, after acquiring the predictions, we store them in a separate column and aggregate or filter the rows to output the results. Yet from a data management perspective, the above workflow causes a lot of data processing on all rows in the database, but only a small fraction of rows appear in the final output. For example, in Q3, if the default rate is only 10%, a reliable proxy or early filter may reject many likely-negative entries before they reach the full AI pipeline. To sum up, a query-centric AI workflow contains two types of speed-up opportunities. First, there is a filter condition (Q2-4) where we only need to focus on a subset of labels. We could derive a simpler pipeline that optimizes the subset’s task. Query clauses like WHERE, JOIN, and LIMIT could trigger such a condition. Secondly, there is also an aggregate condition (Q1) where we could consider different weights on different labels. For example, for conditions like GROUP BY, rows with a particular label might contain similar value distribution on the numerical column so that we do not need to query many of them to get the average. Furthermore, given the needs of feature engineering, the AI pipeline might need to query lots of columns to produce the predictions. Many predictive systems are powered by expensive AI pipelines that contain complex (1) feature engineering and (2)
From these queries, we can observe that some predicted labels or values are more important than others. For example, we only need to focus on those data points whose predicted values are in a certain range. Thus, a query-centric AI workflow, unlike typical Machine Learning (ML) pipelines that assign equal importance to each row’s predictions, results in different preferences on the pipeline output. To understand the different importance among different AI workflow outputs and assign them as weights to the proxy model, our strategies first conduct a query result size estimation based on the query and pipeline, and then apply either an estimator to predict results with sampling or a proxy model that optimizes the subset from the query. We will explain the technical details in Section 4.
3
RELATED WORK
Our work draws on and departs from three bodies of literature. We first review classical approximate query processing in Section 3.1, which provides the statistical foundation for our AQP formulation. We then discuss query optimization systems that handle ML predicates inside a database system in Section 3.2, which share our goal of reducing oracle invocations but assume the oracle is embedded 3
in the query engine. Finally, we review LLM pipeline optimization work in Section 3.3, which reduces the cost of AI workflows but operates at the pipeline level without exploiting query structure. The distinction that cuts across all three is the query-centric view: by treating AI agentic workflows as declarative queries, we expose optimization opportunities that are invisible from the pipeline view alone—namely, that aggregate functions, LIMIT clauses, and WHERE predicates constrain how many records need to reach the oracle at all. This is what enables our two strategies: an AQP formulation that terminates oracle calls early once the aggregate estimate converges, and a proxy model whose outputs are expressed as SQL predicates evaluated natively by the query engine.
3.1
runtime and dynamically reorder predicate evaluation, achieving up to 6.4× speedup over a static plan. Both categories therefore require either inserting a neural network into the query pipeline, or modifying the query engine to incorporate the ML oracle into the query plan. Our work identifies that AI agentic workflows can be expressed as declarative queries, and this observation opens a different opportunity: both our AQP formulation and proxy-based filtering can invoke an existing query engine directly without modifying it. The AQP formulation issues progressive sampling queries natively to the engine. The decisiontree proxy produces structural predicates expressible as SQL WHERE clauses, which the engine evaluates directly—the only external step is a one-time, lightweight training phase on a small sample of oracle-labeled records, after which the proxy runs entirely within the query execution framework.
Approximate Query Processing
Online aggregation [8] introduced the idea of computing aggregate results progressively as rows are scanned: after each batch of samples, the system updates the aggregate estimate and its confidence interval, and terminates early once the interval width falls within a user-specified error bound. The key guarantee is that the estimate is unbiased and the confidence interval is valid at every step. WanderJoin [16] extends this framework to queries with joins. The core difficulty is that uniform sampling from each table independently produces few join matches when join selectivity is low, making the estimator high-variance. WanderJoin resolves this by performing random walks over the join graph: each walk traces one path through the joined tables and produces one unbiased sample from the join result without materializing the full join. Our AQP formulation applies this framework to AI workflows, where the setting differs from standard AQP in one critical way: in classical AQP, each sampled row’s value is obtained by a cheap column read; in query-centric AI workflows, each row’s value requires an expensive oracle or reward-model call that dominates the total query cost. We therefore treat the number of oracle invocations, rather than rows scanned, as the optimization target, and terminate sampling as soon as the confidence interval converges to within the user-specified error bound.
3.2
3.3
LLM Pipeline Optimization
A parallel line of work optimizes LLM-powered workflows from the machine learning and systems perspective. These efforts fall into three categories: designing a declarative structure of LLM pipelines, routing queries to cheaper models, and reducing generation cost in Best-of-𝑁 sampling. First, LOTUS [20] introduces semantic operators, a declarative programming model that extends relational algebra with LLMbased operations such as semantic filtering, joining, and ranking over unstructured text. Abacus [21] builds a cost-based optimizer on top of LOTUS that selects among operator implementations under a cost budget. TAG [3] frames natural language questions as tableaugmented generation tasks, enabling accurate analytical query answering without Text-to-SQL. DocETL [22] takes an agentic approach to pipeline design, rewriting user-specified LLM pipelines into decomposed plans that are 25–80% more accurate than handengineered baselines on complex document processing tasks. These systems optimize what LLM operations to perform and how to structure them across a dataset. Our work is complementary: we take the query plan as fixed and ask how many records need to reach the oracle to answer it accurately. For optimization through model routing, FrugalGPT [4] routes each query to the cheapest LLM that can answer it accurately, achieving up to 98% cost reduction over always using GPT-4 with no accuracy loss on several benchmarks. Task Cascades [23] decomposes a task into a cascade of cheaper sub-operations, escalating only uncertain records to the expensive oracle; across eight document-processing tasks at a 90% target accuracy, it reduces endto-end cost by an average of 36% over a standard model-cascade baseline. Both approaches reduce cost by substituting a cheaper model for the expensive one, but still invoke a model on every record in the dataset. Our approach is orthogonal: rather than finding a cheaper model to call, we reduce the number of records that require any model invocation at all. Lastly, best-of-𝑁 sampling generates 𝑁 candidate responses per prompt and retains the top-scoring one, and is a standard component of LLM post-training pipelines [7, 26]. Speculative Rejection [24] reduces the cost of generating candidates: it stops generating tokens for a candidate early if an interim reward estimate falls below a threshold, requiring between 16 and 32 times fewer
Query Optimization with ML Predicates
Prior work on optimizing queries with expensive ML predicates falls into two categories. The first targets video analytics. NoScope [12], BlazeIt [11], and TASTI [14] accelerate video queries by inserting a lightweight neural network ahead of the expensive oracle model: the proxy scores each frame and filters out records unlikely to satisfy the predicate before they reach the full model. This approach requires training and running a separate neural network at query time, which itself incurs GPU inference on every record that reaches the proxy. The second category targets general ML predicates expressed as UDFs. PP [19] and CORE [27] rewrite the query plan to insert proxy filters ahead of costly ML UDFs inside the DBMS, requiring the oracle to be registered as part of the query execution engine. CORE relaxes PP’s independence assumption across predicates using branch-and-bound search, achieving up to 63% higher throughput than PP. ABae [13] embeds a trained neural proxy into the query plan to stratify records for budget allocation. Aero [9] goes further by modifying the query engine itself to monitor UDF statistics at 4
4.1.1 Illustration: AQP with SUM. For concreteness, we first address the following SUM query:
GPU resources than standard Best-of-𝑁 to achieve comparable reward. Kalayci et al. [10] connect Best-of-𝑁 to the Pandora’s Box optimal stopping problem and propose a UCB-style algorithm that reduces the number of candidates generated by 15–35%. Both Speculative Rejection and the Pandora’s Box formulation operate on the generation side. Our work instead optimizes the scoring side of best-of-𝑁 RLHF: given a fixed pool of 𝑁 candidates, we reduce the number of reward-model invocations needed to identify the top-𝑘 responses. The candidate-generation cost is fixed across all methods we compare and is orthogonal to our optimization; our scoring-side reduction is therefore complementary to—and can be combined with—generation-side methods such as Speculative Rejection [24] and the Pandora’s Box approach [10].
4
SELECT SUM(c1) WHERE c2 = "label 1"
Suppose the predicted “label 1” rows 𝛾 are sampled with the probability 𝑝 (𝛾), and the aggregated function on 𝛾 (i.e., SUM) is Í 𝜐 (𝛾). Then, 𝜐 (𝛾)/𝑝 (𝛾) is an unbiased estimator of 𝛾 𝜐 (𝛾), which is the SUM function we want to compute. As we obtain samples and compute 𝜐 (𝛾𝑖 )/𝑝𝑖 throughout the iterations, we take its average (still an unbiased estimator of SUM) with a reducing variance as more samples are collected. Moreover, with expectation and variance, we could use standard statistical formulas to calculate the confidence interval to decide whether we are satisfied with the query results and terminate the pipeline. In this way, we could reduce the number of rows to the AI workflow systematically.
QUERY-CENTRIC AI WORKFLOW INFERENCE
4.1.2 AGG Functions. AGG functions like SUM, AVG, VAR and COUNT are similar to the SUM example. They are all based on sample points to estimate their distribution. Let the observation value for these functions on the 𝑖-th trial be 𝑣𝑖 . Then the estimator Í of 𝑌ˆ with 𝑛 trials is: 𝑌ˆ = 𝑛1 𝑛𝑖=1 𝑣𝑖 . To use Equation 1 to decide when to stop sampling √︃ and return results, we estimate the sample 1 Í𝑛 ˆ 2 standard deviation: 𝑠 = 𝑛−1 𝑖=1 (𝑣 𝑖 − 𝑌 ) . With the estimator for 𝑌ˆ and the sample standard deviation 𝑠, the achieved half-width of the confidence interval in Equation 1 is 𝑍√𝛼𝑛𝑠 ; we stop sampling once this half-width stabilizes across rounds, and accept the estimate only if it has stabilized within the user-specified bound 𝜖. Other AGG functions, such as MIN and MAX, are difficult for the sampling method. They usually do not follow a certain distribution like Gaussian. Assuming these values depend on one unique row only, the expected number √︃ and standard deviation of sampling trials
The core idea of an inference optimization from a query-centric AI workflow is that we do not need to fetch all the data points and invoke the AI pipeline on each of them to obtain accurate results that require predictions. For example, if the query aims at collecting the AVG(c1) where the predicted values at c2 are equal to "label 1" while all values in c1 are equal to 0, we do not even need to invoke the AI pipeline on a single row to acquire 0 as the correct ML query result. On the other hand, if we do not aggregate the rows in the query result (e.g. SELECT), we still do not need to invoke the AI pipeline on much of the database if the query has a filter by prediction labels. Therefore, there exists an execution plan that decides how many AI-pipeline invocations we need, as well as how they should be constructed. In the following sections, we describe two main strategies to process the queries without invoking the AI pipeline on all data from the database, their tradeoffs, and how to interplay them together to optimize query performance.
4.1
needed will be 𝑁𝑏 and (1 − 𝑁𝑏 )/( 𝑁𝑏 ) 2 where 𝑁 is the number of rows and 𝑏 is the sample size. If 𝑏 ≪ 𝑁 , the 95% CI for the number of trials will be large. If 𝑏 is not much smaller than 𝑁 , sampling becomes comparable to a full table scan. In both cases, a full scan is preferable. To address the limitation on these AGG functions, we introduce the second strategy that we should construct queries to select points with target labels with small cost, in the next section.
Strategy AQP: Treating the Workflow as an Approximate Query Processing Problem
The first strategy treats these workflows as an Approximate Query Processing (AQP) problem and estimates aggregate results from samples. We use an online aggregation approach [8] to provide approximate results with confidence intervals and gradually improve with an increasing number of samples. Given an AGG(expression) and any samples where AGG could be SUM, AVG, COUNT or VAR and expression could involve any columns in the database, the ML query should provide an estimator 𝑌ˆ with a confidence interval: 𝑃𝑟 [|𝑌ˆ − AGG(expression)| ≤ 𝜖] ≥ 𝛼 .
4.1.3 LIMIT Queries. Beyond AGG functions, Strategy AQP extends naturally to LIMIT queries of the form: SELECT * ORDER BY score DESC LIMIT k
This pattern appears directly in LLM post-training pipelines: for each prompt, 𝑁 candidate responses are generated and scored by a reward model, and only the top-𝑘 responses are retained [7]. Rather than invoking the oracle on all 𝑁 candidates, we ask: how many candidates must we sample so that the global top-𝑘 are likely to be included? Stopping Criterion via Order Statistics.. Let the 𝑁 candidates have i.i.d. continuous reward scores drawn from an unknown distribution 𝐹 . By a standard order-statistics argument, the probability that all 𝑘 global top candidates are contained in a uniform random sample of 𝑚 out of 𝑁 is:
(1)
𝜖 is the user-defined half width of the confidence interval and 𝛼 is the confidence level. The algorithm that leverages online aggregations is as follows. Iteratively, we uniformly random-sample (i.e., bootstrap) some data points from the database with probability 𝑝, calculate 𝑌ˆ and the confidence-interval half-width from the samples, and continue to sample until the half-width stabilizes, i.e., it no longer decreases significantly across rounds. If the stabilized half-width is within the user-specified bound 𝜖, the estimate is accepted; if it stabilizes above 𝜖, Strategy AQP is not well suited to this query.
𝑃𝑘 (𝑚, 𝑁 ) =
𝑘 −1 Ö 𝑚 −𝑖 𝑖=0
5
𝑁 −𝑖
.
(2)
This result is distribution-free: it holds regardless of 𝐹 . A conservative distribution-free stopping rule would terminate sampling when 𝑃𝑘 (𝑚, 𝑁 ) ≥ 𝛼, where 𝛼 is the aforementioned confidence level in Equation 1. The special case 𝑘 = 1 recovers the simpler bound 𝑃1 (𝑚, 𝑁 ) = 𝑚/𝑁 ≥ 𝛼, i.e., sampling a fraction 𝛼 of candidates suffices to include the global best with probability 𝛼. For large 𝑁 , Equation 2 is well approximated by: 𝑚 ≈ 𝑁 · 𝛼 1/𝑘 ,
(3)
showing that larger 𝑘 requires sampling a larger fraction of candidates— since all 𝑘 of the global top scorers must be retained—which makes this worst-case bound conservative and motivates the adaptive criterion below. Adaptive Stopping in Practice.. While Equation 2 provides a distribution-free guarantee, it is a worst-case bound: it assumes the top-𝑘 candidates may appear anywhere in the population with equal probability, requiring a sample of size 𝑚 ≈ 𝑁 ·𝛼 1/𝑘 regardless of the actual score distribution. In practice, this can be overly conservative. Following standard practice in approximate query processing [8], we instead allow the user to specify a tolerance 𝛿 > 0, and terminate sampling when the worst score among the current top-𝑘 stabilizes: (𝑚) (𝑚−1) |𝑟 (𝑘 − 𝑟 (𝑘 | ≤ 𝛿, ) )
(4)
(𝑚) where 𝑟 (𝑘 denotes the 𝑘-th highest reward score observed after ) drawing 𝑚 (𝑚 ≥ 𝑘 + 1) candidates. This criterion exploits the actual score distribution rather than assuming the worst case, and therefore terminates earlier whenever high-score candidates are concentrated in a small fraction of the population. This adaptive stopping rule is empirical rather than distribution-free: it does not provide the worst-case guarantee in Equation 2, but in our workloads we find that once the 𝑘-th best observed score stabilizes, the global top-𝑘 candidates are captured with high probability. In our evaluation on LLM fine-tuning pipelines (Section 5.2), we find that this adaptive stopping criterion achieves a substantially better cost–accuracy trade-off than the theoretical bound, reducing oracle model calls by up to 80% while maintaining downstream quality within 5% of the full-oracle baseline.
4.2
Figure 3: Overview of Strategy PM. A one-time training phase (left) labels a small uniform sample with the reward model and trains a decision-tree proxy on surface features, calibrating a variance threshold 𝜎th and a score threshold 𝜏. At inference (right), the proxy scores all 𝑁 candidates per prompt; when calibrated as reliable, the variance gate skips uninformative prompts at zero oracle cost, and the candidate filter forwards only high-scoring survivors to the reward model for the final top-𝑘 selection. mimic the ML pipeline as efficient SQL to select data points with the predicted label “1” from the database. The routine is as follows: First, we uniformly randomly sample a few rows from the database, pass them to the AI workflow to obtain the labels, and train a lightweight, query-specific proxy model trained on-the-fly at query execution time. We use a cascade approach to determine to what extent the data needs to go through the AI pipeline based on the output characteristics (e.g. when the prediction confidence is low on the data).
Strategy PM: Proxy Model to Increase Throughput on Model Inference
The second strategy to avoid feeding all data to the AI workflow is to develop a predictive model that handles the “easy” input before passing the remaining data to the AI pipelines. We define easy data as rows whose pipeline outcome can be predicted with high confidence by a rule-based model, or whose prediction error does not affect the final query result. That being said, even if the outputs from the proxy model might have wide percentage errors, if they do not affect the top-k final results in the LIMIT clause for example, we can safely terminate the proxy model routine. Thus, the idea is to train an efficient proxy model with limited ground truth samples and process the majority of the data without affecting the quality of the final output, thus increasing the throughput of the query.
4.2.2 Overview of Proxy Model. To implement this strategy (Figure 3), we train a rule-based proxy model on-the-fly at query execution time. Given a query, an input relation, and a targeted AI pipeline, we first draw a small uniform random sample of records, pass them through the AI pipeline to collect ground-truth outputs (labels or scores), and use this labeled set to train a lightweight, query-specific surrogate. For simplicity, we use a Decision Tree (DT), though the routine applies equally to other rule-based models such as Bayesian rules or random forests. The key advantage of a DT is that training is fast and the resulting model is expressible as a set of SQL-compatible predicates, keeping the proxy natively within
4.2.1 Illustration: SELECT rows with label “1”. Let us consider a simple query: SELECT * FROM table where label = ’1’ The goal is to train a rule-based model as a proxy model which could 6
the query execution framework. For example, in an LLM rejection sampling pipeline, the sample is scored by a reward model (RM) and the DT is trained to predict RM scores from surface features of the candidate responses. In practice, we require a minimum of 5% inferred points from this pipeline.
predictions are weakly discriminative, for example in the general instruction-following domain, where free-form responses make surface features less informative and the continuous reward target requires the RM to perform the final ranking over a broad survivor set. Strategy B (gate + filter) activates both tests with an aggressive gate (𝑞𝜎 = 0.99, skipping the large majority of low-dispersion prompts) and a lower filter quantile (𝑞𝜏 = 0.5). This is used for strong generators on structured tasks (Math, Code), where the tree predicts RM scores well and most prompts are resolved without any RM call. The choice between the two is itself made automatically from the training sample using the mean training RM score 𝑠¯ as a proxy for generator strength: the general domain always uses A, while otherwise B is selected when 𝑠¯ exceeds a calibrated threshold (a strong generator whose answers are broadly high-scoring) and A otherwise. Tying the cascade’s aggressiveness to the observed difficulty of the input batch in this way, rather than to a fixed reduction rate, is what allows the same gate to yield up to 19× RM-call reduction on structured domains while remaining safe on harder, free-form ones.
4.2.3 Feature Engineering. Rather than relying on the pipeline’s internal representations, the proxy is trained on cheap, pipelineagnostic surface features extracted directly from each record — covering structural properties (e.g., formatting completeness, symbol density), lexical statistics (e.g., lexical diversity, compression ratio), and domain hints (e.g., presence of numeric results or structured syntax). These features compute in negligible time relative to a single pipeline call and require no model invocation. Crucially, they function as query-derived structural predicates: lightweight WHERE-clause filters that are invisible to inference-based surrogates operating on semantic content alone. For example, in fine-tuning LLMs, such features include response length, equation density, and answer markers—signals a lightweight proxy can use to discard weak candidates for specialized tasks such as mathematical reasoning or coding before invoking an expensive Reward Model.
4.2.5 Fallback or Early Termination. When the proxy assigns nearuniform scores but its calibrated reliability is low, the system falls back to invoking the full pipeline, preserving correctness. Conversely, when the proxy is both near-uniform and reliable on the calibration sample, or when all proxy scores are comfortably within the query’s acceptance region, the pipeline can be bypassed entirely. This ties the cascade’s aggressiveness to the actual difficulty of the input batch rather than a fixed reduction rate. For example, a reliably low-scoring batch signals that none of the candidates merit fine-tuning, and the round can be skipped.
4.2.4 Cascade Gating. We now expand the rejection-sampling cascade, which is the form used throughout our LLM evaluation (Section 5.2). The query is SELECT * ORDER BY rm_score DESC LIMIT k over the 𝑁 candidates of each prompt, and the proxy is a decision tree that predicts the RM score from surface features. Because each prompt is an independent query group of 𝑁 candidates, the cascade is applied per prompt and specializes into two coupled tests. Variance gate (group-level skip). The tree first scores all 𝑁 candidates of a prompt and we measure the within-prompt dispersion of these predictions, 𝜎pred = std(𝑠ˆ1, . . . , 𝑠ˆ𝑁 ). When 𝜎pred falls below a threshold 𝜎th , the candidates are deemed interchangeable, meaning no candidate is confidently better than the rest, so the entire prompt is skipped: we return a single candidate (chosen by a random number generator with a fixed seed for reproducibility) and issue zero RM calls for that prompt. This is the LLM-specific instance of the fallback rule discussed below. When the proxy assigns near-uniform values across a group, ranking is uninformative and oracle scoring is wasted. Candidate filter (within-group pruning). If a prompt survives the variance gate, we avoid scoring all 𝑁 candidates by pruning the obviously weak ones. We retain the survivor set 𝑆 = { 𝑐 : 𝑠ˆ(𝑐) ≥ 𝜏 }, where 𝜏 is a score threshold, forward only 𝑆 to the RM, and select arg max𝑐 ∈𝑆 over the true RM scores. If 𝑆 is empty, meaning no candidate clears 𝜏, we fall back to forwarding all 𝑁 candidates to preserve correctness. This costs |𝑆 | RM calls in place of 𝑁 . Both thresholds are calibrated on the training sample alone and depend only on the proxy’s predictions, never on oracle accuracy, so the gating is independent of ground truth. Concretely, 𝜎th is set to the 𝑞𝜎 -quantile of the per-prompt prediction dispersions over the training prompts, and 𝜏 to the 𝑞𝜏 -quantile of the tree’s predicted scores over all training candidates. Two cascade presets. The two tests above are exposed as two presets that are special cases of the same mechanism. Strategy A (filter-only) disables the variance gate (𝑞𝜎 = 0, so no prompt is skipped) and runs only the candidate filter at a moderately high quantile (𝑞𝜏 = 0.6). This conservative setting is used when proxy
4.3
Deciding Which Strategy to Use
The two strategies are complementary and the choice between them is guided by the query structure and the observed data distribution. Query structure. If the query contains no aggregation (e.g., a bare SELECT clause), Strategy AQP does not apply and we default to Strategy PM. If the query contains an aggregation, both strategies are candidates; since Strategy AQP is analytical and runs in 𝑂 (1) after sampling, it is always cheap to attempt and serves as the baseline. Confidence interval width. We rely on Strategy AQP as the primary estimator when its confidence interval is within the userspecified error bound 𝜖. When the CI is too wide — typically due to a skewed label distribution or high output variance — we switch to or augment with Strategy PM, which avoids the variance accumulation problem by routing hard records to the full pipeline directly. Minimum sample size. Both strategies require a minimum sample of 5–10% of the input relation to produce reliable estimates; this threshold is the empirically observed lower bound below which neither the CI nor the proxy model generalizes usefully. In practice, we draw this initial sample once and reuse it to initialize both strategies, amortizing the sampling cost across the two.
4.4
Multiple Query Functions
For queries that produce multiple independent outputs like SELECT SUM(c1), AVG(c2), ..., we can treat them separately, running the 7
same strategies in parallel. However, when a single query involves more than one aggregate function, as in the following example: SELECT SUM(c1) + SUM(c2) WHERE c3 = "label 1" There exists an opportunity to terminate some computations early if some functions’ results do not contribute significantly to the variance of the final result. Formally, let the query produce aggregate estimates 𝐴ˆ1, . . . , 𝐴ˆ𝑘 , and let the final output be a function 𝑌 = 𝑓 (𝐴1, . . . , 𝐴𝑘 ). Each 𝐴ˆ𝑖 is maintained with an error estimate (e.g., variance or confidence interval) from Strategy AQP. At runtime, we estimate the contribution of each aggregate 𝐴𝑖 to the uncertainty of 𝑌 . Specifically, we approximate: impact𝑖 =
are largely black-box: their internal predicates and data distributions are opaque, making controlled experimentation difficult. Second, TPC-DS is a widely adopted benchmark that covers representative SQL query patterns with diverse aggregate functions, providing a well-understood setting in which we can precisely control label distributions and isolate the effect of each strategy. Specifically, we select eight representative TPC-DS queries (Q7, 13, 50, 52, 53, 54, 55, 56) spanning three aggregate functions (SUM, AVG, and COUNT). For each query, we split the original WHERE predicate into two parts: the predicted predicate that we could predict using columns in the table and the other structural filters and joins. We construct a data table by joining the relevant fact and dimension tables according to the original query’s join conditions, materializing the filters explicitly. The predictive predicate is converted into a binary row label. We then evaluate a single scalar aggregate (SUM, AVG, or COUNT) over all rows satisfying the learned predicate, stripping the original query’s GROUP BY, ORDER BY, and LIMIT clauses to isolate the core problem—accurately estimating an aggregate under a predictive predicate with limited oracle calls that mimics an AI workflow. Since the value distributions of TPC-DS are referenced from real world data tables, applying the original WHERE predicate to dimension columns (e.g., i_category, i_class for Q53) to produce a natural label distribution limits our opportunities to stress test our strategies with a comprehensive set of ranges of label distributions. Thus, we also set the positive-label rate to 30% and 50% to simulate more balanced distributions that may arise in practice. We define the oracle budget as the number of calls to retrieve these labels from the column. For each strategy and label distribution, we vary the oracle budget from 0.1% to 20% of total rows, repeat each configuration 10 times, and report the mean aggregate error with confidence intervals. Both strategies use uniform random sampling to select the data for calling the oracle initially.
𝜕𝑓 · Δ𝐴𝑖 , 𝜕𝐴𝑖
where Δ𝐴𝑖 is the current error bound of 𝐴𝑖 . The term 𝜕𝐴𝑖 captures how sensitive the final result is to changes in 𝐴𝑖 . When 𝑓 is simple (e.g., sum, difference, ratio), this sensitivity can be computed analytically. For general expressions, including UDFs, we approximate it using finite differences by perturbing the aggregate value: 𝜕𝑓
𝜕𝑓 𝑓 (𝐴ˆ1, . . . , 𝐴ˆ𝑖 + ℎ𝑖 , . . . ) − 𝑓 (𝐴ˆ1, . . . , 𝐴ˆ𝑖 − ℎ𝑖 , . . . ) ≈ . 𝜕𝐴𝑖 2ℎ𝑖 Here, we choose ℎ𝑖 to be proportional to the current uncertainty of 𝐴𝑖 (i.e., ℎ𝑖 = Δ𝐴𝑖 ) so that the sensitivity estimate is evaluated at the scale of the current confidence interval. Using this estimate, the system can selectively stop refining aggregates whose contribution is small. In particular, we stop sampling for 𝐴𝑖 when: impact𝑖 ≤ 𝜖, for a user- or system-defined tolerance 𝜖. This strategy naturally prioritizes aggregates based on their influence on the final result. For example, in additive queries, aggregates with small variance quickly become irrelevant and can be terminated early. In contrast, in ratio queries, the denominator often dominates the error and continues to receive samples.
5
Results for Strategy AQP. We present the results of Strategy AQP under the settings described above in Figures 4 and 5. Figure 4 measures the real errors of the queries under different percentages of oracle calls and Figure 5 measures the estimated CI with increasing numbers of sampling rounds. For Figure 4, the dots along the line represent the increasing number of sampling trials (e.g., the 4th dot means we have done four trials of sampling). Each trial bootstraps 1% of the data so that the total of oracle calls might be less than the total number of samples due to overlapping. Overall, the results show that if the label distributions are uniform (p=30%, 50%), the error of the AGG results drops below 10% with less than 10% of oracle calls. For real distributions based on the data characteristics in TPC-DS, larger tables still achieve similar results, whereas small tables with fewer than 10,000 rows (e.g., Q13, 54) might require more than 20% of oracle calls. We hypothesize that our sampling strategy follows the Law of Large Numbers; for small tables, running the original AI workflow may be preferable. Figure 5 shows that all queries reach their adaptive stopping point at around 10–13 rounds of sampling, an 85–90% reduction in oracle calls. Under balanced distributions the confidence interval (Section 4.1.2) has narrowed below 10% by then, whereas under the natural distribution the small-table cases stop with the interval still wide, indicating that Strategy AQP alone may be unreliable and
EVALUATION
In this section, we evaluate our query-centric optimization on two kinds of AI workflows. Our goals are: • Using AI workflows derived from TPC-DS queries, we demonstrate that under balanced distributions AGG results can be approximated within 10% error while the adaptive stopping point is reached at 10–15% of tuples—an 85–90% reduction in oracle calls. • Using LLM post-training tasks with Reward Models as an example, we reduce reward-model scoring time by up to 19× while staying within 10% of full-RM accuracy on structured Math and Code tasks and within ∼15% of the full-RM reward score on the open-ended General task.
5.1
Evaluating AGG Queries with TPC-DS
Experimental Setup. We choose TPC-DS as the evaluation workload for two reasons. First, real-world AI-agentic query workflows 8