This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 3df58d5280 perf: intern QueryContext SQL text into a per-plan pool (up
to 20x smaller serialized plans for TPC-DS) (#5204)
3df58d5280 is described below
commit 3df58d52808fc125b2f6d132a712ed9745b9f022
Author: Andy Grove <[email protected]>
AuthorDate: Sun Aug 2 07:32:29 2026 -0600
perf: intern QueryContext SQL text into a per-plan pool (up to 20x smaller
serialized plans for TPC-DS) (#5204)
---
.github/workflows/pr_build_linux.yml | 1 +
.github/workflows/pr_build_macos.yml | 1 +
native/common/src/query_context.rs | 7 +-
native/core/src/execution/jni_api.rs | 1 +
native/core/src/execution/planner.rs | 128 +++++++++++-----
native/proto/src/proto/expr.proto | 11 +-
native/proto/src/proto/operator.proto | 10 ++
.../apache/comet/serde/QueryContextInterner.scala | 137 +++++++++++++++++
.../spark/sql/comet/CometNativeScanExec.scala | 11 +-
.../org/apache/spark/sql/comet/operators.scala | 37 +++--
.../apache/comet/QueryContextInternerSuite.scala | 163 +++++++++++++++++++++
11 files changed, 443 insertions(+), 64 deletions(-)
diff --git a/.github/workflows/pr_build_linux.yml
b/.github/workflows/pr_build_linux.yml
index ff69a2cc68..2d01ecc104 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -343,6 +343,7 @@ jobs:
org.apache.comet.CometNativeSuite
org.apache.comet.CometConfSuite
org.apache.comet.CometPublicApiSuite
+ org.apache.comet.QueryContextInternerSuite
org.apache.comet.CometSetOpWithGroupBySuite
org.apache.comet.CometSparkSessionExtensionsSuite
org.apache.spark.CometPluginsSuite
diff --git a/.github/workflows/pr_build_macos.yml
b/.github/workflows/pr_build_macos.yml
index f3f3f08269..5f60c07451 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.github/workflows/pr_build_macos.yml
@@ -159,6 +159,7 @@ jobs:
org.apache.comet.CometNativeSuite
org.apache.comet.CometConfSuite
org.apache.comet.CometPublicApiSuite
+ org.apache.comet.QueryContextInternerSuite
org.apache.comet.CometSetOpWithGroupBySuite
org.apache.comet.CometSparkSessionExtensionsSuite
org.apache.spark.CometPluginsSuite
diff --git a/native/common/src/query_context.rs
b/native/common/src/query_context.rs
index 10ecef6550..405ce88d04 100644
--- a/native/common/src/query_context.rs
+++ b/native/common/src/query_context.rs
@@ -59,9 +59,12 @@ pub struct QueryContext {
}
impl QueryContext {
+ /// `sql_text` accepts either an owned `String` or an already-shared
`Arc<String>`. The plan
+ /// deserializer passes an `Arc` cloned from the plan's SQL text pool so
that all the contexts
+ /// in one plan share a single allocation instead of each cloning the full
query text.
#[allow(clippy::too_many_arguments)]
pub fn new(
- sql_text: String,
+ sql_text: impl Into<Arc<String>>,
start_index: i32,
stop_index: i32,
object_type: Option<String>,
@@ -70,7 +73,7 @@ impl QueryContext {
start_position: i32,
) -> Self {
Self {
- sql_text: Arc::new(sql_text),
+ sql_text: sql_text.into(),
start_index,
stop_index,
object_type,
diff --git a/native/core/src/execution/jni_api.rs
b/native/core/src/execution/jni_api.rs
index 09e5cf6cd5..f798b80335 100644
--- a/native/core/src/execution/jni_api.rs
+++ b/native/core/src/execution/jni_api.rs
@@ -781,6 +781,7 @@ pub unsafe extern "system" fn
Java_org_apache_comet_Native_executePlan(
let planner =
PhysicalPlanner::new(Arc::clone(&exec_context.session_ctx), partition)
.with_exec_id(exec_context_id)
+ .with_sql_text_pool(&exec_context.spark_plan)
.with_task_context(exec_context.task_context.clone());
let (scans, shuffle_scans, root_op) = planner.create_plan(
&exec_context.spark_plan,
diff --git a/native/core/src/execution/planner.rs
b/native/core/src/execution/planner.rs
index 5accd31609..f20dadf7f3 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -137,6 +137,7 @@ use datafusion_comet_spark_expr::{
};
use itertools::Itertools;
use jni::objects::{Global, JObject};
+use log::warn;
use num::{BigInt, ToPrimitive};
use object_store::path::Path;
use std::cmp::max;
@@ -240,6 +241,11 @@ pub struct PhysicalPlanner {
partition: i32,
session_ctx: Arc<SessionContext>,
query_context_registry: Arc<datafusion_comet_spark_expr::QueryContextMap>,
+ /// SQL texts referenced by `QueryContext.sql_text_idx`, taken from the
root operator's
+ /// `sql_text_pool`. Held as `Arc`s so that every context in the plan
shares one allocation
+ /// rather than each cloning the full query text out of the proto. Empty
when the plan was
+ /// serialized without interning, in which case contexts carry `sql_text`
inline.
+ sql_text_pool: Vec<Arc<String>>,
/// Captured at `createPlan` time on `ExecutionContext`; see that struct
for the
/// propagation rationale. `None` when no driving Spark task is available.
task_context: Option<Arc<Global<JObject<'static>>>>,
@@ -258,10 +264,81 @@ impl PhysicalPlanner {
session_ctx,
partition,
query_context_registry:
datafusion_comet_spark_expr::create_query_context_map(),
+ sql_text_pool: vec![],
task_context: None,
}
}
+ /// Load the SQL text pool from the root operator of the plan about to be
planned. Must be
+ /// called with the *root* operator: the JVM only populates the pool
there, and
+ /// `QueryContext.sql_text_idx` values are indices into it.
+ pub fn with_sql_text_pool(mut self, root: &Operator) -> Self {
+ self.sql_text_pool = root
+ .sql_text_pool
+ .iter()
+ .map(|text| Arc::new(text.clone()))
+ .collect();
+ self
+ }
+
+ /// Register the `QueryContext` carried by a serialized expression, if
any, so that native ANSI
+ /// errors raised by it can render Spark's `== SQL ... ==` block. Shared
by `create_expr` and
+ /// `create_agg_expr`.
+ fn register_query_context(
+ &self,
+ expr_id: Option<u64>,
+ ctx_proto: Option<&spark_expression::QueryContext>,
+ ) {
+ if let (Some(expr_id), Some(ctx_proto)) = (expr_id, ctx_proto) {
+ self.query_context_registry
+ .register(expr_id, self.build_query_context(ctx_proto));
+ }
+ }
+
+ /// Resolve a serialized `QueryContext` into a `QueryContext`, sharing the
pooled SQL text
+ /// when the context refers to the pool by index. Falls back to the inline
`sql_text` for
+ /// plans serialized without interning -- the paths that serialize an
operator directly rather
+ /// than through `convertBlock` (e.g. `CometNativeWriteExec`, the native
shuffle writer) and
+ /// the bare-`Expr` Parquet filter path, which has no root operator to
hang a pool on.
+ fn build_query_context(
+ &self,
+ ctx_proto: &spark_expression::QueryContext,
+ ) -> datafusion_comet_spark_expr::QueryContext {
+ let pooled = ctx_proto
+ .sql_text_idx
+ .map(|idx| match usize::try_from(idx).ok() {
+ Some(idx) => self.sql_text_pool.get(idx).map(Arc::clone),
+ None => None,
+ });
+
+ let sql_text: Arc<String> = match pooled {
+ Some(Some(text)) => text,
+ // An index we cannot resolve means the plan was interned against
a pool this planner
+ // never saw, so `sql_text` is empty and the error block would
silently come out blank.
+ // Warn rather than fail: a degraded error message is better than
a failed query.
+ Some(None) => {
+ warn!(
+ "QueryContext sql_text_idx {:?} is out of range for a SQL
text pool of {} \
+ entries; SQL context will be missing from any error
raised by this expression",
+ ctx_proto.sql_text_idx,
+ self.sql_text_pool.len()
+ );
+ Arc::new(ctx_proto.sql_text.clone())
+ }
+ None => Arc::new(ctx_proto.sql_text.clone()),
+ };
+
+ datafusion_comet_spark_expr::QueryContext::new(
+ sql_text,
+ ctx_proto.start_index,
+ ctx_proto.stop_index,
+ ctx_proto.object_type.clone(),
+ ctx_proto.object_name.clone(),
+ ctx_proto.line,
+ ctx_proto.start_position,
+ )
+ }
+
pub fn with_exec_id(mut self, exec_context_id: i64) -> Self {
self.exec_context_id = exec_context_id;
self
@@ -345,25 +422,7 @@ impl PhysicalPlanner {
spark_expr: &Expr,
input_schema: SchemaRef,
) -> Result<Arc<dyn PhysicalExpr>, ExecutionError> {
- // Register QueryContext if present
- if let (Some(expr_id), Some(ctx_proto)) =
- (spark_expr.expr_id, spark_expr.query_context.as_ref())
- {
- // Deserialize QueryContext from protobuf
- let query_ctx = datafusion_comet_spark_expr::QueryContext::new(
- ctx_proto.sql_text.clone(),
- ctx_proto.start_index,
- ctx_proto.stop_index,
- ctx_proto.object_type.clone(),
- ctx_proto.object_name.clone(),
- ctx_proto.line,
- ctx_proto.start_position,
- );
-
- // Register query context for error reporting
- let registry = &self.query_context_registry;
- registry.register(expr_id, query_ctx);
- }
+ self.register_query_context(spark_expr.expr_id,
spark_expr.query_context.as_ref());
// Try to use the modular registry first - this automatically handles
any registered expression types
if ExpressionRegistry::global().can_handle(spark_expr) {
@@ -2391,25 +2450,7 @@ impl PhysicalPlanner {
spark_expr: &AggExpr,
schema: SchemaRef,
) -> Result<AggregateFunctionExpr, ExecutionError> {
- // Register QueryContext if present
- if let (Some(expr_id), Some(ctx_proto)) =
- (spark_expr.expr_id, spark_expr.query_context.as_ref())
- {
- // Deserialize QueryContext from protobuf
- let query_ctx = datafusion_comet_spark_expr::QueryContext::new(
- ctx_proto.sql_text.clone(),
- ctx_proto.start_index,
- ctx_proto.stop_index,
- ctx_proto.object_type.clone(),
- ctx_proto.object_name.clone(),
- ctx_proto.line,
- ctx_proto.start_position,
- );
-
- // Register query context for error reporting
- let registry = &self.query_context_registry;
- registry.register(expr_id, query_ctx);
- }
+ self.register_query_context(spark_expr.expr_id,
spark_expr.query_context.as_ref());
match spark_expr.expr_struct.as_ref().unwrap() {
AggExprStruct::Count(expr) => {
@@ -4481,6 +4522,7 @@ mod tests {
fn test_unpack_dictionary_primitive() {
let op_scan = Operator {
plan_id: 0,
+ sql_text_pool: vec![],
children: vec![],
op_struct: Some(OpStruct::Scan(spark_operator::Scan {
fields: vec![spark_expression::DataType {
@@ -4546,6 +4588,7 @@ mod tests {
fn test_unpack_dictionary_string() {
let op_scan = Operator {
plan_id: 0,
+ sql_text_pool: vec![],
children: vec![],
op_struct: Some(OpStruct::Scan(spark_operator::Scan {
fields: vec![spark_expression::DataType {
@@ -4696,6 +4739,7 @@ mod tests {
Operator {
plan_id: 0,
+ sql_text_pool: vec![],
children: vec![child_op],
op_struct: Some(OpStruct::Filter(spark_operator::Filter {
predicate: Some(expr),
@@ -4722,6 +4766,7 @@ mod tests {
let op_scan = create_scan();
let op_join = Operator {
plan_id: 0,
+ sql_text_pool: vec![],
children: vec![op_scan.clone(), op_scan.clone()],
op_struct: Some(OpStruct::HashJoin(spark_operator::HashJoin {
left_join_keys: vec![create_bound_reference(0)],
@@ -4758,6 +4803,7 @@ mod tests {
fn create_scan() -> Operator {
Operator {
plan_id: 0,
+ sql_text_pool: vec![],
children: vec![],
op_struct: Some(OpStruct::Scan(spark_operator::Scan {
fields: vec![create_proto_datatype()],
@@ -4787,6 +4833,7 @@ mod tests {
// ScanExec: source=[CometScan parquet (unknown)], schema=[col_0:
Int32]
let op_scan = Operator {
plan_id: 0,
+ sql_text_pool: vec![],
children: vec![],
op_struct: Some(OpStruct::Scan(spark_operator::Scan {
fields: vec![
@@ -4834,6 +4881,7 @@ mod tests {
let projection = Operator {
children: vec![op_scan],
plan_id: 0,
+ sql_text_pool: vec![],
op_struct: Some(OpStruct::Projection(spark_operator::Projection {
project_list: vec![spark_expression::Expr {
expr_struct:
Some(ExprStruct::ScalarFunc(spark_expression::ScalarFunc {
@@ -4908,6 +4956,7 @@ mod tests {
// Mock scan operator with 3 INT32 columns
let op_scan = Operator {
plan_id: 0,
+ sql_text_pool: vec![],
children: vec![],
op_struct: Some(OpStruct::Scan(spark_operator::Scan {
fields: vec![
@@ -4958,6 +5007,7 @@ mod tests {
let projection = Operator {
children: vec![op_scan],
plan_id: 0,
+ sql_text_pool: vec![],
op_struct: Some(OpStruct::Projection(spark_operator::Projection {
project_list: vec![spark_expression::Expr {
expr_struct:
Some(ExprStruct::ScalarFunc(spark_expression::ScalarFunc {
@@ -5394,6 +5444,7 @@ mod tests {
// Create a Scan operator with Date32 (DATE) and Int8 (TINYINT) columns
let op_scan = Operator {
plan_id: 0,
+ sql_text_pool: vec![],
children: vec![],
op_struct: Some(OpStruct::Scan(spark_operator::Scan {
fields: vec![
@@ -5458,6 +5509,7 @@ mod tests {
let projection = Operator {
children: vec![op_scan],
plan_id: 1,
+ sql_text_pool: vec![],
op_struct: Some(OpStruct::Projection(spark_operator::Projection {
project_list: vec![subtract_expr],
})),
diff --git a/native/proto/src/proto/expr.proto
b/native/proto/src/proto/expr.proto
index 4c16186f57..04b20b5df3 100644
--- a/native/proto/src/proto/expr.proto
+++ b/native/proto/src/proto/expr.proto
@@ -109,7 +109,11 @@ message Expr {
// QueryContext provides SQL query context for error messages.
// Mirrors Spark's SQLQueryContext for rich error reporting.
message QueryContext {
- // Full SQL query text
+ // Full SQL query text. Left empty when `sql_text_idx` is set, which is the
case for a plan
+ // serialized through `CometNativeExec.convertBlock`. Still used by the
paths that serialize an
+ // operator directly without interning (e.g. `CometNativeWriteExec`, the
native shuffle writer,
+ // the Iceberg scan's `convertBlock` override) and by the bare-`Expr`
Parquet filter path, which
+ // has no root operator to carry a pool.
string sql_text = 1;
// Character offset where expression starts (0-based)
@@ -129,6 +133,11 @@ message QueryContext {
// Column position within the line (0-based)
int32 start_position = 7;
+
+ // Index into the root `Operator.sql_text_pool` holding this context's SQL
text. Set instead of
+ // `sql_text` so that a query text shared by many expressions is serialized
once per plan rather
+ // than once per expression.
+ optional int32 sql_text_idx = 8;
}
message AggExpr {
diff --git a/native/proto/src/proto/operator.proto
b/native/proto/src/proto/operator.proto
index d9515edc99..ced87262f3 100644
--- a/native/proto/src/proto/operator.proto
+++ b/native/proto/src/proto/operator.proto
@@ -36,6 +36,16 @@ message Operator {
// Spark plan ID
uint32 plan_id = 2;
+ // Pool of distinct SQL texts referenced by `QueryContext.sql_text_idx`.
Every expression that
+ // carries a QueryContext needs the full query text to render Spark's `==
SQL ... ==` error
+ // block, and in a typical plan they all want the *same* text -- embedding
it per expression made
+ // the query text over 95% of the serialized plan. The JVM interns the texts
into this pool on
+ // the root operator of each serialized native block (see
`QueryContextInterner`) and leaves
+ // `QueryContext.sql_text` empty.
+ //
+ // Only populated on the root operator; nested operators leave it empty.
+ repeated string sql_text_pool = 3;
+
oneof op_struct {
Scan scan = 100;
Projection projection = 101;
diff --git
a/spark/src/main/scala/org/apache/comet/serde/QueryContextInterner.scala
b/spark/src/main/scala/org/apache/comet/serde/QueryContextInterner.scala
new file mode 100644
index 0000000000..7d53be29f8
--- /dev/null
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryContextInterner.scala
@@ -0,0 +1,137 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet.serde
+
+import scala.collection.mutable
+
+import com.google.protobuf.Descriptors.FieldDescriptor
+import com.google.protobuf.Message
+
+import org.apache.comet.serde.OperatorOuterClass.Operator
+
+/**
+ * Deduplicates `QueryContext.sql_text` across a serialized native plan.
+ *
+ * Every expression that Comet converts carries a `QueryContext` so that
native ANSI errors can
+ * render Spark's `== SQL (line N, position M) ==` block, and that block
prints the *whole* query
+ * text -- so the text cannot be trimmed to a fragment. But in a typical plan
every expression
+ * carries the same text, and embedding it per expression made the query text
the overwhelming
+ * majority of the serialized plan: on a 778-character query with 90
expressions, 56 KB of a 58 KB
+ * plan (96%) was duplicated SQL text.
+ *
+ * This rewrites the tree so each distinct text appears once, in
`Operator.sql_text_pool` on the
+ * root operator, with each `QueryContext` referring to it by `sql_text_idx`.
The plan bytes are
+ * shipped in the stage's task binary, re-parsed and re-serialized per task by
+ * `CometExecRDD.compute` whenever the plan contains a native scan, and
decoded per task by the
+ * native planner, so the saving is paid back on every task.
+ *
+ * The 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 expression variants.
+ * `QueryContextInternerSuite` asserts that no un-pooled context survives in a
real plan.
+ */
+object QueryContextInterner {
+
+ /**
+ * Returns `op` with all `QueryContext.sql_text` values hoisted into the
root's `sql_text_pool`,
+ * or `op` itself when there is nothing to intern.
+ */
+ def intern(op: Operator): Operator = {
+ // Defensive no-op so interning is idempotent. Not reachable from the
single call site in
+ // `CometNativeExec.convertBlock`, which always starts from the
un-interned `nativeOp`.
+ if (op.getSqlTextPoolCount > 0) return op
+
+ val builder = op.toBuilder
+ // Each text is appended to the pool the first time it is seen and its
position becomes the
+ // index every context referring to it carries, so the pool needs no
separate second pass.
+ val indexOf = mutable.Map.empty[String, Int]
+ walk(
+ builder,
+ ctx => {
+ val idx = indexOf.getOrElseUpdate(
+ ctx.getSqlText, {
+ builder.addSqlTextPool(ctx.getSqlText)
+ builder.getSqlTextPoolCount - 1
+ })
+ Some(ctx.toBuilder.clearSqlText().setSqlTextIdx(idx).build())
+ })
+ if (indexOf.isEmpty) {
+ // No expression carried a context (e.g. a DataFrame-API query, where
Spark records no SQL
+ // text). Return the original rather than rebuilding an identical
message.
+ op
+ } else {
+ builder.build()
+ }
+ }
+
+ /**
+ * Returns `expr` with every `QueryContext` in it removed.
+ *
+ * Used to derive content hashes that must agree between the driver and the
executor even though
+ * only one side sees the interned form - see
`NativeScanPlanDataInjector.sourceKey`. Dropping
+ * contexts makes such a hash independent of how they happen to be encoded.
+ */
+ def stripQueryContexts(expr: ExprOuterClass.Expr): ExprOuterClass.Expr = {
+ val builder = expr.toBuilder
+ walk(builder, _ => None)
+ builder.build()
+ }
+
+ /**
+ * Depth-first walk over every nested message, applying `f` to each
`QueryContext` found on an
+ * `Expr` or `AggExpr`. `f` returns the replacement context, or `None` to
clear the field.
+ */
+ private def walk(
+ builder: Message.Builder,
+ f: ExprOuterClass.QueryContext => Option[ExprOuterClass.QueryContext]):
Unit = {
+ builder match {
+ case e: ExprOuterClass.Expr.Builder if e.hasQueryContext =>
+ f(e.getQueryContext) match {
+ case Some(ctx) => e.setQueryContext(ctx)
+ case None => e.clearQueryContext()
+ }
+ case a: ExprOuterClass.AggExpr.Builder if a.hasQueryContext =>
+ f(a.getQueryContext) match {
+ case Some(ctx) => a.setQueryContext(ctx)
+ case None => a.clearQueryContext()
+ }
+ case _ =>
+ }
+
+ val fields = builder.getDescriptorForType.getFields.iterator()
+ while (fields.hasNext) {
+ val field = fields.next()
+ // Map fields are repeated messages on the wire but reject
`getRepeatedFieldBuilder`, and no
+ // Comet map field has a message value type, so there is nothing to
visit inside them.
+ if (field.getJavaType == FieldDescriptor.JavaType.MESSAGE &&
!field.isMapField) {
+ if (field.isRepeated) {
+ var i = 0
+ val count = builder.getRepeatedFieldCount(field)
+ while (i < count) {
+ walk(builder.getRepeatedFieldBuilder(field, i), f)
+ i += 1
+ }
+ } else if (builder.hasField(field)) {
+ walk(builder.getFieldBuilder(field), f)
+ }
+ }
+ }
+ }
+}
diff --git
a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala
b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala
index 5b577f41f4..52990cfa3d 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala
@@ -360,15 +360,8 @@ object CometNativeScanExec {
scan: CometScanExec): CometNativeScanExec = {
// Generate unique key for this scan so PlanDataInjector can match
common+partition data.
// Multiple scans of same table with different projections/filters get
different keys.
- val common = nativeOp.getNativeScan.getCommon
- val source = common.getSource
- val keyComponents = Seq(
- common.getRequiredSchemaList.toString,
- common.getDataFiltersList.toString,
- common.getProjectionVectorList.toString,
- common.getFieldsList.toString)
- val hashCode = keyComponents.mkString("|").hashCode
- val sourceKey = s"${source}_${hashCode}"
+ // Derived by the injector that will look it up, so the two sides cannot
drift apart.
+ val sourceKey =
NativeScanPlanDataInjector.sourceKey(nativeOp.getNativeScan.getCommon)
val batchScanExec = CometNativeScanExec(
nativeOp,
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
index b422d60e67..9c4917e2d6 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
@@ -58,7 +58,7 @@ import org.apache.comet.{CometConf, CometExecIterator,
CometRuntimeException, Co
import org.apache.comet.CometSparkSessionExtensions.{isCometShuffleEnabled,
withFallbackReason}
import org.apache.comet.parquet.CometParquetUtils
import org.apache.comet.rules.CometExecRule
-import org.apache.comet.serde.{CometOperatorSerde, Compatible, Incompatible,
OperatorOuterClass, SupportLevel, Unsupported}
+import org.apache.comet.serde.{CometOperatorSerde, Compatible, Incompatible,
OperatorOuterClass, QueryContextInterner, SupportLevel, Unsupported}
import org.apache.comet.serde.OperatorOuterClass.{AggregateMode =>
CometAggregateMode, Operator}
import org.apache.comet.serde.QueryPlanSerde
import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto,
isStringCollationType, supportedSortType}
@@ -325,17 +325,28 @@ private[comet] object NativeScanPlanDataInjector extends
PlanDataInjector {
op.getNativeScan.hasCommon &&
!op.getNativeScan.hasFilePartition
- override def getKey(op: Operator): Option[String] = {
- // Reconstruct the same sourceKey that was used when storing the data
- val common = op.getNativeScan.getCommon
- val source = common.getSource
+ override def getKey(op: Operator): Option[String] =
Some(sourceKey(op.getNativeScan.getCommon))
+
+ /**
+ * The key under which a native scan's planning data is stored and looked
up. Called on the
+ * driver by `CometNativeScanExec.apply` to store, and on the executor by
[[getKey]] to look up
+ * \- both must derive the identical string from the same scan, so this is
the single definition
+ * rather than two mirrored copies.
+ *
+ * Data filters are stripped of their `QueryContext` before hashing: the
executor reads them
+ * back out of the interned plan (see `QueryContextInterner`) while the
driver holds the
+ * un-interned form, so including the context encoding would make the two
sides disagree. Only
+ * data filters can carry a context, so the other components are hashed
as-is.
+ */
+ private[comet] def sourceKey(common: OperatorOuterClass.NativeScanCommon):
String = {
+ val dataFilters = common.getDataFiltersList.asScala
+ .map(QueryContextInterner.stripQueryContexts(_).toString)
val keyComponents = Seq(
common.getRequiredSchemaList.toString,
- common.getDataFiltersList.toString,
+ dataFilters.mkString("[", ", ", "]"),
common.getProjectionVectorList.toString,
common.getFieldsList.toString)
- val hashCode = keyComponents.mkString("|").hashCode
- Some(s"${source}_${hashCode}")
+ s"${common.getSource}_${keyComponents.mkString("|").hashCode}"
}
override def inject(
@@ -869,12 +880,10 @@ abstract class CometNativeExec extends CometExec {
def convertBlock(): CometNativeExec = {
def transform(arg: Any): AnyRef = arg match {
case serializedPlan: SerializedPlan if serializedPlan.isEmpty =>
- val size = nativeOp.getSerializedSize
- val bytes = new Array[Byte](size)
- val codedOutput = CodedOutputStream.newInstance(bytes)
- nativeOp.writeTo(codedOutput)
- codedOutput.checkNoSpaceLeft()
- SerializedPlan(Some(bytes))
+ // Hoist duplicated QueryContext SQL text into a pool on the root
operator. This is the
+ // point where the whole native block is in hand, which is what the
pool indices are
+ // scoped to.
+
SerializedPlan(Some(CometExec.serializeNativePlan(QueryContextInterner.intern(nativeOp))))
case other: AnyRef => other
case null => null
}
diff --git
a/spark/src/test/scala/org/apache/comet/QueryContextInternerSuite.scala
b/spark/src/test/scala/org/apache/comet/QueryContextInternerSuite.scala
new file mode 100644
index 0000000000..fd33f24994
--- /dev/null
+++ b/spark/src/test/scala/org/apache/comet/QueryContextInternerSuite.scala
@@ -0,0 +1,163 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet
+
+import scala.collection.mutable
+
+import org.apache.spark.sql.{CometTestBase, DataFrame}
+import org.apache.spark.sql.comet.CometNativeExec
+
+import org.apache.comet.serde.{ExprOuterClass, OperatorOuterClass}
+
+class QueryContextInternerSuite extends CometTestBase {
+
+ private val sqlText =
+ """select k1,
+ | sum(d1) as sum_d1,
+ | max(case when i1 > 100 then d1 + 1 else d1 * 2 end) as m1
+ |from t1
+ |where i1 between 1 and 1000 and d1 > 1.0
+ |group by k1
+ |order by sum_d1 desc""".stripMargin
+
+ /** Every `Expr`/`AggExpr` QueryContext in an operator tree, at any nesting
depth. */
+ private def collectContexts(
+ root: OperatorOuterClass.Operator): Seq[ExprOuterClass.QueryContext] = {
+ val found = mutable.ArrayBuffer.empty[ExprOuterClass.QueryContext]
+
+ def walk(m: com.google.protobuf.Message): Unit = {
+ m match {
+ case e: ExprOuterClass.Expr if e.hasQueryContext => found +=
e.getQueryContext
+ case a: ExprOuterClass.AggExpr if a.hasQueryContext => found +=
a.getQueryContext
+ case _ =>
+ }
+ m.getAllFields.forEach { (_, v) =>
+ v match {
+ case child: com.google.protobuf.Message => walk(child)
+ case list: java.util.List[_] =>
+ list.forEach {
+ case child: com.google.protobuf.Message => walk(child)
+ case _ =>
+ }
+ case _ =>
+ }
+ }
+ }
+
+ walk(root)
+ found.toSeq
+ }
+
+ /**
+ * Each native block of `df`'s plan, as the un-interned `nativeOp` (what
would have been
+ * serialized before interning) paired with the bytes actually serialized
for it.
+ */
+ private def nativeBlocks(df: DataFrame): Seq[(OperatorOuterClass.Operator,
Array[Byte])] =
+ stripAQEPlan(df.queryExecution.executedPlan).collect {
+ case n: CometNativeExec if n.serializedPlanOpt.isDefined =>
+ (n.nativeOp, n.serializedPlanOpt.plan.get)
+ }
+
+ private def withTestTable(f: => Unit): Unit = {
+ withTempPath { dir =>
+ spark
+ .range(0, 100)
+ .selectExpr(
+ "cast(id % 17 as int) as k1",
+ "cast(id as decimal(20,4)) as d1",
+ "cast(id % 1000 as int) as i1")
+ .write
+ .mode("overwrite")
+ .parquet(dir.getAbsolutePath)
+ withTempView("t1") {
+ spark.read.parquet(dir.getAbsolutePath).createOrReplaceTempView("t1")
+ f
+ }
+ }
+ }
+
+ test("SQL text is interned into the root operator pool") {
+ withTestTable {
+ val blocks = nativeBlocks(spark.sql(sqlText))
+ assert(blocks.nonEmpty, "expected at least one serialized native block")
+
+ var sawContext = false
+ blocks.foreach { case (_, bytes) =>
+ val root = OperatorOuterClass.Operator.parseFrom(bytes)
+ val contexts = collectContexts(root)
+ if (contexts.nonEmpty) {
+ sawContext = true
+ contexts.foreach { ctx =>
+ assert(
+ ctx.hasSqlTextIdx,
+ s"QueryContext was not interned; carries inline sql_text:
${ctx.getSqlText}")
+ assert(ctx.getSqlText.isEmpty, "interned QueryContext should not
repeat sql_text")
+ assert(ctx.getSqlTextIdx >= 0 && ctx.getSqlTextIdx <
root.getSqlTextPoolCount)
+ }
+ // The query text is stored once for the block, not once per
expression.
+ assert(root.getSqlTextPoolCount == 1, s"pool:
${root.getSqlTextPoolList}")
+ assert(root.getSqlTextPool(0).contains("sum(d1)"))
+ }
+ }
+ assert(sawContext, "expected at least one expression to carry a
QueryContext")
+ }
+ }
+
+ test("interning shrinks the serialized plan") {
+ withTestTable {
+ // `nativeOp` is never interned - only the bytes written for it are - so
its serialized size
+ // is exactly what the plan weighed before this optimization.
+ val blocks = nativeBlocks(spark.sql(sqlText))
+ val before = blocks.map { case (nativeOp, _) =>
nativeOp.getSerializedSize }.sum
+ val after = blocks.map { case (_, bytes) => bytes.length }.sum
+
+ assert(
+ before > after * 4,
+ s"expected interning to shrink the plan substantially: $before ->
$after")
+ }
+ }
+
+ test("ANSI error still reports full SQL context after interning") {
+ // Deliberately its own integer-typed table rather than reusing
`withTestTable`: integer
+ // division routes through the native `CheckedBinaryExpr`, which is what
consults the
+ // QueryContext registry. Decimal division raises a plain compute error
with no SQL context,
+ // so it would not exercise this path at all.
+ withSQLConf("spark.sql.ansi.enabled" -> "true") {
+ withTempPath { dir =>
+ spark
+ .range(1, 10)
+ .selectExpr("cast(id as int) as a", "cast(0 as int) as b")
+ .write
+ .mode("overwrite")
+ .parquet(dir.getAbsolutePath)
+ withTempView("t_div") {
+
spark.read.parquet(dir.getAbsolutePath).createOrReplaceTempView("t_div")
+
+ val query = "select a / b as ratio from t_div"
+ val err = intercept[Exception](spark.sql(query).collect())
+ val msg = Option(err.getMessage).getOrElse("")
+ assert(msg.contains("DIVIDE_BY_ZERO"), msg)
+ // The pooled text must round-trip through native: the summary
echoes the full query.
+ assert(msg.contains(query), s"error message lost the SQL
context:\n$msg")
+ }
+ }
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
