diff --git a/readme.md b/readme.md index 1741f0e75..3b16632d8 100644 --- a/readme.md +++ b/readme.md @@ -6,6 +6,8 @@ Substrait Java is a project that makes it easier to build [Substrait](https://su 3) **Spark** is the module that provides an API for translating a Substrait plan to and from a Spark query plan. The most commonly used logical relations and functions are supported, including those generated from all of the TPC-H and TPC-DS queries. The supported features are formally specified in the Substrait dialect file [spark_dialect.yaml](spark/spark_dialect.yaml). +Spark exports decimal and i64 `SUM`/`AVG` through `extension:substrait:spark`, including window calls. These functions carry Spark's result types, decimal buffer behavior, rounding and overflow modes. Decimal overflow outcomes can still depend on physical buffer materialization, as described in the function definitions. Consumers need the [Spark extension definitions](spark/src/main/resources/spark.yml) and implementations of those semantics. The functions support `INITIAL_TO_RESULT` only; exporting or importing staged calls is rejected. Optimized decimal `AVG` retains its unscaled i64 average, fp64 division and final decimal cast. Decimal optimizations can also reference the Spark `unscaled_value`, `make_decimal` and `try_make_decimal` scalar functions. Optimized SUM retains its i64 accumulator and the final unscaled decimal conversion. + ## Building Building requires a **JDK 17** installation, and Gradle itself must run on JDK 17 (see [Gradle & JDK 17](CONTRIBUTING.md#gradle--jdk-17) in the contributing guide for details). diff --git a/spark/spark_dialect.yaml b/spark/spark_dialect.yaml index f5d5f08ee..15e619dbe 100644 --- a/spark/spark_dialect.yaml +++ b/spark/spark_dialect.yaml @@ -587,6 +587,27 @@ supported_scalar_functions: notation: "FUNCTION" supported_impls: - "date_i32" +- source: "spark" + name: "make_decimal" + system_metadata: + name: "make_decimal" + notation: "FUNCTION" + supported_impls: + - "i64_type" +- source: "spark" + name: "try_make_decimal" + system_metadata: + name: "try_make_decimal" + notation: "FUNCTION" + supported_impls: + - "i64_type" +- source: "spark" + name: "unscaled_value" + system_metadata: + name: "unscaledvalue" + notation: "FUNCTION" + supported_impls: + - "dec" - source: "string" name: "contains" system_metadata: @@ -708,7 +729,6 @@ supported_aggregate_functions: - "fp64" - "i16" - "i32" - - "i64" - "i8" - source: "arithmetic" name: "max" @@ -751,15 +771,7 @@ supported_aggregate_functions: - "fp64" - "i16" - "i32" - - "i64" - "i8" -- source: "arithmetic_decimal" - name: "avg" - system_metadata: - name: "avg" - notation: "FUNCTION" - supported_impls: - - "dec" - source: "arithmetic_decimal" name: "max" system_metadata: @@ -774,13 +786,6 @@ supported_aggregate_functions: notation: "FUNCTION" supported_impls: - "dec" -- source: "arithmetic_decimal" - name: "sum" - system_metadata: - name: "sum" - notation: "FUNCTION" - supported_impls: - - "dec" - source: "datetime" name: "max" system_metadata: @@ -803,7 +808,48 @@ supported_aggregate_functions: - "iyear" - "pts" - "ptstz" +- source: "spark" + name: "avg" + system_metadata: + name: "avg" + notation: "FUNCTION" + supported_impls: + - "dec" + - "i64" +- source: "spark" + name: "sum" + system_metadata: + name: "sum" + notation: "FUNCTION" + supported_impls: + - "dec" + - "i64" supported_window_functions: +- source: "aggregate_approx" + name: "approx_count_distinct" + system_metadata: + name: "approx_count_distinct" + notation: "FUNCTION" + supported_impls: + - "any" +- source: "aggregate_generic" + name: "any_value" + system_metadata: + name: "first" + notation: "FUNCTION" + supported_impls: + - "any" +- source: "arithmetic" + name: "avg" + system_metadata: + name: "avg" + notation: "FUNCTION" + supported_impls: + - "fp32" + - "fp64" + - "i16" + - "i32" + - "i8" - source: "arithmetic" name: "cume_dist" system_metadata: @@ -832,6 +878,30 @@ supported_window_functions: notation: "FUNCTION" supported_impls: - "any_i32_any" +- source: "arithmetic" + name: "max" + system_metadata: + name: "max" + notation: "FUNCTION" + supported_impls: + - "fp32" + - "fp64" + - "i16" + - "i32" + - "i64" + - "i8" +- source: "arithmetic" + name: "min" + system_metadata: + name: "min" + notation: "FUNCTION" + supported_impls: + - "fp32" + - "fp64" + - "i16" + - "i32" + - "i64" + - "i8" - source: "arithmetic" name: "nth_value" system_metadata: @@ -867,3 +937,73 @@ supported_window_functions: notation: "FUNCTION" supported_impls: - "" +- source: "arithmetic" + name: "std_dev" + system_metadata: + name: "stddev_samp" + notation: "FUNCTION" + supported_impls: + - "fp64" +- source: "arithmetic" + name: "sum" + system_metadata: + name: "sum" + notation: "FUNCTION" + supported_impls: + - "fp32" + - "fp64" + - "i16" + - "i32" + - "i8" +- source: "arithmetic_decimal" + name: "max" + system_metadata: + name: "max" + notation: "FUNCTION" + supported_impls: + - "dec" +- source: "arithmetic_decimal" + name: "min" + system_metadata: + name: "min" + notation: "FUNCTION" + supported_impls: + - "dec" +- source: "datetime" + name: "max" + system_metadata: + name: "max" + notation: "FUNCTION" + supported_impls: + - "date" + - "iday" + - "iyear" + - "pts" + - "ptstz" +- source: "datetime" + name: "min" + system_metadata: + name: "min" + notation: "FUNCTION" + supported_impls: + - "date" + - "iday" + - "iyear" + - "pts" + - "ptstz" +- source: "spark" + name: "avg" + system_metadata: + name: "avg" + notation: "FUNCTION" + supported_impls: + - "dec" + - "i64" +- source: "spark" + name: "sum" + system_metadata: + name: "sum" + notation: "FUNCTION" + supported_impls: + - "dec" + - "i64" diff --git a/spark/src/main/resources/spark.yml b/spark/src/main/resources/spark.yml index 48281ea03..0bafd366e 100644 --- a/spark/src/main/resources/spark.yml +++ b/spark/src/main/resources/spark.yml @@ -16,6 +16,42 @@ --- urn: extension:substrait:spark scalar_functions: + - name: unscaled_value + description: >- + Return the decimal's unscaled integer as a signed i64. Return null for + null input and raise an error if the unscaled integer does not fit i64. + impls: + - args: + - name: x + value: DECIMAL + nullability: MIRROR + return: i64 + - name: make_decimal + description: >- + Interpret an i64 as an unscaled decimal with the precision and scale of + the type argument. Return null for null input and raise an error when + the unscaled value exceeds the target precision. + impls: + - args: + - name: value + value: i64 + - name: target + type: DECIMAL + nullability: MIRROR + return: DECIMAL + - name: try_make_decimal + description: >- + Interpret an i64 as an unscaled decimal with the precision and scale of + the type argument. Return null for null input or when the unscaled value + exceeds the target precision. + impls: + - args: + - name: value + value: i64 + - name: target + type: DECIMAL + nullability: DECLARED_OUTPUT + return: DECIMAL? - name: add description: >- Adds days to a date @@ -26,3 +62,69 @@ scalar_functions: - name: days value: i32 return: date + +aggregate_functions: + - name: sum + description: >- + Spark sum. Decimal inputs have a decimal buffer type + and result; i64 inputs use an i64 accumulator and result. Ignore null + inputs and return null for an empty or all-null group. Overflow returns + null for NULL, raises an error for ERROR, or wraps for WRAP (i64 only). + Decimal updates and merges do not check the buffer's declared precision. + Finalization checks the result precision. An out-of-range live decimal + buffer can therefore recover after cancellation, even after an earlier + window result was null. Materialization that checks the buffer precision + can instead store null; a stored null propagates through updates and merges + and is handled by the overflow option at finalization. Decimal overflow + outcomes can depend on physical buffer materialization. For i64 with NULL, + overflow during addition makes the accumulator remain null. Only + INITIAL_TO_RESULT is supported. + impls: + - args: + - name: x + value: DECIMAL + options: + overflow: + values: ["NULL", ERROR] + nullability: DECLARED_OUTPUT + decomposable: NONE + return: DECIMAL? + - args: + - name: x + value: i64 + options: + overflow: + values: [WRAP, "NULL", ERROR] + nullability: DECLARED_OUTPUT + decomposable: NONE + return: i64? + - name: avg + description: >- + Spark average. Decimal inputs use a decimal sum + buffer and an i64 count, then divide with HALF_UP rounding to + decimal. i64 inputs use an fp64 sum and i64 + count and return fp64. Ignore null inputs and return null for an empty + or all-null group. The decimal sum buffer has the same unchecked updates, + merges and materialization behavior as sum:dec. A null sum buffer or result + overflow returns null for NULL and raises an error for ERROR at finalization. + Decimal overflow outcomes can depend on physical buffer materialization. + Only INITIAL_TO_RESULT is supported. + impls: + - args: + - name: x + value: DECIMAL + options: + overflow: + values: ["NULL", ERROR] + nullability: DECLARED_OUTPUT + decomposable: NONE + return: DECIMAL? + - args: + - name: x + value: i64 + options: + overflow: + values: ["NULL", ERROR] + nullability: DECLARED_OUTPUT + decomposable: NONE + return: fp64? diff --git a/spark/src/main/scala/io/substrait/spark/SparkExtension.scala b/spark/src/main/scala/io/substrait/spark/SparkExtension.scala index 27eac65ec..b130b72a8 100644 --- a/spark/src/main/scala/io/substrait/spark/SparkExtension.scala +++ b/spark/src/main/scala/io/substrait/spark/SparkExtension.scala @@ -16,7 +16,7 @@ */ package io.substrait.spark -import io.substrait.spark.expression.{ToAggregateFunction, ToWindowFunction} +import io.substrait.spark.expression.{SparkAggregateFunctions, ToAggregateFunction, ToWindowFunction} import io.substrait.extension.{DefaultExtensionCatalog, SimpleExtension} @@ -42,19 +42,23 @@ object SparkExtension { ret.toSeq } - /** - * Standard extensions only: `spark.yml` declares no aggregates or windows, and - * [[toAggregateFunction]] / [[toWindowFunction]] bind against these same collections, so a - * `spark.yml` aggregate would be neither advertised in the dialect nor bindable at runtime. - */ lazy val StandardAggregateFunctions: Seq[SimpleExtension.AggregateFunctionVariant] = EXTENSION_COLLECTION.aggregateFunctions().asScala.toSeq - /** @see [[StandardAggregateFunctions]] */ lazy val StandardWindowFunctions: Seq[SimpleExtension.WindowFunctionVariant] = EXTENSION_COLLECTION.windowFunctions().asScala.toSeq - val toAggregateFunction: ToAggregateFunction = ToAggregateFunction(StandardAggregateFunctions) + /** Functions bound by both the runtime converters and the generated dialect. */ + lazy val AggregateFunctions: Seq[SimpleExtension.AggregateFunctionVariant] = + StandardAggregateFunctions.filterNot(SparkAggregateFunctions.replaces) ++ + SparkImpls.aggregateFunctions().asScala + + lazy val WindowFunctions: Seq[SimpleExtension.WindowFunctionVariant] = + StandardWindowFunctions.filterNot(SparkAggregateFunctions.replaces) ++ + SparkImpls.windowFunctions().asScala + + val toAggregateFunction: ToAggregateFunction = ToAggregateFunction(AggregateFunctions) + + val toWindowFunction: ToWindowFunction = ToWindowFunction(WindowFunctions) - val toWindowFunction: ToWindowFunction = ToWindowFunction(StandardWindowFunctions) } diff --git a/spark/src/main/scala/io/substrait/spark/expression/FunctionMappings.scala b/spark/src/main/scala/io/substrait/spark/expression/FunctionMappings.scala index 76ec9c2fc..a3b0f24e0 100644 --- a/spark/src/main/scala/io/substrait/spark/expression/FunctionMappings.scala +++ b/spark/src/main/scala/io/substrait/spark/expression/FunctionMappings.scala @@ -113,6 +113,7 @@ class FunctionMappings { } val SCALAR_SIGS: Seq[Sig] = Seq( + s[UnscaledValue]("unscaled_value"), s[Add]("add"), s[Subtract]("subtract"), s[Multiply]("multiply"), diff --git a/spark/src/main/scala/io/substrait/spark/expression/SparkAggregateFunctions.scala b/spark/src/main/scala/io/substrait/spark/expression/SparkAggregateFunctions.scala new file mode 100644 index 000000000..181480ceb --- /dev/null +++ b/spark/src/main/scala/io/substrait/spark/expression/SparkAggregateFunctions.scala @@ -0,0 +1,182 @@ +/* + * 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 io.substrait.spark.expression + +import io.substrait.spark.ToSubstraitType + +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, EvalMode, Expression} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateFunction, Average, Complete, DeclarativeAggregate, Sum} +import org.apache.spark.sql.types.{DecimalType, LongType} + +import io.substrait.`type`.Type +import io.substrait.expression.{Expression => SExpression, FunctionArg, FunctionOption} +import io.substrait.extension.{DefaultExtensionCatalog, SimpleExtension} + +import scala.jdk.CollectionConverters._ + +/** The accumulator, rounding and overflow contracts carried by the Spark extension aggregates. */ +private[spark] object SparkAggregateFunctions { + val URN = "extension:substrait:spark" + val keys: Set[String] = Set("sum:dec", "avg:dec", "sum:i64", "avg:i64") + + def isSparkFunction(function: SimpleExtension.Function): Boolean = + function.urn() == URN && keys.contains(function.key()) + + def replaces(function: SimpleExtension.Function): Boolean = + (function.urn() == DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC_DECIMAL && + Set("sum:dec", "avg:dec").contains(function.key())) || + (function.urn() == DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC && + Set("sum:i64", "avg:i64").contains(function.key())) + + def validateResult( + function: SimpleExtension.Function, + arguments: Seq[FunctionArg], + outputType: Type): Unit = { + if (isSparkFunction(function)) { + val inputTypes = arguments.map { + case expression: SExpression => expression.getType + case other => + throw new IllegalArgumentException(s"Expected a Spark aggregate value, got $other") + } + validateTypes(function, inputTypes, outputType) + } + } + + private def validateTypes( + function: SimpleExtension.Function, + inputTypes: Seq[Type], + outputType: Type): Unit = { + val derived = function.resolveType(inputTypes.asJava) + if (derived != outputType) { + throw new IllegalArgumentException( + s"Spark ${function.key()} declares $outputType but its definition derives $derived") + } + } + + def options( + expression: AggregateExpression, + function: SimpleExtension.Function): Seq[FunctionOption] = { + if (!isSparkFunction(function)) return Seq.empty + if (expression.mode != Complete) { + throw new UnsupportedOperationException( + s"Spark extension ${function.key()} only supports complete aggregation, got ${expression.mode}") + } + val (mode, wraps) = expression.aggregateFunction match { + case sum: Sum => (sum.evalMode, sum.child.dataType == LongType) + case avg: Average => (avg.evalMode, false) + case other => throw new UnsupportedOperationException(s"Unsupported Spark aggregate: $other") + } + val overflow = mode match { + case EvalMode.ANSI => "ERROR" + case EvalMode.LEGACY if wraps => "WRAP" + case _ => "NULL" + } + Seq(FunctionOption.builder().name("overflow").addValues(overflow).build()) + } + + def configure( + aggregate: AggregateFunction, + function: SimpleExtension.Function, + options: java.util.List[FunctionOption], + phase: SExpression.AggregationPhase, + outputType: Type): AggregateFunction = { + if (!isSparkFunction(function)) return aggregate + if (phase != SExpression.AggregationPhase.INITIAL_TO_RESULT) { + throw new UnsupportedOperationException( + s"Spark extension ${function.key()} only supports INITIAL_TO_RESULT, got $phase") + } + val unknown = options.asScala.filterNot(_.getName.equalsIgnoreCase("overflow")) + if (unknown.nonEmpty) { + throw new UnsupportedOperationException( + s"Unsupported Spark aggregate options: ${unknown.map(_.getName).mkString(", ")}") + } + val overflow = options.asScala.filter(_.getName.equalsIgnoreCase("overflow")) + if (overflow.size > 1) { + throw new IllegalArgumentException("Spark aggregate has duplicate overflow options") + } + val supported = + if (function.key() == "sum:i64") Seq("WRAP", "NULL", "ERROR") + else Seq("NULL", "ERROR") + val preferences = overflow.headOption + .map(_.values().asScala.toSeq) + .getOrElse(Seq(supported.head)) + val declared = function.options().get("overflow").getValues.asScala + val undeclared = preferences.filterNot(value => declared.exists(_.equalsIgnoreCase(value))) + if (undeclared.nonEmpty) { + throw new IllegalArgumentException( + s"Undeclared overflow preferences for Spark ${function.key()}: ${undeclared.mkString(", ")}") + } + val selected = preferences.iterator + .flatMap(value => supported.find(_.equalsIgnoreCase(value))) + .take(1) + .toSeq + .headOption + .getOrElse { + throw new UnsupportedOperationException( + s"Unsupported overflow preferences for Spark ${function.key()}: ${preferences.mkString(", ")}") + } + val mode = selected match { + case "ERROR" => EvalMode.ANSI + case "NULL" => EvalMode.TRY + case "WRAP" => EvalMode.LEGACY + } + val configured = aggregate match { + case sum: Sum => sum.copy(evalMode = mode) + case avg: Average => avg.copy(evalMode = mode) + case other => throw new UnsupportedOperationException(s"Unsupported Spark aggregate: $other") + } + validateTypes( + function, + configured.children.map(child => ToSubstraitType(child.dataType, child.nullable)), + outputType) + val actualType = ToSubstraitType(configured.dataType, configured.nullable) + if (actualType != outputType) { + throw new IllegalArgumentException( + s"Spark ${function.key()} declares $outputType but derives $actualType") + } + configured match { + case sum: Sum if sum.child.dataType.isInstanceOf[DecimalType] => + DecimalAggregate(sum) + case avg: Average if avg.child.dataType.isInstanceOf[DecimalType] => + DecimalAggregate(avg) + case other => other + } + } + + def unwrap(expression: AggregateExpression): AggregateExpression = + expression.aggregateFunction match { + case DecimalAggregate(delegate) => expression.copy(aggregateFunction = delegate) + case _ => expression + } + + // DecimalAggregates changes AVG arithmetic and rebuilds SUM with the consumer's overflow + // mode. Imported calls must retain their declared contracts under either optimizer setting. + private case class DecimalAggregate(delegate: DeclarativeAggregate) extends DeclarativeAggregate { + override def children: Seq[Expression] = delegate.children + override def nullable: Boolean = delegate.nullable + override def dataType = delegate.dataType + override lazy val aggBufferAttributes: Seq[AttributeReference] = delegate.aggBufferAttributes + override lazy val inputAggBufferAttributes: Seq[AttributeReference] = + delegate.inputAggBufferAttributes + override lazy val initialValues: Seq[Expression] = delegate.initialValues + override lazy val updateExpressions: Seq[Expression] = delegate.updateExpressions + override lazy val mergeExpressions: Seq[Expression] = delegate.mergeExpressions + override lazy val evaluateExpression: Expression = delegate.evaluateExpression + override protected def withNewChildrenInternal(children: IndexedSeq[Expression]): Expression = + copy(delegate = delegate.withNewChildren(children).asInstanceOf[DeclarativeAggregate]) + } +} diff --git a/spark/src/main/scala/io/substrait/spark/expression/ToAggregateFunction.scala b/spark/src/main/scala/io/substrait/spark/expression/ToAggregateFunction.scala index 5d243193a..09bfb4d91 100644 --- a/spark/src/main/scala/io/substrait/spark/expression/ToAggregateFunction.scala +++ b/spark/src/main/scala/io/substrait/spark/expression/ToAggregateFunction.scala @@ -38,8 +38,9 @@ abstract class ToAggregateFunction(functions: Seq[SimpleExtension.AggregateFunct outputType: Type): AggregateFunctionInvocation = { val sparkAggregate = sparkExp.asInstanceOf[AggregateExpression] + SparkAggregateFunctions.validateResult(function, arguments, outputType) - ExpressionCreator.aggregateFunction( + val invocation = ExpressionCreator.aggregateFunction( function, outputType, ToAggregateFunction.fromSpark(sparkAggregate.mode), @@ -47,13 +48,19 @@ abstract class ToAggregateFunction(functions: Seq[SimpleExtension.AggregateFunct ToAggregateFunction.fromSpark(sparkAggregate.isDistinct), arguments.asJava ) + AggregateFunctionInvocation + .builder() + .from(invocation) + .addAllOptions(SparkAggregateFunctions.options(sparkAggregate, function).asJava) + .build() } def convert( expression: AggregateExpression, operands: Seq[SExpression]): Option[AggregateFunctionInvocation] = { - Option(signatures.get(expression.aggregateFunction.getClass)) - .flatMap(m => m.attemptMatch(expression, operands)) + val native = SparkAggregateFunctions.unwrap(expression) + Option(signatures.get(native.aggregateFunction.getClass)) + .flatMap(m => m.attemptMatch(native, operands)) } def apply( diff --git a/spark/src/main/scala/io/substrait/spark/expression/ToScalarFunction.scala b/spark/src/main/scala/io/substrait/spark/expression/ToScalarFunction.scala index 239d0e034..60b880fd7 100644 --- a/spark/src/main/scala/io/substrait/spark/expression/ToScalarFunction.scala +++ b/spark/src/main/scala/io/substrait/spark/expression/ToScalarFunction.scala @@ -16,9 +16,11 @@ */ package io.substrait.spark.expression -import org.apache.spark.sql.catalyst.expressions.Expression +import io.substrait.spark.ToSubstraitType -import io.substrait.`type`.Type +import org.apache.spark.sql.catalyst.expressions.{Expression, MakeDecimal} + +import io.substrait.`type`.{Type, TypeCreator} import io.substrait.expression.{Expression => SExpression, FunctionArg} import io.substrait.extension.SimpleExtension @@ -32,6 +34,16 @@ abstract class ToScalarFunction(functions: Seq[SimpleExtension.ScalarFunctionVar function: SimpleExtension.ScalarFunctionVariant, arguments: Seq[FunctionArg], outputType: Type): SExpression = { + if (ToScalarFunction.isDecimalConstructor(function)) { + val types = arguments.map { + case e: SExpression => e.getType + case t: Type => t + case other => throw new IllegalArgumentException(s"Unexpected decimal argument: $other") + } + require( + function.resolveType(types.asJava) == outputType, + "Decimal constructor output type mismatch") + } SExpression.ScalarFunctionInvocation .builder() .outputType(outputType) @@ -41,12 +53,32 @@ abstract class ToScalarFunction(functions: Seq[SimpleExtension.ScalarFunctionVar } def convert(expression: Expression, operands: Seq[FunctionArg]): Option[SExpression] = { - Option(signatures.get(expression.getClass)) - .flatMap(m => m.attemptMatch(expression, operands)) + expression match { + case decimal: MakeDecimal => + val name = if (decimal.nullOnOverflow) "try_make_decimal" else "make_decimal" + functions + .find(f => f.urn() == SparkAggregateFunctions.URN && f.key() == s"$name:i64_type") + .map( + f => + generateBinding( + decimal, + f, + operands :+ TypeCreator.REQUIRED.decimal(decimal.precision, decimal.scale), + ToSubstraitType(decimal.dataType, decimal.nullable))) + case _ => + Option(signatures.get(expression.getClass)) + .flatMap(m => m.attemptMatch(expression, operands)) + } } } object ToScalarFunction { + private[spark] val decimalConstructorKeys = + Set("make_decimal:i64_type", "try_make_decimal:i64_type") + + private[spark] def isDecimalConstructor(function: SimpleExtension.Function): Boolean = + function.urn() == SparkAggregateFunctions.URN && decimalConstructorKeys.contains(function.key()) + def apply(functions: Seq[SimpleExtension.ScalarFunctionVariant]): ToScalarFunction = { new ToScalarFunction(functions) { override def getSigs: Seq[Sig] = FunctionMappings.SCALAR_SIGS diff --git a/spark/src/main/scala/io/substrait/spark/expression/ToSparkExpression.scala b/spark/src/main/scala/io/substrait/spark/expression/ToSparkExpression.scala index 749caae6a..093523d91 100644 --- a/spark/src/main/scala/io/substrait/spark/expression/ToSparkExpression.scala +++ b/spark/src/main/scala/io/substrait/spark/expression/ToSparkExpression.scala @@ -279,6 +279,26 @@ class ToSparkExpression( override def visit( expr: SExpression.ScalarFunctionInvocation, context: EmptyVisitationContext): Expression = { + if (ToScalarFunction.isDecimalConstructor(expr.declaration())) { + require(expr.options().isEmpty, "Decimal constructor options are not supported") + expr.arguments().asScala.toSeq match { + case Seq(value: SExpression, target: Type.Decimal) + if value.getType.isInstanceOf[Type.I64] => + val derived = expr.declaration().resolveType(java.util.List.of(value.getType, target)) + require(derived == expr.getType, "Decimal constructor output type mismatch") + val converted = MakeDecimal( + value.accept(this, context), + target.precision(), + target.scale(), + expr.declaration().name() == "try_make_decimal") + require( + io.substrait.spark + .ToSubstraitType(converted.dataType, converted.nullable) == expr.getType, + "Decimal constructor Spark type mismatch") + return converted + case _ => throw new IllegalArgumentException("Expected an i64 value and a decimal type") + } + } val eArgs = expr.arguments().asScala val args = eArgs.zipWithIndex.map { case (arg, i) => diff --git a/spark/src/main/scala/io/substrait/spark/expression/ToSubstraitExpression.scala b/spark/src/main/scala/io/substrait/spark/expression/ToSubstraitExpression.scala index 19ff39a2c..239705172 100644 --- a/spark/src/main/scala/io/substrait/spark/expression/ToSubstraitExpression.scala +++ b/spark/src/main/scala/io/substrait/spark/expression/ToSubstraitExpression.scala @@ -173,7 +173,6 @@ abstract class ToSubstraitExpression extends HasOutputStack[Seq[Attribute]] { translateUp(child) .map(ExpressionCreator .cast(ToSubstraitType.apply(dataType, c.nullable), _, FailureBehavior.THROW_EXCEPTION)) - case UnscaledValue(value) => translateUp(Cast(value, LongType)) case c @ CheckOverflow(child, dataType, _) => // CheckOverflow similar with cast translateUp(child) diff --git a/spark/src/main/scala/io/substrait/spark/expression/ToWindowFunction.scala b/spark/src/main/scala/io/substrait/spark/expression/ToWindowFunction.scala index c29746fb5..4eb4c375a 100644 --- a/spark/src/main/scala/io/substrait/spark/expression/ToWindowFunction.scala +++ b/spark/src/main/scala/io/substrait/spark/expression/ToWindowFunction.scala @@ -41,6 +41,8 @@ abstract class ToWindowFunction(functions: Seq[SimpleExtension.WindowFunctionVar arguments: Seq[FunctionArg], outputType: Type): WindowRelFunctionInvocation = { + SparkAggregateFunctions.validateResult(function, arguments, outputType) + val (frameType, lower, upper) = sparkExp match { case WindowExpression( _, @@ -56,7 +58,7 @@ abstract class ToWindowFunction(functions: Seq[SimpleExtension.WindowFunctionVar case _ => throw new UnsupportedOperationException(s"Unsupported window expression: $sparkExp") } - ExpressionCreator.windowRelFunction( + val invocation = ExpressionCreator.windowRelFunction( function, outputType, SExpression.AggregationPhase.INITIAL_TO_RESULT, // use defaults... @@ -66,18 +68,29 @@ abstract class ToWindowFunction(functions: Seq[SimpleExtension.WindowFunctionVar upper, arguments.asJava ) + val options = sparkExp match { + case WindowExpression(aggregate: AggregateExpression, _) => + SparkAggregateFunctions.options(aggregate, function) + case _ => Seq.empty + } + WindowRelFunctionInvocation.builder().from(invocation).addAllOptions(options.asJava).build() } def convert( expression: WindowExpression, operands: Seq[SExpression]): Option[WindowRelFunctionInvocation] = { - val cls = expression.windowFunction match { + val native = expression.windowFunction match { + case agg: AggregateExpression => + expression.copy(windowFunction = SparkAggregateFunctions.unwrap(agg)) + case _ => expression + } + val cls = native.windowFunction match { case agg: AggregateExpression => agg.aggregateFunction.getClass case other => other.getClass } Option(signatures.get(cls)) - .flatMap(m => m.attemptMatch(expression, operands)) + .flatMap(m => m.attemptMatch(native, operands)) } def apply( diff --git a/spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala b/spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala index 6e505fe5f..40f4487c4 100644 --- a/spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala +++ b/spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala @@ -92,7 +92,7 @@ class ToLogicalPlan(val spark: AnyRef = SparkCompat.instance.getOrCreateSparkSes arguments = ArrayBuffer(Literal(1)).toSeq } - val aggregateFunction = SparkExtension.toAggregateFunction + val mappedAggregateFunction = SparkExtension.toAggregateFunction .getSparkExpressionFromSubstraitFunc(function.declaration.key, arguments) .map(_.asInstanceOf[AggregateFunction]) .getOrElse({ @@ -111,6 +111,13 @@ class ToLogicalPlan(val spark: AnyRef = SparkCompat.instance.getOrCreateSparkSes throw new IllegalArgumentException(msg) }) + val aggregateFunction = SparkAggregateFunctions.configure( + mappedAggregateFunction, + function.declaration(), + function.options(), + function.aggregationPhase(), + function.outputType()) + if (function.sort().asScala.exists(!_.direction().isPresent)) { throw new UnsupportedOperationException( "A sort field using a custom comparison function is not supported") @@ -179,8 +186,14 @@ class ToLogicalPlan(val spark: AnyRef = SparkCompat.instance.getOrCreateSparkSes .map { case win: WindowFunction => win case agg: AggregateFunction => - AggregateExpression( + val configured = SparkAggregateFunctions.configure( agg, + func.declaration(), + func.options(), + func.aggregationPhase(), + func.outputType()) + AggregateExpression( + configured, ToAggregateFunction.toSpark(func.aggregationPhase()), ToAggregateFunction.toSpark(func.invocation()), None) diff --git a/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala b/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala index d50f1ca72..32b4ecc73 100644 --- a/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala +++ b/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala @@ -27,7 +27,7 @@ import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.analysis.ResolvedIdentifier import org.apache.spark.sql.catalyst.catalog.{CatalogTable, HiveTableRelation} import org.apache.spark.sql.catalyst.expressions._ -import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Average, Sum} +import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression import org.apache.spark.sql.catalyst.plans._ import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.execution.LogicalRDD @@ -158,30 +158,7 @@ class ToSubstraitRel extends AbstractLogicalPlanVisitor with Logging with Predic */ override def visitAggregate(agg: Aggregate): relation.Rel = { val input = visit(agg.child) - val actualResultExprs = agg.aggregateExpressions.map { - // eliminate the internal MakeDecimal and UnscaledValue functions by undoing the spark optimisation: - // https://github.com/apache/spark/blob/master/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala#L2223 - case Alias(expr, name) => - Alias( - expr.transform { - case MakeDecimal( - ae @ AggregateExpression(Sum(UnscaledValue(e), _), _, _, _, _), - _, - _, - _) => - ae.copy(aggregateFunction = Sum(e)) - case Cast( - Divide(ae @ AggregateExpression(Average(UnscaledValue(e), _), _, _, _, _), _, _), - _, - _, - _) => - ae.copy(aggregateFunction = Average(e)) - case e => e - }, - name - )() - case e => e - } + val actualResultExprs = agg.aggregateExpressions val actualGroupExprs = agg.groupingExpressions val aggExprToOutputOrdinal = mutable.HashMap.empty[Expression, Int] @@ -244,12 +221,21 @@ class ToSubstraitRel extends AbstractLogicalPlanVisitor with Logging with Predic } override def visitWindow(window: Window): relation.Rel = { - val windowExpressions = window.windowExpressions + val expressions = window.windowExpressions.map { + case alias @ Alias(MakeDecimal(w: WindowExpression, _, _, _), _) => + (w, alias: Expression, true) + case alias @ Alias(Cast(Divide(w: WindowExpression, _: Literal, _), _, _, _), _) + if w.windowFunction.isInstanceOf[AggregateExpression] => + (w, alias: Expression, true) + case w: WindowExpression => (w, w: Expression, false) + case alias @ Alias(w: WindowExpression, _) => (w, alias: Expression, false) + case other => + throw new UnsupportedOperationException(s"Unsupported window expression: $other") + } + val windowExpressions = expressions .map { - case w: WindowExpression => fromWindowCall(w, window.child.output) - case Alias(w: WindowExpression, _) => fromWindowCall(w, window.child.output) - case other => - throw new UnsupportedOperationException(s"Unsupported window expression: $other") + case (w, _, _) => + fromWindowCall(w, window.child.output) } .toSeq .asJava @@ -257,12 +243,33 @@ class ToSubstraitRel extends AbstractLogicalPlanVisitor with Logging with Predic val partitionExpressions = window.partitionSpec.map(toExpression(window.child.output)).toSeq.asJava val sorts = window.orderSpec.map(toSortField(window.child.output)).toSeq.asJava - relation.ConsistentPartitionWindow.builder + val windowRel = relation.ConsistentPartitionWindow.builder .input(visit(window.child)) .addAllWindowFunctions(windowExpressions) .addAllPartitionExpressions(partitionExpressions) .addAllSorts(sorts) .build() + if (!expressions.exists(_._3)) return windowRel + + val results = expressions.zipWithIndex.map { + case ((w, _, _), i) => + AttributeReference(s"window_func_$i", w.dataType, nullable = w.nullable)() + } + val output = window.child.output ++ results + val projects = window.child.output ++ expressions.zipWithIndex.map { + case ((w, original, _), i) => + original.transformDown { + case expression: WindowExpression if expression == w => results(i) + } + } + relation.Project.builder + .input(windowRel) + .remap(relation.Rel.Remap.offset(output.size, projects.size)) + .expressions(projects.map(toExpression(output)).toSeq.asJava) + .hint(Hint.builder + .addAllOutputNames(ToSubstraitType.toNamedStruct(window.schema).names()) + .build()) + .build() } private def asLong(e: Expression): Long = e match { diff --git a/spark/src/main/scala/io/substrait/spark/utils/DialectGenerator.scala b/spark/src/main/scala/io/substrait/spark/utils/DialectGenerator.scala index 5acdb4cdf..58c19bf87 100644 --- a/spark/src/main/scala/io/substrait/spark/utils/DialectGenerator.scala +++ b/spark/src/main/scala/io/substrait/spark/utils/DialectGenerator.scala @@ -1,8 +1,8 @@ package io.substrait.spark.utils import io.substrait.spark.SparkExtension +import io.substrait.spark.expression.{Sig, ToScalarFunction} import io.substrait.spark.expression.FunctionMappings.{AGGREGATE_SIGS, SCALAR_SIGS, WINDOW_SIGS} -import io.substrait.spark.expression.Sig import org.apache.spark.sql.catalyst.expressions.{BinaryOperator, Expression, Literal} import org.apache.spark.sql.types.{ByteType, DateType, DayTimeIntervalType, DoubleType, FloatType, IntegerType, LongType, ShortType, TimestampNTZType, TimestampType, YearMonthIntervalType} @@ -43,9 +43,24 @@ class DialectGenerator( } def generate(): Dialect = { - val scalars = SCALAR_SIGS.flatMap(supportedFunctions(scalarFunctions)) + // Type arguments have no Catalyst literal for the generic signature probe. + val decimalConstructors = scalarFunctions.filter(ToScalarFunction.isDecimalConstructor).map { + f => + SourcedFunction( + f.urn(), + DialectFunction + .builder() + .source(dependencyAlias(f.urn())) + .name(f.name()) + .systemMetadata( + SystemFunctionMetadata.builder().name(f.name()).notation(Notation.FUNCTION).build()) + .addSupportedImpls("i64_type") + .build() + ) + } + val scalars = SCALAR_SIGS.flatMap(supportedFunctions(scalarFunctions)) ++ decimalConstructors val aggregates = AGGREGATE_SIGS.flatMap(supportedFunctions(aggregateFunctions)) - val windows = WINDOW_SIGS.flatMap(supportedFunctions(windowFunctions)) + val windows = (WINDOW_SIGS ++ AGGREGATE_SIGS).flatMap(supportedFunctions(windowFunctions)) val builder = Dialect .builder() @@ -331,8 +346,8 @@ class DialectGenerator( object DialectGenerator extends DialectGenerator( SparkExtension.SparkScalarFunctions, - SparkExtension.StandardAggregateFunctions, - SparkExtension.StandardWindowFunctions) { + SparkExtension.AggregateFunctions, + SparkExtension.WindowFunctions) { def main(args: Array[String]): Unit = { val yaml = generateYaml() diff --git a/spark/src/test/scala/io/substrait/spark/DialectSuite.scala b/spark/src/test/scala/io/substrait/spark/DialectSuite.scala index 637bbde75..be6df2df3 100644 --- a/spark/src/test/scala/io/substrait/spark/DialectSuite.scala +++ b/spark/src/test/scala/io/substrait/spark/DialectSuite.scala @@ -46,9 +46,9 @@ class DialectSuite extends SparkFunSuite with SharedSparkSession with SubstraitP private def generatorWith(extras: SimpleExtension.ExtensionCollection*): DialectGenerator = new DialectGenerator( SparkExtension.SparkScalarFunctions, - SparkExtension.StandardAggregateFunctions ++ + SparkExtension.AggregateFunctions ++ extras.flatMap(_.aggregateFunctions().asScala.toSeq), - SparkExtension.StandardWindowFunctions + SparkExtension.WindowFunctions ) override def beforeAll(): Unit = { @@ -65,6 +65,19 @@ class DialectSuite extends SparkFunSuite with SharedSparkSession with SubstraitP assertResult(java.util.List.of())(errors) } + test("Spark function declarations satisfy the extension schema") { + val registry = SchemaRegistry.withDefaultDialect(SpecificationVersion.DRAFT_2020_12) + val schemaStream = getClass.getResourceAsStream("/substrait/text/simple_extensions_schema.yaml") + val schema = + try registry.getSchema(schemaStream, InputFormat.YAML) + finally schemaStream.close() + val source = Source.fromInputStream(getClass.getResourceAsStream("/spark.yml")) + val definitions = + try source.mkString + finally source.close() + assertResult(java.util.List.of())(schema.validate(definitions, InputFormat.YAML)) + } + test("the CLI writes the published dialect to the file it is given") { // `main` creates the file before it opens a writer on it, so asserting that it exists passes // whatever was written to it -- including nothing. Read the file back instead. @@ -162,26 +175,45 @@ class DialectSuite extends SparkFunSuite with SharedSparkSession with SubstraitP "'extension:io.substrait:functions_extra'")(error.getMessage) } - test("no aggregate or window function outside the standard extensions is advertised") { - // The runtime converters bind aggregates and windows against the standard extensions only, so - // a dialect derived from a wider set -- the collection merged with spark.yml -- would advertise - // functions that then fail with "Unable to find binding for call". Sourced from the catalog - // rather than from SparkExtension so that rewiring the generator to the merged collection is - // what this fails on. It cannot fail while spark.yml declares no aggregate or window function, - // which is exactly the condition that makes the bug latent rather than live. + test( + "aggregate and window declarations include the Spark contracts without stale standard variants") { val standard = DefaultExtensionCatalog.DEFAULT_COLLECTION - val declared: Set[(String, String)] = - (standard.aggregateFunctions().asScala.toSeq.map(f => (f.urn(), f.key())) ++ - standard.windowFunctions().asScala.toSeq.map(f => (f.urn(), f.key()))).toSet - + val sparkUrn = "extension:substrait:spark" + val sparkKeys = Set("sum:dec", "avg:dec", "sum:i64", "avg:i64") + val replaced = Set( + (DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC_DECIMAL, "sum:dec"), + (DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC_DECIMAL, "avg:dec"), + (DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, "sum:i64"), + (DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, "avg:i64") + ) + val declaredAggregate = + standard.aggregateFunctions().asScala.toSeq.map(f => (f.urn(), f.key())).toSet -- replaced ++ + sparkKeys.map(key => (sparkUrn, key)) + val declaredWindow = + standard.windowFunctions().asScala.toSeq.map(f => (f.urn(), f.key())).toSet -- replaced ++ + sparkKeys.map(key => (sparkUrn, key)) val dialect = DialectGenerator.generate() val urns = dialect.dependencies().asScala - val advertised = (dialect.supportedAggregateFunctions().asScala.toSeq ++ - dialect.supportedWindowFunctions().asScala.toSeq).flatMap { - f => f.supportedImpls().asScala.toSeq.map(impl => (urns(f.source()), s"${f.name()}:$impl")) + for ( + (functions, declared) <- Seq( + dialect.supportedAggregateFunctions() -> declaredAggregate, + dialect.supportedWindowFunctions() -> declaredWindow) + ) { + val advertised = functions.asScala.toSeq.flatMap { + f => f.supportedImpls().asScala.toSeq.map(impl => (urns(f.source()), s"${f.name()}:$impl")) + }.toSet + assert(advertised.subsetOf(declared)) + assert(sparkKeys.map(key => (sparkUrn, key)).subsetOf(advertised)) + assert(advertised.intersect(replaced).isEmpty) } - - assert(advertised.nonEmpty) - assertResult(Seq.empty)(advertised.filterNot(declared.contains)) + val windowAdvertised = dialect + .supportedWindowFunctions() + .asScala + .toSeq + .flatMap { + f => f.supportedImpls().asScala.toSeq.map(impl => (urns(f.source()), s"${f.name()}:$impl")) + } + .toSet + assert(windowAdvertised.contains((DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, "max:i64"))) } } diff --git a/spark/src/test/scala/io/substrait/spark/SparkAggregateFunctionsSuite.scala b/spark/src/test/scala/io/substrait/spark/SparkAggregateFunctionsSuite.scala new file mode 100644 index 000000000..c8b87a1f2 --- /dev/null +++ b/spark/src/test/scala/io/substrait/spark/SparkAggregateFunctionsSuite.scala @@ -0,0 +1,596 @@ +/* + * 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 io.substrait.spark + +import io.substrait.spark.compat.SparkCompat +import io.substrait.spark.expression.SparkAggregateFunctions +import io.substrait.spark.logical.{ToLogicalPlan, ToSubstraitRel} + +import org.apache.spark.{SparkFunSuite, SparkThrowable} +import org.apache.spark.sql.Row +import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow} +import org.apache.spark.sql.catalyst.expressions.{BoundReference, DecimalAddNoOverflowCheck, EvalMode, Literal, UnsafeProjection} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Average, Complete, Final, Partial, PartialMerge, Sum} +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{Decimal, DecimalType} + +import io.substrait.`type`.TypeCreator +import io.substrait.expression.{Expression => SExpression, ExpressionCreator, FunctionOption} +import io.substrait.extension.SimpleExtension +import io.substrait.plan.{PlanProtoConverter, ProtoPlanConverter} +import io.substrait.relation.{Aggregate, ConsistentPartitionWindow, Rel, SingleInputRel} + +import scala.annotation.nowarn +import scala.jdk.CollectionConverters._ + +class SparkAggregateFunctionsSuite extends SparkFunSuite with SharedSparkSession { + private val ansiKey = "spark.sql.ansi.enabled" + private val urn = "extension:substrait:spark" + private val input = Literal.create(Decimal(BigDecimal("1.00")), DecimalType(10, 2)) + private val sum = Sum(input, EvalMode.LEGACY) + private val declaration = SparkExtension.COLLECTION.getAggregateFunction( + SimpleExtension.FunctionAnchor.of(urn, "sum:dec")) + private val sumType = TypeCreator.NULLABLE.decimal(20, 2) + + private def option(values: String*): FunctionOption = + FunctionOption.builder().name("overflow").addAllValues(values.asJava).build() + + // Spark 4 deprecates this accessor; Spark 3 does not yet provide getCondition. + @nowarn("cat=deprecation") + private def errorClass(error: SparkThrowable): String = error.getErrorClass + + private def outcome(body: => Seq[Row]): Either[String, Seq[Row]] = { + try Right(body) + catch { + case error: Exception => + var cause: Throwable = error + while (cause.getCause != null) cause = cause.getCause + cause match { + case throwable: SparkThrowable => Left(errorClass(throwable)) + case other => throw other + } + } + } + + private def verifyTypes(rel: Rel): Unit = { + def verify( + function: SimpleExtension.Function, + arguments: java.util.List[io.substrait.expression.FunctionArg], + output: io.substrait.`type`.Type): Unit = { + if (function.urn() == urn) { + val actual = arguments.asScala.map(_.asInstanceOf[SExpression].getType).toList.asJava + assertResult(output)(function.resolveType(actual)) + } + } + rel match { + case aggregate: Aggregate => + aggregate.getMeasures.asScala.foreach { + measure => + val function = measure.getFunction + verify(function.declaration(), function.arguments(), function.outputType()) + } + case window: ConsistentPartitionWindow => + window.getWindowFunctions.asScala.foreach { + function => verify(function.declaration(), function.arguments(), function.outputType()) + } + case _ => + } + rel match { + case single: SingleInputRel => verifyTypes(single.getInput) + case _ => + } + } + + private val decimalRule = "org.apache.spark.sql.catalyst.optimizer.DecimalAggregates" + + private def roundTrip( + sql: String, + ansi: Boolean, + sourceDecimalOptimization: Boolean = true, + targetDecimalOptimization: Boolean = true, + repeatTransport: Boolean = false): Either[String, Seq[Row]] = { + var result: Option[Either[String, Seq[Row]]] = None + withSQLConf( + ansiKey -> ansi.toString, + "spark.sql.optimizer.excludedRules" -> (if (sourceDecimalOptimization) "" else decimalRule)) { + val native = spark.sql(sql) + val exported = new ToSubstraitRel().convert(native.queryExecution.optimizedPlan) + val proto = new PlanProtoConverter().toProto(exported) + val parsed = new ProtoPlanConverter(SparkExtension.COLLECTION) + .from(io.substrait.proto.Plan.parseFrom(proto.toByteArray)) + parsed.getRoots.asScala.foreach(root => verifyTypes(root.getInput)) + val unsupported = intercept[IllegalArgumentException] { + new ProtoPlanConverter().from(proto) + } + assert(unsupported.getMessage.contains(urn)) + val expected = outcome(native.collect().toSeq) + withSQLConf( + ansiKey -> (!ansi).toString, + "spark.sql.optimizer.excludedRules" -> (if (targetDecimalOptimization) "" + else decimalRule)) { + val imported = new ToLogicalPlan(spark).convert(parsed) + val restored = if (repeatTransport) { + val optimized = SparkCompat.instance.createQueryExecution(spark, imported).optimizedPlan + val again = new PlanProtoConverter().toProto(new ToSubstraitRel().convert(optimized)) + new ToLogicalPlan(spark).convert( + new ProtoPlanConverter(SparkExtension.COLLECTION).from(again)) + } else imported + assertResult(native.schema)(restored.schema) + val actual = outcome { + val qe = SparkCompat.instance.createQueryExecution(spark, restored) + val convert = CatalystTypeConverters.createToScalaConverter(restored.schema) + qe.executedPlan.executeCollect().map(row => convert(row).asInstanceOf[Row]).toSeq + } + assertResult(expected)(actual) + } + result = Some(expected) + } + result.get + } + + private def values(p: Int, s: Int, rows: String*): String = + " FROM VALUES " + rows + .map { + value => + val literal = if (value == null) "NULL" else s"'$value'" + s"(CAST($literal AS DECIMAL($p,$s)))" + } + .mkString(",") + " AS t(d)" + + test("decimal aggregates preserve values and types across ANSI modes") { + for (ansi <- Seq(false, true)) { + val ordinary = values(10, 2, "1.00", "2.01", null) + assert(roundTrip("SELECT sum(d) AS s, avg(d) AS a" + ordinary, ansi).isRight) + roundTrip("SELECT sum(DISTINCT d) AS s, avg(DISTINCT d) AS a" + ordinary, ansi) + roundTrip( + "SELECT sum(d) FILTER (WHERE d > 1) AS s, avg(d) FILTER (WHERE d > 1) AS a" + ordinary, + ansi) + roundTrip("SELECT sum(d) AS s, avg(d) AS a" + values(10, 2, null), ansi) + roundTrip("SELECT sum(d) AS s, avg(d) AS a" + ordinary + " WHERE false", ansi) + val large = "9" + "0" * 31 + ".00" + roundTrip("SELECT sum(d) AS s, avg(d) AS a" + values(34, 2, large, large), ansi) + val cancel = "1" + "0" * 34 + ".00" + roundTrip("SELECT avg(d) AS a" + values(38, 2, cancel, "-" + cancel), ansi) + } + } + + test("AVG preserves source arithmetic when producer and consumer optimizers differ") { + withTable("decimal_avg_rounding") { + spark.sql( + "CREATE TABLE decimal_avg_rounding USING parquet AS " + + "SELECT CAST(CASE WHEN id=0 THEN '0.03' ELSE '0.00' END AS DECIMAL(10,2)) d " + + "FROM range(20000)") + for ( + ansi <- Seq(false, true); + optimized <- Seq(false, true); + query <- Seq( + "SELECT avg(d) AS a FROM decimal_avg_rounding", + "SELECT avg(d) OVER () AS a FROM decimal_avg_rounding LIMIT 1", + "SELECT sum(d) OVER () AS s, avg(d) OVER () AS a FROM decimal_avg_rounding LIMIT 1" + ) + ) { + val result = roundTrip(query, ansi, optimized, !optimized, repeatTransport = true) + val average = new java.math.BigDecimal(if (optimized) "0.000001" else "0.000002") + val expected = if (query.contains("sum(d)")) { + Row(new java.math.BigDecimal("0.03"), average) + } else Row(average) + assertResult(Right(Seq(expected)))(result) + } + } + } + + test("decimal AVG retains HALF_UP rounding at the precision limit") { + val tiny = "0." + "0" * 37 + "1" + for (ansi <- Seq(false, true); value <- Seq(tiny, "-" + tiny)) { + val result = roundTrip("SELECT avg(d) AS a" + values(38, 38, value, "0"), ansi) + assertResult(Right(Seq(Row(new java.math.BigDecimal(value)))))(result) + val below = roundTrip("SELECT avg(d) AS a" + values(38, 38, value, "0", "0"), ansi) + assertResult(Right(Seq(Row(new java.math.BigDecimal("0").setScale(38)))))(below) + } + } + + test("decimal overflow preserves NULL versus ERROR, including try aggregates") { + val overflow = values(38, 0, "9" * 38, "9" * 38) + for (ansi <- Seq(false, true); function <- Seq("sum", "avg")) { + val result = roundTrip(s"SELECT $function(d) AS a" + overflow, ansi) + if (ansi) assert(result.isLeft) + else assertResult(Right(Seq(Row(null))))(result) + assertResult(Right(Seq(Row(null))))( + roundTrip(s"SELECT try_$function(d) AS a" + overflow, ansi)) + } + } + + test("decimal SUM keeps its overflow mode under the consumer optimizer") { + withTable("decimal_sum_mode") { + spark.sql( + "CREATE TABLE decimal_sum_mode USING parquet AS " + + "SELECT CAST(id AS DECIMAL(8,0)) d FROM range(2)") + for (ansi <- Seq(false, true); function <- Seq("sum", "try_sum")) { + withSQLConf(ansiKey -> ansi.toString, "spark.sql.optimizer.excludedRules" -> decimalRule) { + val expectedMode = if (ansi && function == "sum") EvalMode.ANSI else EvalMode.TRY + val expectedOption = if (expectedMode == EvalMode.ANSI) "ERROR" else "NULL" + for (window <- Seq(false, true)) { + val over = if (window) " OVER ()" else "" + val native = spark.sql(s"SELECT $function(d)$over AS s FROM decimal_sum_mode") + val proto = new PlanProtoConverter().toProto( + new ToSubstraitRel().convert(native.queryExecution.optimizedPlan)) + val parsed = new ProtoPlanConverter(SparkExtension.COLLECTION).from(proto) + withSQLConf(ansiKey -> (!ansi).toString, "spark.sql.optimizer.excludedRules" -> "") { + val imported = new ToLogicalPlan(spark).convert(parsed) + val optimized = + SparkCompat.instance.createQueryExecution(spark, imported).optimizedPlan + val calls = optimized.flatMap(_.expressions.flatMap(_.collect { + case expression: AggregateExpression => SparkAggregateFunctions.unwrap(expression) + })) + assert(calls.nonEmpty) + calls.foreach { + call => + val delegate = call.aggregateFunction.asInstanceOf[Sum] + assertResult(expectedMode)(delegate.evalMode) + // Exercise finalization at an overflowing buffer value without billions of rows. + val finalized = delegate.evaluateExpression.transformDown { + case attribute: org.apache.spark.sql.catalyst.expressions.AttributeReference + if attribute == delegate.aggBufferAttributes.head => + Literal.create(Decimal(BigDecimal("1000000000000000000")), attribute.dataType) + case attribute: org.apache.spark.sql.catalyst.expressions.AttributeReference + if attribute == delegate.aggBufferAttributes(1) => + Literal(false) + } + if (expectedMode == EvalMode.ANSI) + intercept[ArithmeticException](finalized.eval()) + else assertResult(null)(finalized.eval()) + assertResult(Seq(option(expectedOption)))( + SparkAggregateFunctions.options(call, declaration)) + } + // Re-export the optimized consumer plan, including its protected aggregate. + new ToSubstraitRel().convert(optimized) + } + } + } + } + } + } + + test("decimal window buffers can recover after a null result until materialized") { + val max = "9" * 38 + val rows = Seq(max, max, "-" + max, "-" + max, "1").zipWithIndex + .map { case (value, i) => s"($i, CAST('$value' AS DECIMAL(38,0)))" } + .mkString(",") + val result = roundTrip( + "SELECT k, sum(d) OVER (ORDER BY k ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) s, " + + "avg(d) OVER (ORDER BY k ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) a " + + s"FROM VALUES $rows AS t(k,d) ORDER BY k", + ansi = false, + repeatTransport = true + ) + def decimal(value: String) = new java.math.BigDecimal(value) + assertResult( + Right( + Seq( + Row(0, decimal(max), null), + Row(1, null, null), + Row(2, decimal(max), null), + Row(3, decimal("0"), decimal("0.0000")), + Row(4, decimal("1"), decimal("0.2000")) + )))(result) + + val t = DecimalType(38, 0) + val value = Literal.create(Decimal(BigDecimal(max)), t) + val live = DecimalAddNoOverflowCheck(value, value, t).eval().asInstanceOf[Decimal] + assertResult(new java.math.BigDecimal(max).multiply(new java.math.BigDecimal("2")))( + live.toBigDecimal.bigDecimal) + val materialize = UnsafeProjection.create(Seq(BoundReference(0, t, nullable = true))) + assert(materialize(InternalRow(live)).isNullAt(0)) + } + + test("window SUM and AVG carry their types and overflow modes") { + val ordinary = values(10, 2, "1.00", "2.01", "3.00") + for ( + ansi <- Seq(false, true); + frame <- Seq( + "ROWS BETWEEN 1 PRECEDING AND CURRENT ROW", + "RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW") + ) { + roundTrip( + s"SELECT d, sum(d) OVER (ORDER BY d $frame) AS s, avg(d) OVER (ORDER BY d $frame) AS a" + ordinary + " ORDER BY d", + ansi) + } + for (ansi <- Seq(false, true)) { + val overflow = values(38, 0, "9" * 38, "9" * 38) + val result = roundTrip("SELECT sum(d) OVER () AS s, avg(d) OVER () AS a" + overflow, ansi) + if (ansi) assert(result.isLeft) + else assertResult(Right(Seq(Row(null, null), Row(null, null))))(result) + } + } + + test("i64 optimization paths retain wrapping, errors, try_sum, and floating AVG") { + val input = + " FROM VALUES (CAST('9223372036854775807' AS BIGINT)), (CAST('1' AS BIGINT)) AS t(d)" + for (ansi <- Seq(false, true)) { + val sum = roundTrip("SELECT sum(d) AS s" + input, ansi) + if (ansi) assert(sum.isLeft) + else assertResult(Right(Seq(Row(Long.MinValue))))(sum) + roundTrip("SELECT avg(d) AS a" + input, ansi) + assertResult(Right(Seq(Row(null))))(roundTrip("SELECT try_sum(d) AS s" + input, ansi)) + val window = roundTrip("SELECT try_sum(d) OVER () AS s" + input, ansi) + assertResult(Right(Seq(Row(null), Row(null))))(window) + } + } + + test("decimal aggregate optimization paths preserve try overflow options") { + def functions(rel: Rel): Seq[io.substrait.expression.AggregateFunctionInvocation] = { + val here = rel match { + case aggregate: Aggregate => aggregate.getMeasures.asScala.map(_.getFunction).toSeq + case _ => Seq.empty + } + rel match { + case single: SingleInputRel => here ++ functions(single.getInput) + case _ => here + } + } + withSQLConf(ansiKey -> "true") { + val sql = "SELECT try_sum(d) AS s, try_avg(d) AS a" + values(5, 2, "1.00", "2.00") + // Some Spark versions lose TRY in DecimalAggregates. Supply a plan that still carries it. + val plan = spark.sql(sql).queryExecution.optimizedPlan.transformAllExpressions { + case expression: AggregateExpression => + expression.aggregateFunction match { + case aggregate: Sum => + expression.copy(aggregateFunction = aggregate.copy(evalMode = EvalMode.TRY)) + case aggregate: Average => + expression.copy(aggregateFunction = aggregate.copy(evalMode = EvalMode.TRY)) + case _ => expression + } + } + val exported = new ToSubstraitRel().convert(plan) + val aggregates = exported.getRoots.asScala.toSeq.flatMap(root => functions(root.getInput)) + assert(aggregates.nonEmpty) + aggregates.foreach { + function => + assertResult(urn)(function.declaration().urn()) + assertResult(Seq(option("NULL")))(function.options().asScala.toSeq) + } + } + } + + test("optimized SUM retains wrapping i64 merges before its final decimal conversion") { + withTable("decimal_sum_wrap") { + spark.sql( + "CREATE TABLE decimal_sum_wrap USING parquet AS SELECT CAST(id AS DECIMAL(8,0)) d FROM range(2)") + withSQLConf(ansiKey -> "false") { + val native = spark.sql("SELECT sum(d) s FROM decimal_sum_wrap") + val source = native.queryExecution.optimizedPlan + val proto = new PlanProtoConverter().toProto(new ToSubstraitRel().convert(source)) + val imported = new ToLogicalPlan(spark).convert( + new ProtoPlanConverter(SparkExtension.COLLECTION).from(proto)) + def sums(plan: org.apache.spark.sql.catalyst.plans.logical.LogicalPlan): Seq[Sum] = + plan + .flatMap(_.expressions.flatMap(_.collect { + case ae: AggregateExpression => ae.aggregateFunction + })) + .collect { case sum: Sum => sum } + def merge(sum: Sum, left: Long, right: Long): Long = { + assertResult(org.apache.spark.sql.types.LongType)(sum.dataType) + sum.mergeExpressions.head + .transformDown { + case a: org.apache.spark.sql.catalyst.expressions.AttributeReference + if a == sum.aggBufferAttributes.head => + Literal(left) + case a: org.apache.spark.sql.catalyst.expressions.AttributeReference + if a == sum.inputAggBufferAttributes.head => + Literal(right) + } + .eval() + .asInstanceOf[Long] + } + val sourceSum = sums(source).head + val importedSum = sums(imported).head + // Valid partial i64 sums can total 2^64 even when each input decimal fits precision 8. + // Exercise those buffer merges without allocating billions of input rows. + for (sum <- Seq(sourceSum, importedSum)) { + val total = merge(sum, merge(sum, Long.MaxValue, Long.MaxValue), 2L) + assertResult(0L)(total) + assertResult(Decimal(0))( + org.apache.spark.sql.catalyst.expressions + .MakeDecimal(Literal(total), 18, 0, true) + .eval()) + } + } + } + } + + test("unmodified optimized TRY plans retain the source plan's actual overflow mode") { + withTable("decimal_try_source") { + spark.sql( + "CREATE TABLE decimal_try_source USING parquet AS SELECT CAST(id AS DECIMAL(8,2)) d FROM range(2)") + for (ansi <- Seq(false, true)) { + roundTrip( + "SELECT try_sum(d) s, try_avg(d) a FROM decimal_try_source", + ansi, + repeatTransport = true) + roundTrip( + "SELECT try_sum(d) OVER () s, try_avg(d) OVER () a FROM decimal_try_source", + ansi, + repeatTransport = true) + } + withSQLConf(ansiKey -> "true") { + val query = spark.sql("SELECT try_sum(d) s FROM decimal_try_source") + def calls(plan: org.apache.spark.sql.catalyst.plans.logical.LogicalPlan) = + plan.flatMap(_.expressions.flatMap(_.collect { case ae: AggregateExpression => ae })) + assertResult(EvalMode.TRY)( + calls(query.queryExecution.analyzed).head.aggregateFunction.asInstanceOf[Sum].evalMode) + val optimized = query.queryExecution.optimizedPlan + val sourceMode = calls(optimized).head.aggregateFunction.asInstanceOf[Sum].evalMode + def measures(rel: Rel): Seq[io.substrait.expression.AggregateFunctionInvocation] = { + val here = rel match { + case a: Aggregate => a.getMeasures.asScala.map(_.getFunction).toSeq + case _ => Seq.empty + } + rel match { + case one: SingleInputRel => here ++ measures(one.getInput) + case _ => here + } + } + val exported = new ToSubstraitRel().convert(optimized) + val invocation = exported.getRoots.asScala.toSeq.flatMap(r => measures(r.getInput)).head + val expected = sourceMode match { + case EvalMode.ANSI => "ERROR" + case EvalMode.LEGACY => "WRAP" + case EvalMode.TRY => "NULL" + } + assertResult("sum:i64")(invocation.declaration().key()) + // The converter receives the optimized plan, so it must export that plan's mode. + assertResult(Seq(option(expected)))(invocation.options().asScala.toSeq) + } + } + } + + test("unscaled decimal constructors retain scale, nullability and overflow behavior") { + val scalar = io.substrait.spark.expression.ToScalarFunction(SparkExtension.SparkScalarFunctions) + val converter = new io.substrait.spark.expression.ToSparkExpression(scalar, None) + for ( + nullOnOverflow <- Seq(false, true); + raw <- Seq[java.lang.Long](123L, -123L, Long.MaxValue, null) + ) { + val child = Literal.create(raw, org.apache.spark.sql.types.LongType) + val native = + org.apache.spark.sql.catalyst.expressions.MakeDecimal(child, 3, 2, nullOnOverflow) + val arg = + if (raw == null) ExpressionCreator.typedNull(TypeCreator.NULLABLE.I64) + else ExpressionCreator.i64(false, raw.longValue()) + val exported = scalar.convert(native, Seq(arg)).get + val restored = exported.accept(converter, io.substrait.util.EmptyVisitationContext.INSTANCE) + assertResult(native.nullable)(restored.nullable) + if (raw != null && raw.longValue() == Long.MaxValue && !nullOnOverflow) { + intercept[ArithmeticException](native.eval()) + intercept[ArithmeticException](restored.eval()) + } else assertResult(native.eval())(restored.eval()) + } + } + + test("unscaled_value preserves fractional digits, NULL, and i64 overflow errors") { + val scalar = io.substrait.spark.expression.ToScalarFunction(SparkExtension.SparkScalarFunctions) + val converter = new io.substrait.spark.expression.ToSparkExpression(scalar, None) + for (value <- Seq("1.23", "-1.23", "123456789012345678901234567890123456.78", null)) { + val literal = + Literal.create(if (value == null) null else Decimal(BigDecimal(value)), DecimalType(38, 2)) + val native = org.apache.spark.sql.catalyst.expressions.UnscaledValue(literal) + val argument = + if (value == null) ExpressionCreator.typedNull(TypeCreator.NULLABLE.decimal(38, 2)) + else ExpressionCreator.decimal(false, new java.math.BigDecimal(value), 38, 2) + val exported = scalar.convert(native, Seq(argument)).get + val restored = exported.accept(converter, io.substrait.util.EmptyVisitationContext.INSTANCE) + if (value != null && new java.math.BigDecimal(value).unscaledValue().bitLength() > 63) { + intercept[ArithmeticException](native.eval()) + intercept[ArithmeticException](restored.eval()) + } else assertResult(native.eval())(restored.eval()) + } + } + + test("export refuses a Spark aggregate declaration with a different result type") { + val incompatible = SimpleExtension.load(s"""urn: $urn + |aggregate_functions: + | - name: sum + | impls: + | - args: + | - value: DECIMAL + | nullability: DECLARED_OUTPUT + | decomposable: NONE + | return: DECIMAL?<38,S> + |""".stripMargin) + val converter = io.substrait.spark.expression + .ToAggregateFunction(incompatible.aggregateFunctions().asScala.toSeq) + intercept[IllegalArgumentException] { + converter( + AggregateExpression( + sum, + org.apache.spark.sql.catalyst.expressions.aggregate.Complete, + isDistinct = false, + filter = None), + Seq(ExpressionCreator.decimal(false, new java.math.BigDecimal("1.00"), 10, 2)) + ) + } + } + + test("Spark extension aggregates reject staged export and import") { + for (mode <- Seq(Partial, PartialMerge, Final)) { + val expression = AggregateExpression(sum, mode, isDistinct = false, filter = None) + val error = intercept[UnsupportedOperationException] { + SparkExtension.toAggregateFunction( + expression, + Seq(ExpressionCreator.decimal(false, new java.math.BigDecimal("1.00"), 10, 2))) + } + assert(error.getMessage.contains("complete aggregation")) + } + for ( + phase <- SExpression.AggregationPhase + .values() + .filterNot(_ == SExpression.AggregationPhase.INITIAL_TO_RESULT) + ) { + val error = intercept[UnsupportedOperationException] { + SparkAggregateFunctions.configure(sum, declaration, java.util.List.of(), phase, sumType) + } + assert(error.getMessage.contains("INITIAL_TO_RESULT")) + } + } + + test("overflow preference order, defaults, and unsupported declarations") { + def configured(options: FunctionOption*): Sum = { + val aggregate = SparkAggregateFunctions.configure( + sum, + declaration, + options.toList.asJava, + SExpression.AggregationPhase.INITIAL_TO_RESULT, + sumType) + SparkAggregateFunctions + .unwrap(AggregateExpression(aggregate, Complete, isDistinct = false, filter = None)) + .aggregateFunction + .asInstanceOf[Sum] + } + assertResult(EvalMode.TRY)(configured().evalMode) + assertResult(EvalMode.ANSI)(configured(option("ERROR", "NULL")).evalMode) + assertResult(EvalMode.ANSI)(configured(option("error", "null")).evalMode) + assertResult(EvalMode.ANSI)( + configured(FunctionOption.builder().name("OVERFLOW").addValues("ERROR").build()).evalMode) + assertResult(EvalMode.TRY)( + configured( + FunctionOption.builder().name("OvErFlOw").addValues("nUlL", "eRrOr").build()).evalMode) + assertResult(EvalMode.TRY)(configured(option("NULL", "ERROR")).evalMode) + intercept[IllegalArgumentException](configured(option("WRAP"))) + for (values <- Seq(Seq("FUTURE", "ERROR"), Seq("ERROR", "FUTURE"))) { + val error = intercept[IllegalArgumentException](configured(option(values: _*))) + assert(error.getMessage.contains("Undeclared overflow preferences")) + } + intercept[UnsupportedOperationException](configured(option())) + intercept[IllegalArgumentException](configured(option("NULL"), option("ERROR"))) + intercept[IllegalArgumentException]( + configured( + option("NULL"), + FunctionOption.builder().name("OVERFLOW").addValues("ERROR").build())) + intercept[UnsupportedOperationException] { + configured(FunctionOption.builder().name("unknown").addValues("x").build()) + } + val mismatch = intercept[IllegalArgumentException] { + SparkAggregateFunctions.configure( + sum, + declaration, + java.util.List.of(), + SExpression.AggregationPhase.INITIAL_TO_RESULT, + TypeCreator.NULLABLE.decimal(38, 2)) + } + assert(mismatch.getMessage.contains("derives")) + } +}