support pushdown limit

Author: Yohahaha

fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussAppendPartitionReader.scala

@@ -28,9 +28,10 @@ class FlussAppendPartitionReader(
     tablePath: TablePath,
     projection: Array[Int],
     pushedPredicate: Option[Predicate],
+    limit: Option[Int],
     flussPartition: FlussAppendInputPartition,
     flussConfig: Configuration)
-  extends FlussPartitionReader(tablePath, flussConfig) {
+  extends FlussPartitionReader(tablePath, flussConfig, limit) {
 
   override protected lazy val projectedRowType: RowType = rowType.project(projection)
 

fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussBatch.scala

@@ -39,6 +39,7 @@ abstract class FlussBatch(
     tablePath: TablePath,
     tableInfo: TableInfo,
     readSchema: StructType,
+    limit: Option[Int],
     flussConfig: Configuration)
   extends Batch
   with AutoCloseable {
@@ -114,9 +115,10 @@ class FlussAppendBatch(
     readSchema: StructType,
     pushedPredicate: Option[Predicate],
     partitionPredicate: Option[Predicate],
+    limit: Option[Int],
     options: CaseInsensitiveStringMap,
     flussConfig: Configuration)
-  extends FlussBatch(tablePath, tableInfo, readSchema, flussConfig) {
+  extends FlussBatch(tablePath, tableInfo, readSchema, limit, flussConfig) {
 
   override val startOffsetsInitializer: OffsetsInitializer = {
     FlussOffsetInitializers.startOffsetsInitializer(options, flussConfig)
@@ -202,6 +204,7 @@ class FlussAppendBatch(
       tablePath,
       projection,
       pushedPredicate,
+      limit,
       options,
       flussConfig)
   }
@@ -214,9 +217,10 @@ class FlussUpsertBatch(
     tableInfo: TableInfo,
     readSchema: StructType,
     partitionPredicate: Option[Predicate],
+    limit: Option[Int],
     options: CaseInsensitiveStringMap,
     flussConfig: Configuration)
-  extends FlussBatch(tablePath, tableInfo, readSchema, flussConfig) {
+  extends FlussBatch(tablePath, tableInfo, readSchema, limit, flussConfig) {
 
   override val startOffsetsInitializer: OffsetsInitializer = {
     val offsetsInitializer = FlussOffsetInitializers.startOffsetsInitializer(options, flussConfig)
@@ -253,6 +257,6 @@ class FlussUpsertBatch(
   }
 
   override def createReaderFactory(): PartitionReaderFactory = {
-    new FlussUpsertPartitionReaderFactory(tablePath, projection, options, flussConfig)
+    new FlussUpsertPartitionReaderFactory(tablePath, projection, limit, options, flussConfig)
   }
 }

fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussMicroBatchStream.scala

@@ -271,7 +271,7 @@ class FlussAppendMicroBatchStream(
     checkpointLocation) {
 
   override def createReaderFactory(): PartitionReaderFactory = {
-    new FlussAppendPartitionReaderFactory(tablePath, projection, None, options, flussConfig)
+    new FlussAppendPartitionReaderFactory(tablePath, projection, None, None, options, flussConfig)
   }
 
   override def planInputPartitions(start: Offset, end: Offset): Array[InputPartition] = {
@@ -352,6 +352,6 @@ class FlussUpsertMicroBatchStream(
   }
 
   override def createReaderFactory(): PartitionReaderFactory = {
-    new FlussUpsertPartitionReaderFactory(tablePath, projection, options, flussConfig)
+    new FlussUpsertPartitionReaderFactory(tablePath, projection, None, options, flussConfig)
   }
 }

fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussPartitionReader.scala

@@ -34,7 +34,10 @@ import org.apache.spark.sql.connector.read.PartitionReader
 
 import java.time.Duration
 
-abstract class FlussPartitionReader(tablePath: TablePath, flussConfig: Configuration)
+abstract class FlussPartitionReader(
+    tablePath: TablePath,
+    flussConfig: Configuration,
+    limit: Option[Int])
   extends PartitionReader[InternalRow]
   with Logging {
 
@@ -57,6 +60,9 @@ abstract class FlussPartitionReader(tablePath: TablePath, flussConfig: Configura
   def next0(): Boolean
 
   override def next(): Boolean = {
+    if (limit.isDefined && numRowsRead >= limit.get) {
+      return false
+    }
     val hasNext = next0()
     if (hasNext) {
       numRowsRead += 1

fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussPartitionReaderFactory.scala

@@ -30,6 +30,7 @@ class FlussAppendPartitionReaderFactory(
     tablePath: TablePath,
     projection: Array[Int],
     pushedPredicate: Option[Predicate],
+    limit: Option[Int],
     options: CaseInsensitiveStringMap,
     flussConfig: Configuration)
   extends PartitionReaderFactory {
@@ -40,6 +41,7 @@ class FlussAppendPartitionReaderFactory(
       tablePath,
       projection,
       pushedPredicate,
+      limit,
       flussPartition,
       flussConfig
     )
@@ -50,6 +52,7 @@ class FlussAppendPartitionReaderFactory(
 class FlussUpsertPartitionReaderFactory(
     tablePath: TablePath,
     projection: Array[Int],
+    limit: Option[Int],
     options: CaseInsensitiveStringMap,
     flussConfig: Configuration)
   extends PartitionReaderFactory {
@@ -59,6 +62,7 @@ class FlussUpsertPartitionReaderFactory(
     new FlussUpsertPartitionReader(
       tablePath,
       projection,
+      limit,
       upsertPartition,
       flussConfig
     )

fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussScan.scala

@@ -43,6 +43,8 @@ trait FlussScan extends Scan {
 
   def partitionPredicate: Option[FlussPredicate] = None
 
+  def limit: Option[Int] = None
+
   protected def scanType: String
 
   override def readSchema(): StructType = {
@@ -54,10 +56,14 @@ trait FlussScan extends Scan {
     val withPushed =
       if (pushedSparkPredicates.isEmpty) base
       else s"$base [PushedPredicates: ${pushedSparkPredicates.mkString("[", ", ", "]")}]"
-    partitionPredicate match {
+    val withPartition = partitionPredicate match {
       case Some(p) => s"$withPushed [PartitionFilter: $p]"
       case None => withPushed
     }
+    limit match {
+      case Some(l) => s"$withPartition [Limit: $l]"
+      case None => withPartition
+    }
   }
 
   override def supportedCustomMetrics(): Array[CustomMetric] =
@@ -72,6 +78,7 @@ case class FlussAppendScan(
     pushedPredicate: Option[FlussPredicate],
     override val partitionPredicate: Option[FlussPredicate],
     override val pushedSparkPredicates: Seq[Predicate],
+    override val limit: Option[Int],
     options: CaseInsensitiveStringMap,
     flussConfig: Configuration)
   extends FlussScan {
@@ -85,6 +92,7 @@ case class FlussAppendScan(
       readSchema,
       pushedPredicate,
       partitionPredicate,
+      limit,
       options,
       flussConfig)
   }
@@ -107,6 +115,7 @@ case class FlussLakeAppendScan(
     requiredSchema: Option[StructType],
     pushedPredicate: Option[FlussPredicate],
     override val pushedSparkPredicates: Seq[Predicate],
+    override val limit: Option[Int],
     options: CaseInsensitiveStringMap,
     flussConfig: Configuration)
   extends FlussScan {
@@ -119,6 +128,7 @@ case class FlussLakeAppendScan(
       tableInfo,
       readSchema,
       pushedPredicate,
+      limit,
       options,
       flussConfig)
   }
@@ -140,14 +150,22 @@ case class FlussUpsertScan(
     tableInfo: TableInfo,
     requiredSchema: Option[StructType],
     override val partitionPredicate: Option[FlussPredicate],
+    override val limit: Option[Int],
     options: CaseInsensitiveStringMap,
     flussConfig: Configuration)
   extends FlussScan {
 
   override protected val scanType: String = "Upsert"
 
   override def toBatch: Batch = {
-    new FlussUpsertBatch(tablePath, tableInfo, readSchema, partitionPredicate, options, flussConfig)
+    new FlussUpsertBatch(
+      tablePath,
+      tableInfo,
+      readSchema,
+      partitionPredicate,
+      limit,
+      options,
+      flussConfig)
   }
 
   override def toMicroBatchStream(checkpointLocation: String): MicroBatchStream = {
@@ -168,6 +186,7 @@ case class FlussLakeUpsertScan(
     requiredSchema: Option[StructType],
     pushedPredicate: Option[FlussPredicate],
     override val pushedSparkPredicates: Seq[Predicate],
+    override val limit: Option[Int],
     options: CaseInsensitiveStringMap,
     flussConfig: Configuration)
   extends FlussScan {
@@ -180,6 +199,7 @@ case class FlussLakeUpsertScan(
       tableInfo,
       readSchema,
       pushedPredicate,
+      limit,
       options,
       flussConfig)
   }

fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussScanBuilder.scala

@@ -24,7 +24,7 @@ import org.apache.fluss.spark.read.lake.{FlussLakeBatch, FlussLakeUtils}
 import org.apache.fluss.spark.utils.{SparkPartitionPredicate, SparkPredicateConverter}
 
 import org.apache.spark.sql.connector.expressions.filter.Predicate
-import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, SupportsPushDownRequiredColumns, SupportsPushDownV2Filters}
+import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, SupportsPushDownLimit, SupportsPushDownRequiredColumns, SupportsPushDownV2Filters}
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.util.CaseInsensitiveStringMap
 
@@ -33,13 +33,22 @@ import java.util.{Collections, IdentityHashMap, Set => JSet}
 import scala.collection.JavaConverters._
 
 /** An interface that extends from Spark [[ScanBuilder]]. */
-trait FlussScanBuilder extends ScanBuilder with SupportsPushDownRequiredColumns {
+trait FlussScanBuilder
+  extends ScanBuilder
+  with SupportsPushDownRequiredColumns
+  with SupportsPushDownLimit {
 
   protected var requiredSchema: Option[StructType] = None
+  protected var limit: Option[Int] = None
 
   override def pruneColumns(requiredSchema: StructType): Unit = {
     this.requiredSchema = Some(requiredSchema)
   }
+
+  override def pushLimit(limit: Int): Boolean = {
+    this.limit = Some(limit)
+    true
+  }
 }
 
 /** Extracts a partition-key predicate so the scan can skip partitions that can't match. */
@@ -132,6 +141,7 @@ class FlussAppendScanBuilder(
       pushedPredicate,
       partitionPredicate,
       acceptedPredicates.toSeq,
+      limit,
       options,
       flussConfig)
   }
@@ -152,6 +162,7 @@ class FlussLakeAppendScanBuilder(
       requiredSchema,
       pushedPredicate,
       acceptedPredicates.toSeq,
+      limit,
       options,
       flussConfig)
   }
@@ -166,7 +177,14 @@ class FlussUpsertScanBuilder(
   extends FlussSupportsPushDownPartitionFilters {
 
   override def build(): Scan = {
-    FlussUpsertScan(tablePath, tableInfo, requiredSchema, partitionPredicate, options, flussConfig)
+    FlussUpsertScan(
+      tablePath,
+      tableInfo,
+      requiredSchema,
+      partitionPredicate,
+      limit,
+      options,
+      flussConfig)
   }
 }
 
@@ -185,6 +203,7 @@ class FlussLakeUpsertScanBuilder(
       requiredSchema,
       pushedPredicate,
       acceptedPredicates.toSeq,
+      limit,
       options,
       flussConfig)
   }

fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussUpsertPartitionReader.scala

@@ -46,9 +46,10 @@ import scala.collection.mutable
 class FlussUpsertPartitionReader(
     tablePath: TablePath,
     projection: Array[Int],
+    limit: Option[Int],
     flussPartition: FlussUpsertInputPartition,
     flussConfig: Configuration)
-  extends FlussPartitionReader(tablePath, flussConfig)
+  extends FlussPartitionReader(tablePath, flussConfig, limit)
   with Logging {
 
   override protected lazy val projectedRowType: RowType = rowType.project(projectionWithPks)

fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/lake/FlussLakeAppendBatch.scala

@@ -40,9 +40,10 @@ class FlussLakeAppendBatch(
     tableInfo: TableInfo,
     readSchema: StructType,
     pushedPredicate: Option[FlussPredicate],
+    limit: Option[Int],
     options: CaseInsensitiveStringMap,
     flussConfig: Configuration)
-  extends FlussLakeBatch(tablePath, tableInfo, readSchema, options, flussConfig) {
+  extends FlussLakeBatch(tablePath, tableInfo, readSchema, limit, options, flussConfig) {
 
   // Required by FlussLakeBatch but unused — lake snapshot determines start offsets.
   override val startOffsetsInitializer: OffsetsInitializer = OffsetsInitializer.earliest()
@@ -57,6 +58,7 @@ class FlussLakeAppendBatch(
         tablePath,
         projection,
         logTailPredicate,
+        limit,
         options,
         flussConfig)
     } else {
@@ -66,6 +68,7 @@ class FlussLakeAppendBatch(
         projection,
         pushedPredicate,
         logTailPredicate,
+        limit,
         flussConfig)
     }
   }

fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/lake/FlussLakeAppendPartitionReader.scala

@@ -33,8 +33,9 @@ class FlussLakeAppendPartitionReader(
     partition: FlussLakeInputPartition,
     lakeSource: LakeSource[LakeSplit],
     projection: Array[Int],
+    limit: Option[Int],
     flussConfig: Configuration)
-  extends FlussPartitionReader(tablePath, flussConfig)
+  extends FlussPartitionReader(tablePath, flussConfig, limit)
   with Logging {
 
   private var recordIterator: CloseableIterator[LogRecord] = _