paimon-spark supports row id push down

Author: zhaoshu

paimon-common/src/main/java/org/apache/paimon/predicate/RowIdPredicateVisitor.java

@@ -0,0 +1,103 @@
+/*
+ * 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 org.apache.paimon.predicate;
+
+import java.util.HashSet;
+import java.util.Set;
+
+import static org.apache.paimon.table.SpecialFields.ROW_ID;
+
+/**
+ * The {@link PredicateVisitor} to extract a list of Row IDs from predicates. The returned Row IDs
+ * can be pushed down to manifest readers and file readers to enable efficient random access.
+ *
+ * <p>Note that there is a significant distinction between returning {@code null} and returning an
+ * empty set:
+ *
+ * <ul>
+ *   <li>{@code null} indicates that the predicate cannot be converted into a random-access pattern,
+ *       meaning the filter is not consumable by this visitor.
+ *   <li>An empty set indicates that no rows satisfy the predicate (e.g. {@code WHERE _ROW_ID = 3
+ *       AND _ROW_ID IN (1, 2)}).
+ * </ul>
+ */
+public class RowIdPredicateVisitor implements PredicateVisitor<Set<Long>> {
+
+    @Override
+    public Set<Long> visit(LeafPredicate predicate) {
+        if (ROW_ID.name().equals(predicate.fieldName())) {
+            LeafFunction function = predicate.function();
+            if (function instanceof Equal || function instanceof In) {
+                HashSet<Long> rowIds = new HashSet<>();
+                for (Object literal : predicate.literals()) {
+                    rowIds.add((Long) literal);
+                }
+                return rowIds;
+            }
+        }
+        return null;
+    }
+
+    @Override
+    public Set<Long> visit(CompoundPredicate predicate) {
+        CompoundPredicate.Function function = predicate.function();
+        HashSet<Long> rowIds = null;
+        // `And` means we should get the intersection of all children.
+        if (function instanceof And) {
+            for (Predicate child : predicate.children()) {
+                Set<Long> childSet = child.visit(this);
+                if (childSet == null) {
+                    return null;
+                }
+
+                if (rowIds == null) {
+                    rowIds = new HashSet<>(childSet);
+                } else {
+                    rowIds.retainAll(childSet);
+                }
+
+                // shortcut for intersection
+                if (rowIds.isEmpty()) {
+                    return rowIds;
+                }
+            }
+        } else if (function instanceof Or) {
+            // `Or` means we should get the union of all children
+            rowIds = new HashSet<>();
+            for (Predicate child : predicate.children()) {
+                Set<Long> childSet = child.visit(this);
+                if (childSet == null) {
+                    return null;
+                }
+
+                rowIds.addAll(childSet);
+            }
+        } else {
+            // unexpected function type, just return null
+            return null;
+        }
+        return rowIds;
+    }
+
+    @Override
+    public Set<Long> visit(TransformPredicate predicate) {
+        // do not support transform predicate now.
+        return null;
+    }
+}

paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/paimon/spark/PaimonBasePushDown.scala

@@ -18,14 +18,16 @@
 
 package org.apache.paimon.spark
 
-import org.apache.paimon.predicate.{PartitionPredicateVisitor, Predicate}
-import org.apache.paimon.types.RowType
+import org.apache.paimon.predicate.{PartitionPredicateVisitor, Predicate, RowIdPredicateVisitor}
+import org.apache.paimon.types.{DataField, DataTypes, RowType}
 
 import org.apache.spark.sql.connector.read.SupportsPushDownFilters
 import org.apache.spark.sql.sources.Filter
 
-import java.util.{List => JList}
+import java.lang.{Long => JLong}
+import java.util.{ArrayList, List => JList}
 
+import scala.collection.JavaConverters._
 import scala.collection.mutable
 
 /** Base trait for Paimon scan filter push down. */
@@ -35,6 +37,7 @@ trait PaimonBasePushDown extends SupportsPushDownFilters {
   protected var rowType: RowType
 
   private var pushedSparkFilters = Array.empty[Filter]
+  protected var pushedRowIds: Array[JLong] = null
   protected var pushedPaimonPredicates: Array[Predicate] = Array.empty
   protected var reservedFilters: Array[Filter] = Array.empty
   protected var hasPostScanPredicates = false
@@ -45,30 +48,50 @@ trait PaimonBasePushDown extends SupportsPushDownFilters {
    * must be interpreted as ANDed together.
    */
   override def pushFilters(filters: Array[Filter]): Array[Filter] = {
-    val pushable = mutable.ArrayBuffer.empty[(Filter, Predicate)]
+    val pushableSparkFilters = mutable.ArrayBuffer.empty[Filter]
+    val pushablePaimonPredicates = mutable.ArrayBuffer.empty[Predicate]
+    var pushableRowIds: mutable.Set[JLong] = null
     val postScan = mutable.ArrayBuffer.empty[Filter]
     val reserved = mutable.ArrayBuffer.empty[Filter]
 
-    val converter = new SparkFilterConverter(rowType)
-    val visitor = new PartitionPredicateVisitor(partitionKeys)
+    val dataFieldWithRowId = new ArrayList[DataField](rowType.getFields)
+    dataFieldWithRowId.add(new DataField(rowType.getFieldCount, "_ROW_ID", DataTypes.BIGINT()))
+    val rowTypeWithRowId = rowType.copy(dataFieldWithRowId)
+    val converter = new SparkFilterConverter(rowTypeWithRowId)
+    val partitionVisitor = new PartitionPredicateVisitor(partitionKeys)
+    val rowIdVisitor = new RowIdPredicateVisitor
+
     filters.foreach {
       filter =>
         val predicate = converter.convertIgnoreFailure(filter)
         if (predicate == null) {
           postScan.append(filter)
         } else {
-          pushable.append((filter, predicate))
-          if (predicate.visit(visitor)) {
+          if (predicate.visit(partitionVisitor)) {
+            pushableSparkFilters.append(filter)
+            pushablePaimonPredicates.append(predicate)
             reserved.append(filter)
+          } else if (predicate.visit(rowIdVisitor) != null) {
+            pushableSparkFilters.append(filter)
+            if (pushableRowIds == null) {
+              pushableRowIds = predicate.visit(rowIdVisitor).asScala
+            } else {
+              pushableRowIds.retain(predicate.visit(rowIdVisitor).asScala)
+            }
           } else {
             postScan.append(filter)
           }
         }
     }
 
-    if (pushable.nonEmpty) {
-      this.pushedSparkFilters = pushable.map(_._1).toArray
-      this.pushedPaimonPredicates = pushable.map(_._2).toArray
+    if (pushableSparkFilters.nonEmpty) {
+      this.pushedSparkFilters = pushableSparkFilters.toArray
+    }
+    if (pushablePaimonPredicates.nonEmpty) {
+      this.pushedPaimonPredicates = pushablePaimonPredicates.toArray
+    }
+    if (pushableRowIds != null) {
+      this.pushedRowIds = pushableRowIds.toArray
     }
     if (reserved.nonEmpty) {
       this.reservedFilters = reserved.toArray

paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/paimon/spark/PaimonScan.scala

@@ -27,12 +27,15 @@ import org.apache.spark.sql.connector.read.SupportsRuntimeFiltering
 import org.apache.spark.sql.sources.{Filter, In}
 import org.apache.spark.sql.types.StructType
 
+import java.lang.{Long => JLong}
+
 import scala.collection.JavaConverters._
 
 case class PaimonScan(
     table: InnerTable,
     requiredSchema: StructType,
     filters: Seq[Predicate],
+    override val rowIds: Seq[JLong],
     reservedFilters: Seq[Filter],
     override val pushDownLimit: Option[Int],
     // no usage, just for compile compatibility

paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/paimon/spark/PaimonScanBuilder.scala

@@ -34,6 +34,13 @@ class PaimonScanBuilder(table: InnerTable)
   override protected var rowType: RowType = table.rowType()
 
   override def build(): Scan = {
-    PaimonScan(table, requiredSchema, pushedPaimonPredicates, reservedFilters, None, pushDownTopN)
+    PaimonScan(
+      table,
+      requiredSchema,
+      pushedPaimonPredicates,
+      pushedRowIds,
+      reservedFilters,
+      None,
+      pushDownTopN)
   }
 }

paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/paimon/spark/PaimonSplitScanBuilder.scala

@@ -24,6 +24,6 @@ import org.apache.spark.sql.connector.read.Scan
 
 class PaimonSplitScanBuilder(table: KnownSplitsTable) extends PaimonScanBuilder(table) {
   override def build(): Scan = {
-    PaimonSplitScan(table, table.splits(), requiredSchema, pushedPaimonPredicates)
+    PaimonSplitScan(table, table.splits(), requiredSchema, pushedPaimonPredicates, pushedRowIds)
   }
 }

paimon-spark/paimon-spark-3.2/src/test/scala/org/apache/paimon/spark/sql/RowIdPushDownTest.scala

@@ -0,0 +1,21 @@
+/*
+ * 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 org.apache.paimon.spark.sql
+
+class RowIdPushDownTest extends RowIdPushDownTestBase {}

paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/paimon/spark/PaimonScan.scala

@@ -29,12 +29,15 @@ import org.apache.spark.sql.connector.read.partitioning.{KeyGroupedPartitioning,
 import org.apache.spark.sql.sources.{Filter, In}
 import org.apache.spark.sql.types.StructType
 
+import java.lang.{Long => JLong}
+
 import scala.collection.JavaConverters._
 
 case class PaimonScan(
     table: InnerTable,
     requiredSchema: StructType,
     filters: Seq[Predicate],
+    override val rowIds: Seq[JLong],
     reservedFilters: Seq[Filter],
     override val pushDownLimit: Option[Int],
     override val pushDownTopN: Option[TopN],

paimon-spark/paimon-spark-3.3/src/test/scala/org/apache/paimon/spark/sql/RowIdPushDownTest.scala

@@ -0,0 +1,21 @@
+/*
+ * 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 org.apache.paimon.spark.sql
+
+class RowIdPushDownTest extends RowIdPushDownTestBase {}

paimon-spark/paimon-spark-3.4/src/test/scala/org/apache/paimon/spark/sql/RowIdPushDownTest.scala

@@ -0,0 +1,21 @@
+/*
+ * 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 org.apache.paimon.spark.sql
+
+class RowIdPushDownTest extends RowIdPushDownTestBase {}

paimon-spark/paimon-spark-3.5/src/test/scala/org/apache/paimon/spark/sql/RowIdPushDownTest.scala

@@ -0,0 +1,21 @@
+/*
+ * 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 org.apache.paimon.spark.sql
+
+class RowIdPushDownTest extends RowIdPushDownTestBase {}