From 04612dc9932c2602d909ff38ad37bf17f680a598 Mon Sep 17 00:00:00 2001 From: Tim Meehan Date: Thu, 10 Sep 2026 16:27:26 +0000 Subject: [PATCH] [SQL] Derive pushdown predicates from ANSI integral arithmetic --- .../sql/catalyst/optimizer/Optimizer.scala | 1 + .../sql/catalyst/optimizer/expressions.scala | 128 ++++++++++++++++++ .../sql/catalyst/rules/RuleIdCollection.scala | 1 + .../BinaryComparisonSimplificationSuite.scala | 44 +++++- .../parquet/ParquetFilterSuite.scala | 32 ++++- 5 files changed, 204 insertions(+), 2 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala index f1484ff02e154..e3440f81da9ac 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala @@ -143,6 +143,7 @@ abstract class Optimizer(catalogManager: CatalogManager) BooleanSimplification, SimplifyConditionals, PushFoldableIntoBranches, + DeriveIntegralComparisonPredicates, SimplifyBinaryComparison, ReplaceNullWithFalseInPredicate, PruneFilters, diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala index cdc0444d74cc1..1846de1fdfe66 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala @@ -567,6 +567,134 @@ object BooleanSimplification extends Rule[LogicalPlan] with PredicateHelper { } +/** + * Derives predicates over integral operands from comparisons involving ANSI addition or + * subtraction. The original comparison is retained to preserve overflow errors and exact + * evaluation semantics. + */ +object DeriveIntegralComparisonPredicates extends Rule[LogicalPlan] { + private val derived = TreeNodeTag[Unit]("derived_integral_comparison_predicate") + + private case class IntegralRange(min: BigInt, max: BigInt, literal: BigInt => Literal) + + private def integralRange(dataType: DataType): Option[IntegralRange] = dataType match { + case ByteType => Some(IntegralRange(Byte.MinValue, Byte.MaxValue, + value => Literal(value.toByte))) + case ShortType => Some(IntegralRange(Short.MinValue, Short.MaxValue, + value => Literal(value.toShort))) + case IntegerType => Some(IntegralRange(Int.MinValue, Int.MaxValue, + value => Literal(value.toInt))) + case LongType => Some(IntegralRange(Long.MinValue, Long.MaxValue, + value => Literal(value.toLong))) + case _ => None + } + + private def integralValue(literal: Literal): Option[BigInt] = literal.value match { + case value: Byte => Some(BigInt(value)) + case value: Short => Some(BigInt(value)) + case value: Int => Some(BigInt(value)) + case value: Long => Some(BigInt(value)) + case _ => None + } + + private def normalizedComparison( + comparison: BinaryComparison): Option[(Expression, Literal, BinaryComparison)] = { + comparison match { + case EqualTo(arithmetic, literal: Literal) => Some((arithmetic, literal, comparison)) + case EqualTo(literal: Literal, arithmetic) => + Some((arithmetic, literal, EqualTo(arithmetic, literal))) + case LessThan(arithmetic, literal: Literal) => Some((arithmetic, literal, comparison)) + case LessThan(literal: Literal, arithmetic) => + Some((arithmetic, literal, GreaterThan(arithmetic, literal))) + case LessThanOrEqual(arithmetic, literal: Literal) => + Some((arithmetic, literal, comparison)) + case LessThanOrEqual(literal: Literal, arithmetic) => + Some((arithmetic, literal, GreaterThanOrEqual(arithmetic, literal))) + case GreaterThan(arithmetic, literal: Literal) => Some((arithmetic, literal, comparison)) + case GreaterThan(literal: Literal, arithmetic) => + Some((arithmetic, literal, LessThan(arithmetic, literal))) + case GreaterThanOrEqual(arithmetic, literal: Literal) => + Some((arithmetic, literal, comparison)) + case GreaterThanOrEqual(literal: Literal, arithmetic) => + Some((arithmetic, literal, LessThanOrEqual(arithmetic, literal))) + case _ => None + } + } + + private def arithmeticOperand( + expression: Expression): Option[(Expression, Literal, BigInt)] = expression match { + case add @ Add(operand, literal: Literal, _) + if add.evalMode == EvalMode.ANSI && operand.deterministic => + integralValue(literal).map((operand, literal, _)) + case add @ Add(literal: Literal, operand, _) + if add.evalMode == EvalMode.ANSI && operand.deterministic => + integralValue(literal).map((operand, literal, _)) + case subtract @ Subtract(operand, literal: Literal, _) + if subtract.evalMode == EvalMode.ANSI && operand.deterministic => + integralValue(literal).map(value => (operand, literal, -value)) + case _ => None + } + + private def comparison( + template: BinaryComparison, + left: Expression, + right: Expression): BinaryComparison = template match { + case _: EqualTo => EqualTo(left, right) + case _: LessThan => LessThan(left, right) + case _: LessThanOrEqual => LessThanOrEqual(left, right) + case _: GreaterThan => GreaterThan(left, right) + case _: GreaterThanOrEqual => GreaterThanOrEqual(left, right) + } + + private def comparisonOutsideRange( + template: BinaryComparison, + thresholdIsBelowRange: Boolean): Expression = template match { + case _: EqualTo => FalseLiteral + case _: LessThan | _: LessThanOrEqual => Literal(!thresholdIsBelowRange) + case _: GreaterThan | _: GreaterThanOrEqual => Literal(thresholdIsBelowRange) + } + + private def derivedPredicate(comparisonExpression: BinaryComparison): Option[Expression] = { + for { + (arithmetic, comparisonLiteral, normalized) <- normalizedComparison(comparisonExpression) + (operand, arithmeticLiteral, delta) <- arithmeticOperand(arithmetic) + range <- integralRange(arithmetic.dataType) + comparisonValue <- integralValue(comparisonLiteral) + if operand.dataType == arithmetic.dataType + if arithmeticLiteral.dataType == arithmetic.dataType + if comparisonLiteral.dataType == arithmetic.dataType + } yield { + val threshold = comparisonValue - delta + val algebraic = if (threshold < range.min || threshold > range.max) { + comparisonOutsideRange(normalized, threshold < range.min) + } else { + comparison(normalized, operand, range.literal(threshold)) + } + val overflow = if (delta > 0) { + Some(GreaterThan(operand, range.literal(range.max - delta))) + } else if (delta < 0) { + Some(LessThan(operand, range.literal(range.min - delta))) + } else { + None + } + overflow.map(Or(algebraic, _)).getOrElse(algebraic) + } + } + + override def apply(plan: LogicalPlan): LogicalPlan = plan.transformWithPruning( + _.containsPattern(BINARY_COMPARISON), ruleId) { + case logicalPlan: LogicalPlan => + logicalPlan.transformExpressionsDownWithPruning(_.containsPattern(BINARY_COMPARISON)) { + case comparison: BinaryComparison if comparison.getTagValue(derived).isEmpty => + derivedPredicate(comparison).filterNot(_ == TrueLiteral).map { predicate => + comparison.setTagValue(derived, ()) + And(predicate, comparison) + }.getOrElse(comparison) + } + } +} + + /** * Simplifies binary comparisons with semantically-equal expressions: * 1) Replace '<=>' with 'true' literal. diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala index cd1d14e1bd6c1..ca2f72d777089 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala @@ -139,6 +139,7 @@ object RuleIdCollection { "org.apache.spark.sql.catalyst.optimizer.ConvertToLocalRelation" :: "org.apache.spark.sql.catalyst.optimizer.CostBasedJoinReorder" :: "org.apache.spark.sql.catalyst.optimizer.DecimalAggregates" :: + "org.apache.spark.sql.catalyst.optimizer.DeriveIntegralComparisonPredicates" :: "org.apache.spark.sql.catalyst.optimizer.EliminateAggregateFilter" :: "org.apache.spark.sql.catalyst.optimizer.EliminateLimits" :: "org.apache.spark.sql.catalyst.optimizer.EliminateMapObjects" :: diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/BinaryComparisonSimplificationSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/BinaryComparisonSimplificationSuite.scala index 2bbdb4e689193..6627d9ff94886 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/BinaryComparisonSimplificationSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/BinaryComparisonSimplificationSuite.scala @@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.plans.PlanTest import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.catalyst.rules._ import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{BooleanType, IntegerType, StructField, StructType} +import org.apache.spark.sql.types.{BooleanType, DoubleType, IntegerType, StructField, StructType} class BinaryComparisonSimplificationSuite extends PlanTest { @@ -42,6 +42,7 @@ class BinaryComparisonSimplificationSuite extends PlanTest { NullPropagation, ConstantFolding, BooleanSimplification, + DeriveIntegralComparisonPredicates, SimplifyBinaryComparison, PruneFilters) :: Nil } @@ -53,6 +54,47 @@ class BinaryComparisonSimplificationSuite extends PlanTest { val nonNullableRelation = LocalRelation($"a".int.withNullability(false)) val boolRelation = LocalRelation($"a".boolean, $"b".boolean) + test("derive pruning predicates from ANSI integral arithmetic comparisons") { + val a = nonNullableRelation.output.head + val add = Add(a, Literal(10), EvalMode.ANSI) + val subtract = Subtract(a, Literal(10), EvalMode.ANSI) + val addOverflow = a > Literal(Int.MaxValue - 10) + val subtractOverflow = a < Literal(Int.MinValue + 10) + + val cases = Seq[(Expression, Expression)]( + (add < Literal(100), (a < Literal(90) || addOverflow) && add < Literal(100)), + (add <= Literal(100), (a <= Literal(90) || addOverflow) && add <= Literal(100)), + (add > Literal(100), (a > Literal(90) || addOverflow) && add > Literal(100)), + (add >= Literal(100), (a >= Literal(90) || addOverflow) && add >= Literal(100)), + (add === Literal(100), (a === Literal(90) || addOverflow) && add === Literal(100)), + (add < Literal(Int.MinValue), addOverflow && add < Literal(Int.MinValue)), + (add > Literal(Int.MaxValue), addOverflow && add > Literal(Int.MaxValue)), + (subtract > Literal(100), + (a > Literal(110) || subtractOverflow) && subtract > Literal(100)), + (subtract < Literal(Int.MinValue), + subtractOverflow && subtract < Literal(Int.MinValue)), + (subtract > Literal(Int.MaxValue), + subtractOverflow && subtract > Literal(Int.MaxValue)), + (Literal(100) < add, (a > Literal(90) || addOverflow) && Literal(100) < add)) + + cases.foreach { case (input, expected) => + checkCondition(nonNullableRelation, input, expected) + } + } + + test("do not derive pruning predicates when arithmetic is not checked integral arithmetic") { + val a = nonNullableRelation.output.head + val cases = Seq( + Add(a, Literal(10), EvalMode.LEGACY) > Literal(100), + Multiply(a, Literal(10), EvalMode.ANSI) > Literal(100), + Add(Cast(a, DoubleType), Literal(10.0), EvalMode.ANSI) > Literal(100.0), + Add(a, a, EvalMode.ANSI) > Literal(100)) + + cases.foreach { condition => + checkCondition(nonNullableRelation, condition, condition) + } + } + test("Preserve nullable exprs when constraintPropagation is false") { withSQLConf(SQLConf.CONSTRAINT_PROPAGATION_ENABLED.key -> "false") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala index 9b851134fc0b2..39e49d421998e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala @@ -36,7 +36,7 @@ import org.apache.parquet.hadoop.{ParquetFileReader, ParquetInputFormat, Parquet import org.apache.parquet.hadoop.util.HadoopInputFile import org.apache.parquet.schema.{MessageType, MessageTypeParser} -import org.apache.spark.{SparkConf, SparkException, SparkRuntimeException} +import org.apache.spark.{SparkArithmeticException, SparkConf, SparkException, SparkRuntimeException} import org.apache.spark.sql._ import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.catalyst.expressions._ @@ -2942,6 +2942,36 @@ class ParquetV1FilterSuite extends ParquetFilterSuite { @ExtendedSQLTest class ParquetV2FilterSuite extends ParquetFilterSuite { import testImplicits.ColumnConstructorExt + import testImplicits.toRichColumn + + test("push down derived predicate for ANSI integral arithmetic") { + withParquetDataFrame((1 to 4).map(Tuple1(_))) { df => + val predicate = Add(df("_1").expr, Literal(10), EvalMode.ANSI) > Literal(100) + checkFilterPredicate( + df, + predicate, + classOf[Operators.Or], + checkAnswer(_, _: Seq[Row]), + Seq.empty) + } + } + + test("derived predicate preserves ANSI integral overflow") { + withParquetDataFrame(Seq(Tuple1(Int.MaxValue))) { df => + val predicate = Add(df("_1").expr, Literal(10), EvalMode.ANSI) > Literal(100) + checkError( + exception = intercept[SparkArithmeticException] { + df.filter(Column(predicate)).collect() + }, + condition = "ARITHMETIC_OVERFLOW", + parameters = Map( + "message" -> "overflow", + "alternative" -> " Use 'try_add' to tolerate overflow and return NULL instead.", + "config" -> s""""${SQLConf.ANSI_ENABLED.key}""""), + sqlState = "22003", + context = ExpectedContext("", -1, -1)) + } + } // TODO: enable Parquet V2 write path after file source V2 writers are workable. override protected def sparkConf: SparkConf =