andygrove opened a new pull request, #5204: URL: https://github.com/apache/datafusion-comet/pull/5204
## Which issue does this PR close? Closes https://github.com/apache/datafusion-comet/issues/5200. Part of the planner performance audit epic https://github.com/apache/datafusion-comet/issues/5199. ## Rationale for this change `QueryPlanSerde` attaches a `QueryContext` to **every** expression it converts, and `QueryContext.sql_text` is the *full* SQL text of the query. The same string therefore ends up embedded once per expression in the serialized native plan. On real TPC-DS queries the query text is 85-95% of the plan bytes. The context itself is needed — `QueryContext::format_summary` prints the whole `sql_text` to build Spark's `== SQL (line N, position M) ==` error block, so it cannot be trimmed to a fragment. The *duplication* is the problem. This costs more than just bytes: 1. **Plan size.** The `Array[Byte]` is captured by `CometExecRDD` and shipped in the stage's task binary. 2. **Per-task JVM protobuf round trip.** Whenever a plan contains a native scan — essentially every scan-rooted Comet stage — `CometExecRDD.compute` does `Operator.parseFrom(serializedPlan)` → `PlanDataInjector.injectPlanData` → `serializeOperator`, **per task**. 3. **Per-task native decode.** `PhysicalPlanner::create_expr` / `create_agg_expr` did `ctx_proto.sql_text.clone()` into a fresh `String` per expression and registered it in `QueryContextMap`. For q64 below that is 867 identical heap copies of a 3 KB query, per task, held for the lifetime of the query. ## Benchmarks Real TPC-DS queries (v1.4 SQL from the Spark test resources) planned against empty TPC-DS tables — planning only, no data needed. "before" is the exact byte count obtained by inlining the pooled text back onto every context, so it is the real pre-patch serialization, not an estimate. "round trip" is the parse + re-serialize that `CometExecRDD.compute` performs per task. | query | SQL chars | `Expr` nodes | with context | plan bytes before | plan bytes after | reduction | per-task round trip before | after | | --- | --- | --- | --- | --- | --- | --- | --- | --- | | q23a | 1654 | 178 | 168 | 284 KB | 32 KB | 8.9x | 0.84 ms | 0.37 ms | | q23b | 2009 | 240 | 221 | 441 KB | 47 KB | 9.3x | 0.75 ms | 0.32 ms | | q64 | 3051 | 1112 | 867 | 2,360 KB | 109 KB | 21.5x | 3.02 ms | 0.73 ms | | q67 | 884 | 258 | 249 | 247 KB | 24 KB | 10.3x | 0.41 ms | 0.20 ms | | q72 | 1261 | 282 | 263 | 381 KB | 38 KB | 9.8x | 0.59 ms | 0.23 ms | | q95 | 976 | 168 | 137 | 153 KB | 24 KB | 6.2x | 0.31 ms | 0.17 ms | | q14a | 3716 | 513 | 393 | 1,446 KB | 95 KB | 15.2x | 1.87 ms | 0.47 ms | | q47 | 1731 | 306 | 248 | 433 KB | 29 KB | 14.5x | 0.65 ms | 0.34 ms | | **total** | | | | **5,749 KB** | **401 KB** | **14.3x** | | | The worst case, q64, went from a **2.4 MB** serialized plan to 109 KB. Its per-task parse + re-serialize drops from 3.0 ms to 0.7 ms; at 10k tasks that is ~23 s of JVM protobuf work per stage removed, before counting the native-side decode and the 867 string clones per task. Environment: Spark 4.1 / Scala 2.13, JDK 17, macOS arm64, debug native build. Measured with a throwaway ScalaTest suite (not committed — a proper planner benchmark is tracked in the epic); the committed `QueryContextInternerSuite` asserts the size reduction so the property does not silently regress. ## What changes are included in this PR? Intern the SQL texts into a pool instead of repeating them: - **proto**: `Operator` gains `repeated string sql_text_pool`, populated only on the root operator of a serialized native block. `QueryContext` gains `optional int32 sql_text_idx`; `sql_text` is cleared when the index is set. - **`QueryContextInterner`** (new): rewrites an operator tree so each distinct SQL text appears once in the root's pool. Applied in `CometNativeExec.convertBlock()`, which is where the whole block is in hand and where the pool indices are scoped. Doing it as a post-pass means none of the ~90 `QueryPlanSerde` call sites change and no plan-scoped mutable state has to be threaded through `exprToProto`. - **native**: `PhysicalPlanner` loads the pool from the root operator (`with_sql_text_pool`) and `build_query_context` resolves `sql_text_idx` against it, handing every context a clone of one shared `Arc<String>`. `QueryContext::new` now takes `impl Into<Arc<String>>` so the `Arc` is shared rather than re-allocated per expression — the struct already stored `Arc<String>`, that sharing was just never realised. - `sql_text` remains a working fallback when `sql_text_idx` is absent, so the paths that serialize an operator directly without going through `convertBlock` (e.g. `CometNativeWriteExec`) are unaffected. - **`sourceKey` derivation**: `CometNativeScanExec.apply` hashes a scan's `NativeScanCommon` on the driver to produce the key that `NativeScanPlanDataInjector.getKey` recomputes on the executor from the *serialized* plan. Since only the executor's copy is interned, that hash no longer matched and the scan failed with `Missing planning data for key: ...`. Both sides now hash the context-free form via `QueryContextInterner.stripQueryContexts`, which also makes the key robust to any future change in how contexts are encoded. The interner walk is descriptor-driven rather than a hand-written recursion over `Expr`'s ~75-way `oneof`: a `QueryContext` can sit on any `Expr` or `AggExpr` at any nesting depth, and a hand-written walk would silently miss newly added variants. `QueryContextInternerSuite` asserts that no un-pooled context survives in a real plan, which is what keeps the walk honest. ## How are these changes tested? New `QueryContextInternerSuite`: - every `QueryContext` in a real plan carries `sql_text_idx` and no inline `sql_text`, and the pool holds the text once; - the plan is materially smaller than the inlined equivalent; - an ANSI `DIVIDE_BY_ZERO` error still reports the full SQL text end-to-end, i.e. the pooled text round-trips JVM → proto → native → error JSON → Spark exception unchanged. Regressions run locally on Spark 4.1 / Scala 2.13: `CometExpressionSuite`, `SparkErrorConverterSuite`, `CometExecSuite`, `CometAggregateSuite`, `CometCastSuite`, `CometJoinSuite`, `CometNativeReaderSuite`, `ParquetReadSuite`, and the full `cargo test` suite. Error messages are unchanged; this is purely a wire-format deduplication. ## Additional context The per-expression `QueryContext` was introduced in https://github.com/apache/datafusion-comet/pull/3580 ("feat: [ANSI] Ansi sql error messages"), which added the SQL context needed for Spark-compatible ANSI error messages. That behaviour is preserved here — only the encoding changes. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
