This is an automated email from the ASF dual-hosted git repository. taiyang-li pushed a commit to branch fake_add_bolt_backend in repository https://gitbox.apache.org/repos/asf/gluten.git
commit 151f2420dc1a1487073014a29673d425e2a42098 Author: liyang.127 <[email protected]> AuthorDate: Mon Jun 29 22:07:47 2026 +0800 [GLUTEN][CORE] Share native write test checker base Extract the common native-write listener and assertion flow into backends-core so backend test suites only keep their plan normalization and native-plan detection hooks. --- .../apache/spark/gluten/NativeWriteChecker.scala | 53 ++++++++--------- .../gluten/test/NativeWriteCheckerBase.scala | 69 ++++++++++++++++++++++ .../apache/spark/sql/execution/WriteUtils.scala | 38 ++++-------- 3 files changed, 104 insertions(+), 56 deletions(-) diff --git a/backends-clickhouse/src/test/scala/org/apache/spark/gluten/NativeWriteChecker.scala b/backends-clickhouse/src/test/scala/org/apache/spark/gluten/NativeWriteChecker.scala index 49168988fa..416debaab8 100644 --- a/backends-clickhouse/src/test/scala/org/apache/spark/gluten/NativeWriteChecker.scala +++ b/backends-clickhouse/src/test/scala/org/apache/spark/gluten/NativeWriteChecker.scala @@ -18,11 +18,11 @@ package org.apache.spark.gluten import org.apache.gluten.config.GlutenConfig import org.apache.gluten.execution.{CHColumnarToCarrierRowExec, GlutenClickHouseWholeStageTransformerSuite} +import org.apache.gluten.test.NativeWriteCheckerBase import org.apache.spark.sql.{Dataset, Row} -import org.apache.spark.sql.execution.{ColumnarWriteFilesExec, QueryExecution} +import org.apache.spark.sql.execution.{ColumnarWriteFilesExec, QueryExecution, SparkPlan} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper -import org.apache.spark.sql.util.QueryExecutionListener import java.io.File @@ -30,6 +30,7 @@ import scala.reflect.runtime.universe.TypeTag trait NativeWriteChecker extends GlutenClickHouseWholeStageTransformerSuite + with NativeWriteCheckerBase with AdaptiveSparkPlanHelper { private val formats: Seq[String] = Seq("orc", "parquet") @@ -39,37 +40,31 @@ trait NativeWriteChecker protected def getWarehouseDir: String = dataHome + "/中文/spark-warehouse" // scalastyle:on nonascii - def withNativeWriteCheck(checkNative: Boolean)(block: => Unit): Unit = { - var nativeUsed = false + override protected def waitUntilListenerBusEmpty(): Unit = { + spark.sparkContext.listenerBus.waitUntilEmpty() + } - val queryListener = new QueryExecutionListener { - override def onFailure(f: String, qe: QueryExecution, e: Exception): Unit = { - fail("query failed", e) - } - override def onSuccess(funcName: String, qe: QueryExecution, duration: Long): Unit = { - if (!nativeUsed) { - val executedPlan = stripAQEPlan(qe.executedPlan) - nativeUsed = if (isSparkVersionGE("3.5")) { - executedPlan.find(_.isInstanceOf[ColumnarWriteFilesExec]).isDefined - } else { - executedPlan.find(_.isInstanceOf[CHColumnarToCarrierRowExec]).isDefined - } - } - } - } - try { - spark.listenerManager.register(queryListener) - block - spark.sparkContext.listenerBus.waitUntilEmpty() - assertResult(checkNative)(nativeUsed) - } finally { - spark.listenerManager.unregister(queryListener) + override protected def normalizeExecutedPlan(qe: QueryExecution): SparkPlan = { + stripAQEPlan(qe.executedPlan) + } + + override protected def isNativeWritePlan(plan: SparkPlan): Boolean = { + if (isSparkVersionGE("3.5")) { + plan.find(_.isInstanceOf[ColumnarWriteFilesExec]).isDefined + } else { + plan.find(_.isInstanceOf[CHColumnarToCarrierRowExec]).isDefined } } + + override protected def onNativeWriteQueryFailure( + funcName: String, + qe: QueryExecution, + error: Exception): Unit = { + fail("query failed", error) + } + def checkInsertQuery(sqlStr: String, checkNative: Boolean): Unit = - withNativeWriteCheck(checkNative) { - spark.sql(sqlStr) - } + checkNativeWrite(sqlStr, checkNative) def withDestinationTable(table: String, createTableSql: Option[String] = None)( f: => Unit): Unit = { diff --git a/backends-core/src/main/scala/org/apache/gluten/test/NativeWriteCheckerBase.scala b/backends-core/src/main/scala/org/apache/gluten/test/NativeWriteCheckerBase.scala new file mode 100644 index 0000000000..6e8507b44e --- /dev/null +++ b/backends-core/src/main/scala/org/apache/gluten/test/NativeWriteCheckerBase.scala @@ -0,0 +1,69 @@ +/* + * 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.test + +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.execution.{QueryExecution, SparkPlan} +import org.apache.spark.sql.util.QueryExecutionListener + +trait NativeWriteCheckerBase { + + protected def spark: SparkSession + + protected def normalizeExecutedPlan(qe: QueryExecution): SparkPlan = qe.executedPlan + + protected def isNativeWritePlan(plan: SparkPlan): Boolean + + protected def waitUntilListenerBusEmpty(): Unit + + protected def onNativeWriteQueryFailure( + funcName: String, + qe: QueryExecution, + error: Exception): Unit = {} + + protected def withNativeWriteCheck(expectedNativeWrite: Boolean)(block: => Unit): Unit = { + var nativeUsed = false + val queryListener = new QueryExecutionListener { + override def onFailure(funcName: String, qe: QueryExecution, e: Exception): Unit = { + onNativeWriteQueryFailure(funcName, qe, e) + } + + override def onSuccess(funcName: String, qe: QueryExecution, duration: Long): Unit = { + if (!nativeUsed) { + nativeUsed = isNativeWritePlan(normalizeExecutedPlan(qe)) + } + } + } + + try { + spark.listenerManager.register(queryListener) + block + waitUntilListenerBusEmpty() + assert( + nativeUsed == expectedNativeWrite, + s"Expected native write usage to be $expectedNativeWrite, but was $nativeUsed") + } finally { + spark.listenerManager.unregister(queryListener) + } + } + + protected def checkNativeWrite(sqlStr: String, expectNative: Boolean = true): Unit = { + withNativeWriteCheck(expectNative) { + spark.sql(sqlStr) + } + } +} diff --git a/backends-velox/src/test/scala/org/apache/spark/sql/execution/WriteUtils.scala b/backends-velox/src/test/scala/org/apache/spark/sql/execution/WriteUtils.scala index 416ce41f57..6789c6088b 100644 --- a/backends-velox/src/test/scala/org/apache/spark/sql/execution/WriteUtils.scala +++ b/backends-velox/src/test/scala/org/apache/spark/sql/execution/WriteUtils.scala @@ -17,47 +17,31 @@ package org.apache.spark.sql.execution import org.apache.gluten.execution.VeloxColumnarToCarrierRowExec +import org.apache.gluten.test.NativeWriteCheckerBase import org.apache.spark.sql.{DataFrame, GlutenQueryTest} import org.apache.spark.sql.catalyst.expressions.{BitwiseAnd, Expression, HiveHash, Literal, Pmod, UnsafeProjection} import org.apache.spark.sql.functions._ import org.apache.spark.sql.test.SQLTestUtils -import org.apache.spark.sql.util.QueryExecutionListener import java.io.File -trait WriteUtils extends GlutenQueryTest with SQLTestUtils { +trait WriteUtils extends GlutenQueryTest with SQLTestUtils with NativeWriteCheckerBase { + + override protected def waitUntilListenerBusEmpty(): Unit = { + spark.sparkContext.listenerBus.waitUntilEmpty() + } def tableDir(table: String): File = { val identifier = spark.sessionState.sqlParser.parseTableIdentifier(table) new File(spark.sessionState.catalog.defaultTablePath(identifier)) } - def checkNativeWrite(sqlStr: String, expectNative: Boolean = true): Unit = { - var nativeUsed = false - val queryListener = new QueryExecutionListener { - override def onFailure(f: String, qe: QueryExecution, e: Exception): Unit = {} - override def onSuccess(funcName: String, qe: QueryExecution, duration: Long): Unit = { - if (!nativeUsed) { - nativeUsed = if (isSparkVersionGE("3.4")) { - qe.executedPlan.exists(_.isInstanceOf[ColumnarWriteFilesExec]) - } else { - qe.executedPlan.exists(_.isInstanceOf[VeloxColumnarToCarrierRowExec]) - } - } - } - } - try { - spark.listenerManager.register(queryListener) - spark.sql(sqlStr) - spark.sparkContext.listenerBus.waitUntilEmpty() - if (expectNative) { - assert(nativeUsed) - } else { - assert(!nativeUsed) - } - } finally { - spark.listenerManager.unregister(queryListener) + override protected def isNativeWritePlan(plan: SparkPlan): Boolean = { + if (isSparkVersionGE("3.4")) { + plan.exists(_.isInstanceOf[ColumnarWriteFilesExec]) + } else { + plan.exists(_.isInstanceOf[VeloxColumnarToCarrierRowExec]) } } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
