[spark] supports converting some SparkPredicate to Paimon between LeafPredicate

Author: steFaiz

paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/SparkV2FilterConverter.scala

@@ -18,14 +18,16 @@
 
 package org.apache.paimon.spark
 
-import org.apache.paimon.predicate.{Predicate, PredicateBuilder, Transform}
+import org.apache.paimon.predicate.{GreaterOrEqual, LeafPredicate, LessOrEqual, Predicate, PredicateBuilder, Transform}
 import org.apache.paimon.spark.util.SparkExpressionConverter.{toPaimonLiteral, toPaimonTransform}
 import org.apache.paimon.types.RowType
 
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.connector.expressions.{Expression, Literal}
 import org.apache.spark.sql.connector.expressions.filter.{And, Not, Or, Predicate => SparkPredicate}
 
+import java.util.Objects
+
 import scala.collection.JavaConverters._
 
 /** Conversion from [[SparkPredicate]] to [[Predicate]]. */
@@ -125,7 +127,12 @@ case class SparkV2FilterConverter(rowType: RowType) extends Logging {
 
       case AND =>
         val and = sparkPredicate.asInstanceOf[And]
-        PredicateBuilder.and(convert(and.left), convert(and.right()))
+        val leftPredicate = convert(and.left)
+        val rightPredicate = convert(and.right())
+        convertToBetweenFunction(leftPredicate, rightPredicate) match {
+          case Some(predicate) => predicate
+          case _ => PredicateBuilder.and(leftPredicate, rightPredicate)
+        }
 
       case OR =>
         val or = sparkPredicate.asInstanceOf[Or]
@@ -169,6 +176,42 @@ case class SparkV2FilterConverter(rowType: RowType) extends Logging {
     }
   }
 
+  private def convertToBetweenFunction(
+      leftPredicate: Predicate,
+      rightPredicate: Predicate): Option[Predicate] = {
+    def toBetweenLeafPredicate(
+        transform: Transform,
+        lowerBoundInclusive: Object,
+        upperBoundInclusive: Object): Predicate = {
+      builder.between(transform, lowerBoundInclusive, upperBoundInclusive)
+    }
+
+    (leftPredicate, rightPredicate) match {
+      case (left: LeafPredicate, right: LeafPredicate) =>
+        // left and right should have the same transform
+        if (!Objects.equals(left.transform(), right.transform())) {
+          return None
+        }
+        (left.function(), right.function()) match {
+          case (_: GreaterOrEqual, _: LessOrEqual) =>
+            Some(
+              toBetweenLeafPredicate(
+                left.transform(),
+                left.literals().get(0),
+                right.literals().get(0)))
+          case (_: LessOrEqual, _: GreaterOrEqual) =>
+            Some(
+              toBetweenLeafPredicate(
+                left.transform(),
+                right.literals().get(0),
+                left.literals().get(0)))
+          case _ => None
+        }
+      case _ =>
+        None
+    }
+  }
+
   private object UnaryPredicate {
     def unapply(sparkPredicate: SparkPredicate): Option[Transform] = {
       sparkPredicate.children() match {

paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/RowIdPushDownTestBase.scala

@@ -65,6 +65,10 @@ abstract class RowIdPushDownTestBase extends PaimonSparkTestBase {
         sql("SELECT * FROM t WHERE _ROW_ID IN (6, 7)"),
         Seq()
       )
+      checkAnswer(
+        sql("SELECT * FROM t WHERE _ROW_ID BETWEEN 0 AND 2"),
+        Seq(Row(0, 0, "0"), Row(1, 1, "1"), Row(2, 2, "2"))
+      )
 
       // 2.CompoundPredicate
       checkAnswer(

paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/SparkV2FilterConverterTestBase.scala

@@ -19,7 +19,7 @@
 package org.apache.paimon.spark.sql
 
 import org.apache.paimon.data.{BinaryString, Decimal, Timestamp}
-import org.apache.paimon.predicate.PredicateBuilder
+import org.apache.paimon.predicate.{Between, LeafPredicate, PredicateBuilder}
 import org.apache.paimon.spark.{PaimonSparkTestBase, SparkV2FilterConverter}
 import org.apache.paimon.spark.util.shim.TypeUtils.treatPaimonTimestampTypeAsSparkTimestampType
 import org.apache.paimon.table.source.DataSplit
@@ -295,6 +295,39 @@ abstract class SparkV2FilterConverterTestBase extends PaimonSparkTestBase {
     assert(scanFilesCount(filter) == 3)
   }
 
+  test("V2Filter: Between") {
+    // 1. test basic between
+    val filter = "string_col BETWEEN 'a' AND 'r'"
+    val actual = converter.convert(v2Filter(filter)).get
+    assert(
+      actual.equals(builder.between(0, BinaryString.fromString("a"), BinaryString.fromString("r"))))
+    checkAnswer(
+      sql(s"SELECT string_col from test_tbl WHERE $filter ORDER BY string_col"),
+      Seq(Row("hello"), Row("hi"), Row("paimon"))
+    )
+
+    // 2. >= and <= on same transform should also be converted to between
+    val filter1 = "CONCAT(string_col, '_suffix') >= 'a' AND CONCAT(string_col, '_suffix') <= 'r'"
+    val actual1 = converter.convert(v2Filter(filter1)).get
+    assert(actual1.isInstanceOf[LeafPredicate])
+    val function = actual1.asInstanceOf[LeafPredicate].function
+    assert(function.isInstanceOf[Between])
+    checkAnswer(
+      sql(s"SELECT string_col from test_tbl WHERE $filter1 ORDER BY string_col"),
+      Seq(Row("hello"), Row("hi"), Row("paimon"))
+    )
+
+    // 3. >= and <= on different transform should not be converted to between
+    val filter2 = "CONCAT(string_col, '_suffix1') >= 'a' AND CONCAT(string_col, '_suffix2') <= 'r'"
+    val actual2 = converter.convert(v2Filter(filter2)).get
+    assert(!actual2.isInstanceOf[LeafPredicate])
+
+    // 4. >= and <= on different columns should not be converted to between
+    val filter3 = "string_col >= 'a' AND int_col <= 2"
+    val actual3 = converter.convert(v2Filter(filter3)).get
+    assert(!actual3.isInstanceOf[LeafPredicate])
+  }
+
   test("V2Filter: And") {
     val filter = "int_col > 1 AND int_col < 3"
     val actual = converter.convert(v2Filter(filter)).get