[FEATURE]Support a configuration-driven query optimization runtime for Kyuubi Spark SQL extensions - For ai

[jueshali]

Code of Conduct

• I agree to follow this project's Code of Conduct

Search before asking

• I have searched in the issues and found no similar issues.

Describe the feature

Kyuubi already provides several useful Spark SQL auxiliary optimization rules, such as repartition/rebalance before writing, forcing shuffle before join for AQE skew optimization, final-stage config isolation, final write stage resource isolation, and SQL classification driven dynamic behavior.
These capabilities show that Kyuubi Spark extensions can already apply execution-side optimizations through Spark parser / analyzer / optimizer / planner / query-stage rules.
I would like to propose a more general capability on top of these existing extensions:
Support a configuration-driven query optimization runtime for Kyuubi Spark SQL extensions.
The idea is that Kyuubi should be able to consume a query-level optimization plan passed from upstream systems through configuration, and apply that plan using existing Spark SQL extension points.
This is not about moving optimization decision logic into Kyuubi.
Instead, it is about making Kyuubi a better execution-side runtime for upstream optimization systems.
A query-level optimization plan may include things such as:
write optimization strategy
repartition / coalesce guidance
AQE-related overrides
join strategy bias
resource isolation directives
trace / dual-run identifiers
This becomes increasingly useful in the AI-assisted SQL optimization era, where optimization suggestions may be generated by upstream rule engines, historical optimizers, or AI systems, but still need a standard and observable execution-side runtime to apply them safely.

Motivation

In many real-world deployments, query optimization decisions are often made by an upstream gateway or query platform, for example based on:
workload classification
historical workload analysis
workload-specific resource sizing
heavy query diagnosis
AI-assisted SQL optimization
dual-run validation before rollout
In such an architecture, the upstream system should remain the decision maker, while Kyuubi acts as the execution-side runtime.
Today, Kyuubi has multiple independent optimization rules, but there is no unified runtime abstraction for consuming a query-level optimization plan from upstream systems.
As a result:
upstream systems can only influence execution through scattered Spark confs
there is no common way to express optimization intent across different optimization rules
there is no unified execution trace showing:what optimization plan was received
which directives were applied
where they were applied
whether they were fully or partially applied

This is especially important for AI-assisted SQL optimization workflows, because AI-generated optimization suggestions need:
a standard runtime contract
execution-side observability
safe rollout support
dual-run / validation traceability
So the motivation is not only feature completeness, but also making Kyuubi more suitable as a general execution runtime for upstream optimization systems, including AI-assisted ones.

Describe the solution

A possible implementation could introduce a lightweight internal optimization runtime abstraction, for example:

1. Unified optimization plan abstraction
   An internal OptimizationPlan could represent things such as:
   optimization enabled / disabled
   write optimization strategy
   repartition / coalesce guidance
   AQE overrides
   join strategy bias
   resource isolation directives
   trace id / source / dual-run id
   This plan should be resolved from a well-defined set of Spark/Kyuubi configuration keys.
2. Optimization plan resolver
   A common resolver could parse raw conf into a structured OptimizationPlan, so that each rule does not parse raw string confs independently.
3. Layered runtime application
   Different optimization categories could be applied at different layers:
   Session / query conf only
   AQE toggles
   shuffle partitions
   advisory partition size
   skew join switches
   Post-hoc resolution / logical rewrite
   write-side repartition / rebalance
   final-stage tagging
   logical hint-like guidance
   Planner / query stage prep
   resource profile injection
   final-stage resource isolation
   physical strategy bias
4. Execution trace
   Kyuubi could expose execution-side trace information such as:
   what optimization plan was received
   which directives were consumed
   which rule/layer applied them
   whether they were fully or partially applied
   I think this should be built on top of existing Spark SQL extension points already used by Kyuubi, rather than introducing a completely separate mechanism.

Additional context

No response

Are you willing to submit PR?

• Yes. I would be willing to submit a PR with guidance from the Kyuubi community to improve.
• No. I cannot submit a PR at this time.

[github-actions[bot]]

Hello @jueshali,
Thanks for finding the time to report the issue!
We really appreciate the community's efforts to improve Apache Kyuubi.

[wangzhigang1999]

Thanks for the proposal. This is a broad architectural idea rather than an actionable feature issue at this stage.

I'm converting it to a Discussion so we can collect concrete use cases and narrow down the scope first.

[jueshali]

Example

Original SQL

INSERT OVERWRITE TABLE dws_user_pay_day PARTITION(dt='20260714')
SELECT
  a.user_id,
  a.game,
  SUM(b.pay_amount) AS total_pay
FROM ods_user_login_day a
JOIN dwd_pay_detail_day b
  ON a.user_id = b.user_id
WHERE a.dt = '20260714'
  AND b.dt = '20260714'
GROUP BY a.user_id, a.game

Upstream analysis

Historical statistics show:

• duration: 1200s
• shuffle: 95GB
• spill disk: 42GB
• diagnoses: TASK_SKEW, BIG_SPILL_DISK

AI suggests:

• rebalance before write
• enable skew join
• tune shuffle partitions
• isolate final write stage resources
• send to dual-run validation

Optimization intent

{
  "version": 1,
  "source": "AI",
  "traceId": "trace-20260714-001",
  "dualRunId": "dual-20260714-abc",
  "writeStrategy": {
    "mode": "rebalance"
  },
  "runtimeStrategy": {
    "shufflePartitions": 128,
    "advisoryPartitionSizeInBytes": "64MB",
    "skewJoinEnabled": true
  },
  "resourceStrategy": {
    "finalStageResourceIsolationEnabled": true,
    "finalStageExecutorCores": 2,
    "finalStageExecutorMemory": "4g"
  }
}

What Kyuubi does

1. Parse the optimization plan into an internal OptimizationPlan
2. Apply runtime configs:
   • spark.sql.shuffle.partitions=128
   • spark.sql.adaptive.advisoryPartitionSizeInBytes=64MB
   • spark.sql.adaptive.skewJoin.enabled=true
3. Apply write optimization:
   • rebalance before write
4. Apply final-stage optimization:
   • final-stage resource isolation
   • custom resource profile

Execution trace

{
  "traceId": "trace-20260714-001",
  "dualRunId": "dual-20260714-abc",
  "source": "AI",
  "appliedActions": [
    {"layer": "session-conf", "action": "set shuffle partitions to 128"},
    {"layer": "session-conf", "action": "enable skew join"},
    {"layer": "post-hoc-resolution", "action": "apply write rebalance"},
    {"layer": "query-stage-prep", "action": "enable final-stage resource isolation"}
  ],
  "partialApplied": false
}

Dual-run result

{
  "dualRunId": "dual-20260714-abc",
  "resultValidation": "PASS",
  "durationImprovement": "18%",
  "shuffleReduction": "12%",
  "spillDiskReduction": "35%",
  "recommendation": "PROMOTE"
}

Summary

• historical statistics + AI generate optimization intent
• Kyuubi executes that intent
• trace records what was applied
• dual-run validates correctness and benefit