pith. sign in

arxiv: 2605.17831 · v1 · pith:7U7QM4XMnew · submitted 2026-05-18 · 💻 cs.LG · cs.DB

Agentic Cost-Aware Query Planning with Knowledge Distillation for Big Data Analytics

Mahdi Naser-Moghadasi This is my paper

Pith reviewed 2026-05-20 12:46 UTC · model grok-4.3

classification 💻 cs.LG cs.DB
keywords query optimizationknowledge distillationbig data analyticscost-aware planningbandit explorationresource constraintsSQL planning
0
0 comments X

The pith

An agentic planner with UCB1 exploration and knowledge distillation cuts big data query latency by 23 percent.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper presents a query planning method that pairs a rule-based teacher using six optimization strategies with UCB1 bandit search and a Random Forest cost model to select plans that respect explicit memory and latency limits. The teacher decisions are then distilled into a lightweight student model such as logistic regression or gradient boosting, which replicates the choices during fast inference. This setup addresses the problem that conventional optimizers become too slow or violate constraints in resource-limited big data environments. Evaluation on the NYC Taxi and IMDB datasets shows the combined approach delivers lower latency than default planners while preserving high constraint satisfaction, and the student model runs substantially quicker than the full teacher-bandit loop.

Core claim

The central claim is that a rule-based teacher planner augmented by UCB1 exploration and a learned cost model can generate resource-aware SQL plans, and that distilling those decisions into a simple student classifier produces near-equivalent plans at much higher speed. The teacher applies six fixed strategies, uses bandit search to balance exploration under constraints, and relies on Random Forest predictions of latency from plan features. The distilled student then mimics the teacher's selections for immediate use.

What carries the argument

The agentic query planning pipeline that couples a rule-based teacher with UCB1 bandit exploration, Random Forest latency prediction, and knowledge distillation to a lightweight student model.

If this is right

  • Resource-constrained machines can run big-data analytics with measurably shorter query times than default planners allow.
  • A distilled student model can reproduce teacher-bandit decisions at 15 times the inference speed while retaining 89 percent accuracy.
  • A single-file implementation makes the full planning stack reproducible on ordinary hardware without external dependencies.
  • Explicit resource constraints can be satisfied at 94 percent rate while still achieving 23 percent lower average latency.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • The same distillation pattern could be applied to other planning domains that currently rely on expensive search or simulation.
  • Replacing the Random Forest with a cheaper surrogate inside the teacher loop might further reduce the initial training cost.
  • The approach suggests a route to embedding learned planners directly inside database engines that must operate under strict memory caps.

Load-bearing premise

The Random Forest cost model supplies accurate enough latency predictions from plan features and the six strategies together with UCB1 search cover the plans that matter under the given resource limits.

What would settle it

A test on additional datasets where actual measured latencies deviate markedly from the Random Forest predictions, resulting in no latency improvement or constraint satisfaction below 94 percent.

Figures

Figures reproduced from arXiv: 2605.17831 by Mahdi Naser-Moghadasi.

Figure 1
Figure 1. Figure 1: System Architecture: Agentic query optimization framework combin [PITH_FULL_IMAGE:figures/full_fig_p002_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: illustrates performance breakdown by dataset, showing consistent improvements across both NYC Taxi and IMDB workloads. D. Constraint Satisfaction Analysis Table IV demonstrates the effectiveness of our constraint￾aware approach in maintaining feasible execution under re￾source limits. NYC Taxi IMDB 250 300 350 Dataset Median Latency (ms) DuckDB Default Teacher Only Teacher+Bandit+Cost Student (HGBC) [PITH… view at source ↗
Figure 3
Figure 3. Figure 3: Cost model calibration plot (MAE: 18.4ms, [PITH_FULL_IMAGE:figures/full_fig_p004_3.png] view at source ↗
Figure 4
Figure 4. Figure 4: Ablation study: component contributions to latency improvement [PITH_FULL_IMAGE:figures/full_fig_p005_4.png] view at source ↗
read the original abstract

Query optimization in big data analytics remains computationally expensive, particularly for resource-constrained environments where traditional optimizers fail to satisfy memory and latency constraints. We present an agentic query planning system that combines a rule-based teacher planner, UCB1 bandit exploration, cost-aware prediction, and knowledge distillation to a lightweight student planner. Our teacher planner generates SQL plans using six key optimization strategies, while UCB1 bandit search efficiently explores the plan space under explicit resource constraints. A Random Forest cost model predicts query latency from plan features, enabling cost-aware decisions. A distilled student planner (Logistic Regression or Gradient Boosting) learns to mimic teacher-bandit decisions for fast inference. Evaluation on NYC Taxi and IMDB datasets demonstrates 23% latency reduction compared to default planners while maintaining 94% constraint satisfaction. The student planner achieves 89% accuracy in replicating optimal plans with 15x faster inference time. Our single-file implementation enables reproducible big-data analytics on resource-limited machines and is publicly available at https://github.com/mahdinaser/agentic-kd-planner.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

2 major / 1 minor

Summary. The paper introduces an agentic system for query planning in big data analytics. It combines a rule-based teacher planner using six optimization strategies and UCB1 bandit exploration with a Random Forest cost model for latency prediction. Knowledge distillation is used to train a lightweight student planner (Logistic Regression or Gradient Boosting) to mimic the teacher's decisions. On NYC Taxi and IMDB datasets, it reports 23% latency reduction vs default planners, 94% constraint satisfaction, 89% accuracy for the student, and 15x faster inference.

Significance. This work could be significant for practical deployment in resource-limited settings if the empirical results are robust, as it addresses the computational expense of traditional query optimizers by using bandit exploration and distillation. The open-source implementation is a positive aspect for reproducibility.

major comments (2)
  1. [Evaluation section] Evaluation section: The reported performance numbers (23% latency reduction and 94% constraint satisfaction) are presented without details on the specific baseline planners, data splits, number of trials, or statistical measures like standard deviation or error bars. This undermines the ability to assess the reliability of the central empirical claims.
  2. [Cost model description] Cost model description: The Random Forest is described as predicting query latency from plan features to enable cost-aware UCB1 decisions, but no validation metrics for its predictive accuracy (such as R², MAE, or hold-out validation against measured runtimes) are provided. Given that the latency reduction relies on these predictions steering the search, this is a load-bearing omission for the headline results.
minor comments (1)
  1. [Abstract] Abstract: The six key optimization strategies are mentioned but not listed or briefly described; including a short enumeration would improve clarity for readers.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for the constructive comments on our manuscript. The feedback has identified important gaps in the presentation of our experimental results and cost model validation. We address each major comment below and have updated the manuscript to improve transparency and reproducibility.

read point-by-point responses

  1. Referee: [Evaluation section] Evaluation section: The reported performance numbers (23% latency reduction and 94% constraint satisfaction) are presented without details on the specific baseline planners, data splits, number of trials, or statistical measures like standard deviation or error bars. This undermines the ability to assess the reliability of the central empirical claims.

    Authors: We agree that these details are essential for assessing reliability. In the revised manuscript we have expanded the Evaluation section to explicitly name the baseline planners (Spark SQL default Catalyst optimizer and a non-bandit rule-based planner), describe the data splits (80/20 train/validation on the query logs from each dataset), state the number of independent trials (five runs per dataset with different random seeds), and report all key metrics as means accompanied by standard deviations together with error bars on the corresponding figures. revision: yes

  2. Referee: [Cost model description] Cost model description: The Random Forest is described as predicting query latency from plan features to enable cost-aware UCB1 decisions, but no validation metrics for its predictive accuracy (such as R², MAE, or hold-out validation against measured runtimes) are provided. Given that the latency reduction relies on these predictions steering the search, this is a load-bearing omission for the headline results.

    Authors: We acknowledge the omission of quantitative validation for the Random Forest cost model. We have added a dedicated paragraph and table in the revised manuscript reporting an R² of 0.86, MAE of 14.2 ms, and RMSE of 19.7 ms on a 20% hold-out set of measured runtimes. We also include a scatter plot of predicted versus actual latencies to demonstrate that the model is sufficiently accurate to support cost-aware UCB1 decisions. revision: yes

Circularity Check

0 steps flagged

No circularity: headline claims are direct empirical measurements on external datasets

full rationale

The paper describes a composite system (rule-based teacher with six strategies + UCB1 + Random Forest cost model + distilled student) and reports measured outcomes (23% latency reduction, 94% constraint satisfaction, 89% student accuracy, 15x speedup) from runs on NYC Taxi and IMDB workloads. No equations, derivations, or self-citations are present in the provided text that would make any reported quantity equivalent to a fitted parameter or prior self-result by construction. The Random Forest is used as an internal component whose accuracy is not claimed to be proven by the final numbers; the final numbers are external benchmark results. This is the normal case of a self-contained empirical paper.

Axiom & Free-Parameter Ledger

2 free parameters · 1 axioms · 0 invented entities

The approach relies on standard ML models and bandit algorithms whose hyperparameters are fitted to data; the teacher planner assumes its six strategies suffice for the target workloads.

free parameters (2)
  • UCB1 exploration constant
    Bandit parameter chosen to balance exploration versus exploitation during plan search.
  • Random Forest hyperparameters
    Tuned parameters of the latency prediction model.
axioms (1)
  • domain assumption The six key optimization strategies plus UCB1 search cover the relevant plan space for the evaluated datasets.
    Invoked when the teacher planner generates candidate plans.

pith-pipeline@v0.9.0 · 5713 in / 1250 out tokens · 39724 ms · 2026-05-20T12:46:28.188323+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Lean theorems connected to this paper

Citations machine-checked in the Pith Canon. Every link opens the source theorem in the public Lean library.

What do these tags mean?
matches
The paper's claim is directly supported by a theorem in the formal canon.
supports
The theorem supports part of the paper's argument, but the paper may add assumptions or extra steps.
extends
The paper goes beyond the formal theorem; the theorem is a base layer rather than the whole result.
uses
The paper appears to rely on the theorem as machinery.
contradicts
The paper's claim conflicts with a theorem or certificate in the canon.
unclear
Pith found a possible connection, but the passage is too broad, indirect, or ambiguous to say the theorem truly supports the claim.

Reference graph

Works this paper leans on

26 extracted references · 26 canonical work pages · 1 internal anchor

  1. [1]

    P. G. Selinger et al., ”Access path selection in a relational database management system,” inProc. ACM SIGMOD, 1979, pp. 23-34

    work page 1979

  2. [2]

    PostgreSQL Global Development Group, ”PostgreSQL Documentation,” https://www.postgresql.org/docs/, 2024

    work page 2024

  3. [3]

    Avnur and J

    R. Avnur and J. M. Hellerstein, ”Eddies: continuously adaptive query processing,” inProc. ACM SIGMOD, 2000, pp. 261-272

    work page 2000

  4. [4]

    Marcus et al., ”Neo: A learned query optimizer,”Proc

    R. Marcus et al., ”Neo: A learned query optimizer,”Proc. VLDB Endow., vol. 12, no. 11, pp. 1705-1718, 2019

    work page 2019

  5. [5]

    Marcus et al., ”Bao: Making learned query optimization practical,” inProc

    R. Marcus et al., ”Bao: Making learned query optimization practical,” inProc. ACM SIGMOD, 2021, pp. 1275-1288

    work page 2021

  6. [6]

    Yang et al., ”Balsa: Learning a query optimizer without expert demonstrations,” inProc

    Z. Yang et al., ”Balsa: Learning a query optimizer without expert demonstrations,” inProc. ACM SIGMOD, 2022, pp. 931-944

    work page 2022

  7. [7]

    Marcus and O

    R. Marcus and O. Papaemmanouil, ”Plan-structured deep neural network models for query performance prediction,”Proc. VLDB Endow., vol. 12, no. 11, pp. 1733-1746, 2019

    work page 2019

  8. [8]

    Distilling the Knowledge in a Neural Network

    G. Hinton et al., ”Distilling the knowledge in a neural network,”arXiv preprint arXiv:1503.02531, 2015

    work page internal anchor Pith review Pith/arXiv arXiv 2015

  9. [9]

    Chen et al., ”Heterogeneous knowledge distillation for simultaneous infrared-visible image fusion and super-resolution,”IEEE Trans

    L. Chen et al., ”Heterogeneous knowledge distillation for simultaneous infrared-visible image fusion and super-resolution,”IEEE Trans. Multi- media, vol. 24, pp. 3123-3136, 2022

    work page 2022

  10. [10]

    Zhao et al., ”Implicit compatibility in heterogeneous teacher-student knowledge distillation,”Pattern Recognition, vol

    H. Zhao et al., ”Implicit compatibility in heterogeneous teacher-student knowledge distillation,”Pattern Recognition, vol. 135, p. 109181, 2023

    work page 2023

  11. [11]

    Williams et al., ”Theoretical foundations of knowledge distillation for sequential decision making,” inProc

    K. Williams et al., ”Theoretical foundations of knowledge distillation for sequential decision making,” inProc. ICML, 2023, pp. 37241-37262

    work page 2023

  12. [12]

    Thompson et al., ”Pairwise knowledge distillation with basic evalu- ation metrics,”J

    M. Thompson et al., ”Pairwise knowledge distillation with basic evalu- ation metrics,”J. Machine Learning Research, vol. 25, no. 45, pp. 1-28, 2024

    work page 2024

  13. [13]

    Zhang et al., ”Learned database system optimization through knowl- edge distillation,” inProc

    X. Zhang et al., ”Learned database system optimization through knowl- edge distillation,” inProc. VLDB, 2024, pp. 1842-1855

    work page 2024

  14. [14]

    Kumar et al., ”Distilled adaptive query processing for real-time analytics,”IEEE Trans

    A. Kumar et al., ”Distilled adaptive query processing for real-time analytics,”IEEE Trans. Knowledge Data Eng., vol. 36, no. 8, pp. 3421- 3435, 2024

    work page 2024

  15. [15]

    Wu et al., ”Neural query optimization with transformer architectures,” inProc

    Y . Wu et al., ”Neural query optimization with transformer architectures,” inProc. ACM SIGMOD, 2024, pp. 567-580

    work page 2024

  16. [16]

    Patel et al., ”Adaptive learned query optimization for cloud databases,”Proc

    R. Patel et al., ”Adaptive learned query optimization for cloud databases,”Proc. VLDB Endow., vol. 17, no. 6, pp. 1123-1136, 2024

    work page 2024

  17. [17]

    Rodriguez et al., ”Multi-armed bandit approaches for autonomous database configuration,” inProc

    C. Rodriguez et al., ”Multi-armed bandit approaches for autonomous database configuration,” inProc. IEEE ICDE, 2024, pp. 891-904

    work page 2024

  18. [18]

    Chen et al., ”Transformer-based query optimization with attention mechanisms,”ACM Trans

    L. Chen et al., ”Transformer-based query optimization with attention mechanisms,”ACM Trans. Database Syst., vol. 49, no. 2, pp. 1-28, 2024

    work page 2024

  19. [19]

    Li et al., ”Foundation models for database query optimization,” in Proc

    H. Li et al., ”Foundation models for database query optimization,” in Proc. NeurIPS, 2024, pp. 12456-12469

    work page 2024

  20. [20]

    Wang et al., ”Deep learning for cardinality estimation in modern database systems,”VLDB J., vol

    M. Wang et al., ”Deep learning for cardinality estimation in modern database systems,”VLDB J., vol. 33, no. 4, pp. 789-806, 2024

    work page 2024

  21. [21]

    Smith et al., ”Runtime adaptive query processing with reinforcement learning,” inProc

    J. Smith et al., ”Runtime adaptive query processing with reinforcement learning,” inProc. ACM SIGMOD, 2024, pp. 234-247

    work page 2024

  22. [22]

    Johnson et al., ”Cloud-native query optimization for distributed analytics,”IEEE Trans

    K. Johnson et al., ”Cloud-native query optimization for distributed analytics,”IEEE Trans. Cloud Computing, vol. 12, no. 3, pp. 456-470, 2024

    work page 2024

  23. [23]

    Liu et al., ”Query optimization for vector databases in AI applica- tions,” inProc

    S. Liu et al., ”Query optimization for vector databases in AI applica- tions,” inProc. VLDB, 2024, pp. 2134-2147

    work page 2024

  24. [24]

    Martinez et al., ”Adaptive query planning for large-scale graph databases,”ACM Trans

    D. Martinez et al., ”Adaptive query planning for large-scale graph databases,”ACM Trans. Graph Data, vol. 2, no. 1, pp. 15-32, 2024

    work page 2024

  25. [25]

    Leis et al., ”How good are query optimizers, really?”Proc

    V . Leis et al., ”How good are query optimizers, really?”Proc. VLDB Endow., vol. 9, no. 3, pp. 204-215, 2015

    work page 2015

  26. [26]

    Auer et al., ”Finite-time analysis of the multiarmed bandit problem,” Machine Learning, vol

    P. Auer et al., ”Finite-time analysis of the multiarmed bandit problem,” Machine Learning, vol. 47, no. 2-3, pp. 235-256, 2002

    work page 2002