diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala index f117372b3c..5af23382cf 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala @@ -25,11 +25,15 @@ import org.apache.spark.SparkConf import org.apache.spark.internal.Logging import org.apache.spark.network.util.ByteUnit import org.apache.spark.sql.{SparkSession, SparkSessionExtensions} +import org.apache.spark.sql.catalyst.FunctionIdentifier +import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionInfo, Literal} import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.TreeNode import org.apache.spark.sql.comet._ import org.apache.spark.sql.execution._ import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.StringType +import org.apache.spark.unsafe.types.UTF8String import org.apache.comet.CometConf._ import org.apache.comet.rules.{CometExecRule, CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometScanRule, CometSpark34AqeDppFallbackRule, EliminateRedundantTransitions, RevertNativeForTransitionHeavyStages} @@ -99,6 +103,7 @@ class CometSparkSessionExtensions extensions.injectQueryStagePrepRule { session => CometExecRule(session) } injectQueryStageOptimizerRuleShim(extensions, CometPlanAdaptiveDynamicPruningFilters) injectQueryStageOptimizerRuleShim(extensions, CometReuseSubquery) + extensions.injectFunction(CometSparkSessionExtensions.cometVersionFunction) } case class CometScanColumnar(session: SparkSession) extends ColumnarRule { @@ -119,6 +124,40 @@ class CometSparkSessionExtensions object CometSparkSessionExtensions extends Logging { lazy val isBigEndian: Boolean = ByteOrder.nativeOrder().equals(ByteOrder.BIG_ENDIAN) + /** + * SQL function `comet_version()` returning the Comet build version as a string. Registered via + * `SparkSessionExtensions.injectFunction` so users can query the loaded Comet version, e.g. + * `SELECT comet_version()`. + */ + private[comet] val cometVersionFunction + : (FunctionIdentifier, ExpressionInfo, Seq[Expression] => Expression) = { + val name = "comet_version" + val versionLiteral = Literal(UTF8String.fromString(COMET_VERSION), StringType) + // Populate the full ExpressionInfo so that Spark's ExpressionInfoSuite (SPARK-32870, and the + // example-output check) is satisfied for this injected function. The example output is derived + // from COMET_VERSION so it matches the value the function actually returns across releases. + val info = new ExpressionInfo( + classOf[Literal].getCanonicalName, + "", + name, + "_FUNC_() - Returns the Apache DataFusion Comet version.", + "", + s"\n Examples:\n > SELECT _FUNC_();\n $COMET_VERSION\n ", + "", + "misc_funcs", + "1.0.0", + "", + "") + val builder = (args: Seq[Expression]) => { + if (args.nonEmpty) { + throw new IllegalArgumentException( + s"Function $name does not accept arguments (got ${args.length})") + } + versionLiteral + } + (FunctionIdentifier(name), info, builder) + } + /** * Checks whether Comet extension should be loaded for Spark. */ diff --git a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala index 268fdf94eb..7f61cbac22 100644 --- a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala @@ -130,4 +130,20 @@ class CometSparkSessionExtensionsSuite extends CometTestBase { getCometShuffleMemorySize(conf, sqlConf) == getBytesFromMib((1024 * 0.2).toLong)) } + + test("comet_version() returns the Comet build version") { + withSQLConf(CometConf.COMET_EXEC_ENABLED.key -> "false") { + val result = spark.sql("SELECT comet_version()").collect() + assert(result.length == 1) + val version = result(0).getString(0) + assert(version == COMET_VERSION, s"Expected $COMET_VERSION but got $version") + } + } + + test("comet_version() rejects arguments") { + val ex = intercept[Exception] { + spark.sql("SELECT comet_version(1)").collect() + } + assert(ex.getMessage.contains("comet_version")) + } }