[SQL] Derive pushdown predicates from ANSI integral arithmetic

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,

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.

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" ::

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") {

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 =