Copilot commented on code in PR #12163: URL: https://github.com/apache/gluten/pull/12163#discussion_r3400573851
########## gluten-substrait/src/main/scala/org/apache/gluten/execution/GlutenQueryContext.scala: ########## @@ -0,0 +1,78 @@ +/* + * 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.gluten.execution + +import org.apache.spark.SparkContext +import org.apache.spark.sql.execution.SQLExecution + +import java.util.UUID + +object GlutenQueryContext { + val QueryIdLocalProperty: String = "gluten.query.id" + private val QueryExecutionIdLocalProperty: String = "gluten.query.execution.id" + + private def sanitize(value: String): String = value.replaceAll("[^A-Za-z0-9_\\-]", "_") + + /** + * Initializes a query-scoped id on the driver before any query jobs are submitted. + * + * The id is cached in Spark local properties so all later driver-side async paths can inherit the + * same value. A companion execution id property prevents reusing a stale query id after the + * driver thread switches to a different SQL execution. + */ + def getOrCreateQueryId(sparkContext: SparkContext): String = synchronized { + val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) + val existingQueryId = sparkContext.getLocalProperty(QueryIdLocalProperty) + val existingExecutionId = sparkContext.getLocalProperty(QueryExecutionIdLocalProperty) + if (existingQueryId != null && existingExecutionId == executionId) { + return existingQueryId + } Review Comment: When `executionId` is `null`, `existingExecutionId == executionId` can still be true (both null), causing a previously generated random query id to be reused across unrelated (non-SQLExecution) runs. That can lead to cache-key collisions and incorrect cross-query reuse. Fix by only reusing the cached query id when `executionId != null`, or by storing a non-null sentinel/token in `QueryExecutionIdLocalProperty` when `executionId` is missing so successive calls don’t match on null. ########## gluten-substrait/src/main/scala/org/apache/gluten/execution/GlutenQueryContext.scala: ########## @@ -0,0 +1,78 @@ +/* + * 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.gluten.execution + +import org.apache.spark.SparkContext +import org.apache.spark.sql.execution.SQLExecution + +import java.util.UUID + +object GlutenQueryContext { + val QueryIdLocalProperty: String = "gluten.query.id" + private val QueryExecutionIdLocalProperty: String = "gluten.query.execution.id" + + private def sanitize(value: String): String = value.replaceAll("[^A-Za-z0-9_\\-]", "_") + + /** + * Initializes a query-scoped id on the driver before any query jobs are submitted. + * + * The id is cached in Spark local properties so all later driver-side async paths can inherit the + * same value. A companion execution id property prevents reusing a stale query id after the + * driver thread switches to a different SQL execution. + */ + def getOrCreateQueryId(sparkContext: SparkContext): String = synchronized { + val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) + val existingQueryId = sparkContext.getLocalProperty(QueryIdLocalProperty) + val existingExecutionId = sparkContext.getLocalProperty(QueryExecutionIdLocalProperty) + if (existingQueryId != null && existingExecutionId == executionId) { + return existingQueryId + } + + val queryId = if (executionId != null) { + s"Gluten_Query_${sanitize(sparkContext.applicationId)}_Execution_$executionId" + } else { + val randomId = sanitize(UUID.randomUUID().toString) + s"Gluten_Query_${sanitize(sparkContext.applicationId)}_$randomId" + } + + setQueryId(sparkContext, queryId, executionId) Review Comment: When `executionId` is `null`, `existingExecutionId == executionId` can still be true (both null), causing a previously generated random query id to be reused across unrelated (non-SQLExecution) runs. That can lead to cache-key collisions and incorrect cross-query reuse. Fix by only reusing the cached query id when `executionId != null`, or by storing a non-null sentinel/token in `QueryExecutionIdLocalProperty` when `executionId` is missing so successive calls don’t match on null. ########## gluten-substrait/src/main/scala/org/apache/gluten/execution/GlutenQueryContext.scala: ########## @@ -0,0 +1,78 @@ +/* + * 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.gluten.execution + +import org.apache.spark.SparkContext +import org.apache.spark.sql.execution.SQLExecution + +import java.util.UUID + +object GlutenQueryContext { + val QueryIdLocalProperty: String = "gluten.query.id" + private val QueryExecutionIdLocalProperty: String = "gluten.query.execution.id" + + private def sanitize(value: String): String = value.replaceAll("[^A-Za-z0-9_\\-]", "_") + + /** + * Initializes a query-scoped id on the driver before any query jobs are submitted. + * + * The id is cached in Spark local properties so all later driver-side async paths can inherit the + * same value. A companion execution id property prevents reusing a stale query id after the + * driver thread switches to a different SQL execution. + */ + def getOrCreateQueryId(sparkContext: SparkContext): String = synchronized { + val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) + val existingQueryId = sparkContext.getLocalProperty(QueryIdLocalProperty) + val existingExecutionId = sparkContext.getLocalProperty(QueryExecutionIdLocalProperty) + if (existingQueryId != null && existingExecutionId == executionId) { + return existingQueryId + } + + val queryId = if (executionId != null) { + s"Gluten_Query_${sanitize(sparkContext.applicationId)}_Execution_$executionId" + } else { + val randomId = sanitize(UUID.randomUUID().toString) + s"Gluten_Query_${sanitize(sparkContext.applicationId)}_$randomId" + } + + setQueryId(sparkContext, queryId, executionId) + queryId + } + + def getQueryId(sparkContext: SparkContext): Option[String] = + Option(sparkContext.getLocalProperty(QueryIdLocalProperty)) + + def setQueryId(sparkContext: SparkContext, queryId: String, executionId: String): Unit = { + sparkContext.setLocalProperty(QueryIdLocalProperty, queryId) + sparkContext.setLocalProperty(QueryExecutionIdLocalProperty, executionId) + } Review Comment: `executionId` is modeled as a non-null `String`, but call sites pass the value from `sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)` which can be null. Even though Spark’s API accepts null values, the current signature makes null usage implicit and easy to misuse. Consider changing `setQueryId/withQueryId` to accept `Option[String]` (or provide an overload) and explicitly handle the unset case (e.g., clear the property or store a sentinel). ########## cpp/velox/jni/VeloxJniWrapper.cc: ########## @@ -1138,30 +1146,52 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_native } hashTableBuilders[0]->setHashTable(std::move(mainTable)); + + auto* cache = facebook::velox::exec::HashTableCache::instance(); + if (!cache->hasTable(hashTableId)) { + std::cout << "VeloxJniWrapper inject hash table id is " << hashTableId << "\n"; + cache->injectTable( + hashTableId, + hashTableBuilders[0]->hashTable(), + hashTableBuilders[0]->joinHasNullKeys(), + defaultLeafVeloxMemoryPool()); + } + return gluten::getHashTableObjStore()->save(hashTableBuilders[0]); JNI_METHOD_END(kInvalidObjectHandle) } JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_cloneHashTable( // NOLINT JNIEnv* env, jclass, + jstring cacheKey, jlong tableHandler) { JNI_METHOD_START + auto cacheKeyStr = jStringToCString(env, cacheKey); auto hashTableHandler = ObjectStore::retrieve<gluten::HashTableBuilder>(tableHandler); + auto* cache = facebook::velox::exec::HashTableCache::instance(); + if (!cache->hasTable(cacheKeyStr)) { + cache->injectTable( + cacheKeyStr, hashTableHandler->hashTable(), hashTableHandler->joinHasNullKeys(), defaultLeafVeloxMemoryPool()); + } + return gluten::getHashTableObjStore()->save(hashTableHandler); JNI_METHOD_END(kInvalidObjectHandle) } JNIEXPORT void JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_clearHashTable( // NOLINT JNIEnv* env, jclass, + jstring cacheKey, jlong tableHandler) { JNI_METHOD_START - auto hashTableHandler = ObjectStore::retrieve<gluten::HashTableBuilder>(tableHandler); - hashTableHandler->hashTable()->clear(true); + auto cacheKeyStr = jStringToCString(env, cacheKey); + std::cout << "VeloxJniWrapper clear hash table id is " << cacheKeyStr << "\n"; + facebook::velox::exec::HashTableCache::instance()->drop(cacheKeyStr); Review Comment: The `hasTable` + `injectTable` sequence is not atomic; concurrent builds for the same key can race and inject twice (depending on Velox cache semantics). Prefer an atomic API (if available) or protect the check/inject with a mutex around the critical section. Also, `std::cout` logging in JNI paths is noisy and hard to control in production—use the project’s logging facility (e.g., `LOG(INFO/WARNING)`) so logs can be filtered and routed consistently. ########## cpp/velox/substrait/SubstraitToVeloxPlan.cc: ########## @@ -448,29 +448,11 @@ core::PlanNodePtr SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait:: } else if ( sJoin.has_advanced_extension() && SubstraitParser::configSetInOptimization(sJoin.advanced_extension(), "isBHJ=")) { - std::string hashTableId = sJoin.hashtableid(); - - std::shared_ptr<core::OpaqueHashTable> opaqueSharedHashTable = nullptr; - bool joinHasNullKeys = false; - - try { - auto hashTableBuilder = ObjectStore::retrieve<gluten::HashTableBuilder>(getJoin(hashTableId)); - joinHasNullKeys = hashTableBuilder->joinHasNullKeys(); - auto originalShared = hashTableBuilder->hashTable(); - opaqueSharedHashTable = std::shared_ptr<core::OpaqueHashTable>( - originalShared, reinterpret_cast<core::OpaqueHashTable*>(originalShared.get())); - - LOG(INFO) << "Successfully retrieved and aliased HashTable for reuse. ID: " << hashTableId; - } catch (const std::exception& e) { - LOG(WARNING) - << "Error retrieving HashTable from ObjectStore: " << e.what() - << ". Falling back to building new table. To ensure correct results, please verify that spark.gluten.velox.buildHashTableOncePerExecutor.enabled is set to false."; - opaqueSharedHashTable = nullptr; - } - + const auto& hashTableId = sJoin.hashtableid(); + const auto joinNodeId = hashTableId.empty() ? nextPlanNodeId() : hashTableId; // Create HashJoinNode node return std::make_shared<core::HashJoinNode>( - nextPlanNodeId(), + joinNodeId, Review Comment: Using `hashTableId` as the Velox `PlanNode` id risks creating duplicate node IDs when the same build side is reused for multiple join nodes within the same plan. Velox plan node IDs are typically expected to be unique within a plan, and collisions can break plan execution/diagnostics or corrupt any per-node state keyed by id. Prefer keeping `nextPlanNodeId()` for uniqueness and passing the cache key through a dedicated field/config (or embed uniqueness into the id, e.g., prefix/suffix with the generated node id). ########## ep/build-velox/src/get-velox.sh: ########## @@ -17,9 +17,9 @@ set -exu CURRENT_DIR=$(cd "$(dirname "$BASH_SOURCE")"; pwd) -VELOX_REPO=https://github.com/IBM/velox.git -VELOX_BRANCH=dft-2026_06_06 -VELOX_ENHANCED_BRANCH=ibm-2026_06_06 +VELOX_REPO=https://github.com/JkSelf/velox.git +VELOX_BRANCH=dft-2026_06_06-hashtable-cache +VELOX_ENHANCED_BRANCH=ibm-2026_06_06-hashtable-cache Review Comment: This switches the Velox dependency to a personal fork, which is a supply-chain risk and makes builds less reproducible/auditable. Prefer an official upstream/org fork (or a pinned commit SHA in a trusted repo). If you need to keep this temporarily, make the repo/branch configurable via environment variables and default to the canonical upstream so CI/release builds don’t depend on a personal namespace. ########## backends-velox/src/main/scala/org/apache/gluten/execution/HashJoinExecTransformer.scala: ########## @@ -135,8 +135,10 @@ case class BroadcastHashJoinExecTransformer( override def columnarInputRDDs: Seq[RDD[ColumnarBatch]] = { val streamedRDD = getColumnarInputRDDs(streamedPlan) val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) + val queryId = GlutenQueryContext.getOrCreateQueryId(sparkContext) + val cacheKey = s"$queryId:$buildHashTableId" if (executionId != null) { - GlutenDriverEndpoint.collectResources(executionId, buildBroadcastTableId) + GlutenDriverEndpoint.collectResources(executionId, cacheKey) Review Comment: The cache key is composed ad-hoc with a `:` delimiter. Since this key now crosses multiple layers (driver endpoint, JNI, native cache, and may appear in logs/metrics), it would be safer and easier to evolve if it were built in one place with explicit escaping/sanitization (e.g., a helper that guarantees allowed characters and stable formatting). This prevents future breakage if either component starts emitting or requiring different formats. ########## cpp/velox/jni/VeloxJniWrapper.cc: ########## @@ -1060,6 +1061,13 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_native nullptr); builder->setHashTable(std::move(mainTable)); + auto* cache = facebook::velox::exec::HashTableCache::instance(); + + if (!cache->hasTable(hashTableId)) { + std::cout << "VeloxJniWrapper inject hash table id is " << hashTableId << "\n"; + cache->injectTable(hashTableId, builder->hashTable(), builder->joinHasNullKeys(), defaultLeafVeloxMemoryPool()); + } Review Comment: The `hasTable` + `injectTable` sequence is not atomic; concurrent builds for the same key can race and inject twice (depending on Velox cache semantics). Prefer an atomic API (if available) or protect the check/inject with a mutex around the critical section. Also, `std::cout` logging in JNI paths is noisy and hard to control in production—use the project’s logging facility (e.g., `LOG(INFO/WARNING)`) so logs can be filtered and routed consistently. ########## gluten-arrow/src/main/java/org/apache/gluten/vectorized/NativePlanEvaluator.java: ########## @@ -113,4 +119,31 @@ public long spill(MemoryTarget self, Spiller.Phase phase, long size) { private ColumnarBatchOutIterator createOutIterator(Runtime runtime, long itrHandle) { return new ColumnarBatchOutIterator(runtime, itrHandle); } + + private static long getExecutionId(TaskContext taskContext) { + final String executionId = taskContext.getLocalProperty(SPARK_EXECUTION_ID_KEY); + if (executionId == null) { + return INVALID_EXECUTION_ID; + } + try { + return Long.parseLong(executionId); + } catch (NumberFormatException e) { + LOGGER.warn( + "Invalid Spark execution id '{}', fallback to {}", executionId, INVALID_EXECUTION_ID); + return INVALID_EXECUTION_ID; + } + } + + private static String getQueryId(TaskContext taskContext) { + final String queryId = taskContext.getLocalProperty(GLUTEN_QUERY_ID_KEY); + if (queryId != null) { + return queryId; + } + + final String executionId = taskContext.getLocalProperty(SPARK_EXECUTION_ID_KEY); + if (executionId != null) { + return "Gluten_Execution_" + executionId; + } + return "Gluten_Execution_"; + } Review Comment: If both `gluten.query.id` and Spark execution id are missing, the fallback returns the constant `\"Gluten_Execution_\"`, which can cause unrelated tasks/queries to share the same query identifier and collide in any native caching keyed by query id. Use a collision-resistant fallback (e.g., include stageId + taskAttemptId, or generate a UUID) when execution id is unavailable. -- 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]
