[spark] Consolidate scan builders into Append/Upsert with SplitPlanner

fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/SparkTable.scala

@@ -21,7 +21,7 @@ import org.apache.fluss.client.admin.Admin
 import org.apache.fluss.config.{Configuration => FlussConfiguration}
 import org.apache.fluss.metadata.{TableInfo, TablePath}
 import org.apache.fluss.spark.catalog.{AbstractSparkTable, SupportsFlussPartitionManagement}
-import org.apache.fluss.spark.read.{FlussAppendScanBuilder, FlussLakeAppendScanBuilder, FlussLakeUpsertScanBuilder, FlussUpsertScanBuilder}
+import org.apache.fluss.spark.read.{FlussAppendScanBuilder, FlussUpsertScanBuilder}
 import org.apache.fluss.spark.write.{FlussAppendWriteBuilder, FlussUpsertWriteBuilder}
 
 import org.apache.spark.sql.catalyst.SQLConfHelper
@@ -61,25 +61,10 @@ class SparkTable(
 
   override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = {
     populateSparkConf(flussConfig)
-    val isDataLakeEnabled = tableInfo.getTableConfig.isDataLakeEnabled
-    val startupMode = options
-      .getOrDefault(
-        SparkFlussConf.SCAN_START_UP_MODE.key(),
-        flussConfig.get(SparkFlussConf.SCAN_START_UP_MODE))
-      .toUpperCase
-    val isFullMode = startupMode == SparkFlussConf.StartUpMode.FULL.toString
     if (tableInfo.getPrimaryKeys.isEmpty) {
-      if (isDataLakeEnabled && isFullMode) {
-        new FlussLakeAppendScanBuilder(tablePath, tableInfo, options, flussConfig)
-      } else {
-        new FlussAppendScanBuilder(tablePath, tableInfo, options, flussConfig)
-      }
+      new FlussAppendScanBuilder(tablePath, tableInfo, options, flussConfig)
     } else {
-      if (isDataLakeEnabled) {
-        new FlussLakeUpsertScanBuilder(tablePath, tableInfo, options, flussConfig)
-      } else {
-        new FlussUpsertScanBuilder(tablePath, tableInfo, options, flussConfig)
-      }
+      new FlussUpsertScanBuilder(tablePath, tableInfo, options, flussConfig)
     }
   }
 }

fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/catalog/WithFlussAdmin.scala

@@ -22,7 +22,6 @@ import org.apache.fluss.client.admin.Admin
 import org.apache.fluss.config.{Configuration => FlussConfiguration}
 import org.apache.fluss.utils.{IOUtils, Preconditions}
 
-import org.apache.spark.sql.connector.catalog.CatalogPlugin
 import org.apache.spark.sql.util.CaseInsensitiveStringMap
 
 import java.util
@@ -35,9 +34,6 @@ trait WithFlussAdmin extends AutoCloseable {
   private var _admin: Admin = _
   private var _flussConfig: FlussConfiguration = _
 
-  // TODO: init lake spark catalog
-  protected var lakeCatalog: CatalogPlugin = _
-
   protected def initFlussClient(options: CaseInsensitiveStringMap): Unit = {
     val flussConfigs = new util.HashMap[String, String]()
     options.entrySet().asScala.foreach {

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

@@ -17,44 +17,33 @@
 
 package org.apache.fluss.spark.read
 
-import org.apache.fluss.client.{Connection, ConnectionFactory}
-import org.apache.fluss.client.admin.Admin
-import org.apache.fluss.client.initializer.{BucketOffsetsRetrieverImpl, OffsetsInitializer, SnapshotOffsetsInitializer}
-import org.apache.fluss.client.metadata.KvSnapshots
-import org.apache.fluss.client.table.scanner.log.LogScanner
 import org.apache.fluss.config.Configuration
-import org.apache.fluss.metadata.{PartitionInfo, TableBucket, TableInfo, TablePath}
-import org.apache.fluss.predicate.Predicate
-import org.apache.fluss.spark.SparkFlussConf
-import org.apache.fluss.spark.utils.SparkPartitionPredicate
+import org.apache.fluss.metadata.{TableInfo, TablePath}
+import org.apache.fluss.predicate.{Predicate => FlussPredicate}
+import org.apache.fluss.spark.read.lake.FlussLakePartitionReaderFactory
 
 import org.apache.spark.sql.connector.read.{Batch, InputPartition, PartitionReaderFactory}
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.util.CaseInsensitiveStringMap
 
-import java.util
-
 import scala.collection.JavaConverters._
 
+/**
+ * Base class for planner-backed batch scans. The planner (constructed at ScanBuilder-time) owns the
+ * Fluss client Connection/Admin lifecycle and produces the [[InputPartition]]s; the Batch is a thin
+ * adapter that hands the plan to Spark and dispatches [[createReaderFactory]] based on whether the
+ * planner is unioning with a lake snapshot (lake-union) or reading Fluss only (log-only).
+ */
 abstract class FlussBatch(
     tablePath: TablePath,
     tableInfo: TableInfo,
     readSchema: StructType,
     limit: Option[Int],
-    flussConfig: Configuration)
+    flussConfig: Configuration,
+    planner: SplitPlanner)
   extends Batch
   with AutoCloseable {
 
-  lazy val conn: Connection = ConnectionFactory.createConnection(flussConfig)
-
-  lazy val admin: Admin = conn.getAdmin
-
-  lazy val partitionInfos: util.List[PartitionInfo] = admin.listPartitionInfos(tablePath).get()
-
-  def startOffsetsInitializer: OffsetsInitializer
-
-  def stoppingOffsetsInitializer: OffsetsInitializer
-
   protected def projection: Array[Int] = {
     val columnNameToIndex = tableInfo.getSchema.getColumnNames.asScala.zipWithIndex.toMap
     readSchema.fields.map {
@@ -65,47 +54,10 @@ abstract class FlussBatch(
     }
   }
 
-  protected def createUpsertPartitions(
-      partitionName: String,
-      kvSnapshots: KvSnapshots,
-      bucketOffsetsRetriever: BucketOffsetsRetrieverImpl): Array[InputPartition] = {
-    val tableId = kvSnapshots.getTableId
-    val partitionId = kvSnapshots.getPartitionId
-    val bucketIds = kvSnapshots.getBucketIds
-    val bucketIdToLogOffset =
-      stoppingOffsetsInitializer.getBucketOffsets(partitionName, bucketIds, bucketOffsetsRetriever)
-    bucketIds.asScala
-      .map {
-        bucketId =>
-          val tableBucket = new TableBucket(tableId, partitionId, bucketId)
-          val snapshotIdOpt = kvSnapshots.getSnapshotId(bucketId)
-          val logStartingOffsetOpt = kvSnapshots.getLogOffset(bucketId)
-          val logEndingOffset = bucketIdToLogOffset.get(bucketId)
-
-          if (snapshotIdOpt.isPresent) {
-            assert(
-              logStartingOffsetOpt.isPresent,
-              "Log offset must be present when snapshot id is present")
-            FlussUpsertInputPartition(
-              tableBucket,
-              snapshotIdOpt.getAsLong,
-              logStartingOffsetOpt.getAsLong,
-              logEndingOffset)
-          } else {
-            FlussUpsertInputPartition(tableBucket, -1L, LogScanner.EARLIEST_OFFSET, logEndingOffset)
-          }
-      }
-      .map(_.asInstanceOf[InputPartition])
-      .toArray
-  }
+  override def planInputPartitions(): Array[InputPartition] = planner.plan()
 
   override def close(): Unit = {
-    if (admin != null) {
-      admin.close()
-    }
-    if (conn != null) {
-      conn.close()
-    }
+    if (planner != null) planner.close()
   }
 }
 
@@ -114,182 +66,59 @@ class FlussAppendBatch(
     tablePath: TablePath,
     tableInfo: TableInfo,
     readSchema: StructType,
-    pushedPredicate: Option[Predicate],
-    partitionPredicate: Option[Predicate],
+    pushedPredicate: Option[FlussPredicate],
     limit: Option[Int],
     options: CaseInsensitiveStringMap,
-    flussConfig: Configuration)
-  extends FlussBatch(tablePath, tableInfo, readSchema, limit, flussConfig) {
-
-  override val startOffsetsInitializer: OffsetsInitializer = {
-    FlussOffsetInitializers.startOffsetsInitializer(options, flussConfig)
-  }
-
-  override val stoppingOffsetsInitializer: OffsetsInitializer = {
-    FlussOffsetInitializers.stoppingOffsetsInitializer(true, options, flussConfig)
-  }
-
-  override def planInputPartitions(): Array[InputPartition] = {
-    val maxRecordsPerPartition: Option[Long] = {
-      val value = flussConfig.getLong(SparkFlussConf.SCAN_MAX_RECORDS_PER_PARTITION, 0)
-      if (value > 0) Some(value) else None
-    }
-
-    val bucketOffsetsRetrieverImpl = maxRecordsPerPartition match {
-      case Some(_) => new BucketOffsetsRetrieverImpl(admin, tablePath, true)
-      case _ => new BucketOffsetsRetrieverImpl(admin, tablePath)
-    }
-    val buckets = (0 until tableInfo.getNumBuckets).toSeq
-
-    def splitOffsetRange(
-        tableBucket: TableBucket,
-        startOffset: Long,
-        stopOffset: Long,
-        maxRecords: Long): Seq[InputPartition] = {
-      if (
-        startOffset < 0 || stopOffset <= startOffset || stopOffset <= (startOffset + maxRecords)
-      ) {
-        return Seq(FlussAppendInputPartition(tableBucket, startOffset, stopOffset))
-      }
-      val rangeSize = stopOffset - startOffset
-      val numSplits = ((rangeSize + maxRecords - 1) / maxRecords).toInt
-      val step = (rangeSize + numSplits - 1) / numSplits
+    flussConfig: Configuration,
+    appendPlanner: AppendSplitPlanner)
+  extends FlussBatch(tablePath, tableInfo, readSchema, limit, flussConfig, appendPlanner) {
 
-      Iterator
-        .from(0)
-        .take(numSplits)
-        .map(i => startOffset + i * step)
-        .map {
-          from => FlussAppendInputPartition(tableBucket, from, math.min(from + step, stopOffset))
-        }
-        .toSeq
-    }
-
-    def createPartitions(
-        partitionId: Option[Long],
-        startBucketOffsets: Map[Integer, Long],
-        stoppingBucketOffsets: Map[Integer, Long]): Array[InputPartition] = {
-      buckets.flatMap {
-        bucketId =>
-          val (startOffset, stopOffset) =
-            (startBucketOffsets(bucketId), stoppingBucketOffsets(bucketId))
-          val tableBucket = partitionId match {
-            case Some(pid) => new TableBucket(tableInfo.getTableId, pid, bucketId)
-            case None => new TableBucket(tableInfo.getTableId, bucketId)
-          }
-          maxRecordsPerPartition match {
-            case Some(maxRecs) =>
-              splitOffsetRange(tableBucket, startOffset, stopOffset, maxRecs)
-            case _ =>
-              Seq(FlussAppendInputPartition(tableBucket, startOffset, stopOffset))
-          }
-      }.toArray
-    }
-
-    if (tableInfo.isPartitioned) {
-      val matching =
-        SparkPartitionPredicate.filterPartitions(
-          tableInfo,
-          partitionInfos.asScala.toSeq,
-          partitionPredicate)
-      matching
-        .map {
-          partitionInfo =>
-            val startBucketOffsets = startOffsetsInitializer.getBucketOffsets(
-              partitionInfo.getPartitionName,
-              buckets.map(Integer.valueOf).asJava,
-              bucketOffsetsRetrieverImpl)
-            val stoppingBucketOffsets = stoppingOffsetsInitializer.getBucketOffsets(
-              partitionInfo.getPartitionName,
-              buckets.map(Integer.valueOf).asJava,
-              bucketOffsetsRetrieverImpl)
-            (
-              partitionInfo.getPartitionId,
-              startBucketOffsets.asScala.map(e => (e._1, Long2long(e._2))),
-              stoppingBucketOffsets.asScala.map(e => (e._1, Long2long(e._2))))
-        }
-        .flatMap {
-          case (partitionId, startBucketOffsets, stoppingBucketOffsets) =>
-            createPartitions(
-              Some(partitionId),
-              startBucketOffsets.toMap,
-              stoppingBucketOffsets.toMap)
-        }
-        .toArray
+  override def createReaderFactory(): PartitionReaderFactory = {
+    if (appendPlanner.hasLakeSnapshot) {
+      new FlussLakePartitionReaderFactory(
+        tableInfo.getProperties.toMap,
+        tablePath,
+        projection,
+        pushedPredicate,
+        appendPlanner.logTailPredicate,
+        limit,
+        flussConfig)
     } else {
-      val startBucketOffsets = startOffsetsInitializer.getBucketOffsets(
-        null,
-        buckets.map(Integer.valueOf).asJava,
-        bucketOffsetsRetrieverImpl)
-      val stoppingBucketOffsets = stoppingOffsetsInitializer.getBucketOffsets(
-        null,
-        buckets.map(Integer.valueOf).asJava,
-        bucketOffsetsRetrieverImpl)
-      createPartitions(
-        None,
-        startBucketOffsets.asScala.map(e => (e._1, Long2long(e._2))).toMap,
-        stoppingBucketOffsets.asScala.map(e => (e._1, Long2long(e._2))).toMap)
+      new FlussAppendPartitionReaderFactory(
+        tablePath,
+        projection,
+        pushedPredicate,
+        limit,
+        options,
+        flussConfig)
     }
   }
-
-  override def createReaderFactory(): PartitionReaderFactory = {
-    new FlussAppendPartitionReaderFactory(
-      tablePath,
-      projection,
-      pushedPredicate,
-      limit,
-      options,
-      flussConfig)
-  }
-
 }
 
 /** Batch for reading primary key table (upsert table). */
 class FlussUpsertBatch(
     tablePath: TablePath,
     tableInfo: TableInfo,
     readSchema: StructType,
-    partitionPredicate: Option[Predicate],
+    pushedPredicate: Option[FlussPredicate],
     limit: Option[Int],
     options: CaseInsensitiveStringMap,
-    flussConfig: Configuration)
-  extends FlussBatch(tablePath, tableInfo, readSchema, limit, flussConfig) {
+    flussConfig: Configuration,
+    upsertPlanner: UpsertSplitPlanner)
+  extends FlussBatch(tablePath, tableInfo, readSchema, limit, flussConfig, upsertPlanner) {
 
-  override val startOffsetsInitializer: OffsetsInitializer = {
-    val offsetsInitializer = FlussOffsetInitializers.startOffsetsInitializer(options, flussConfig)
-    if (!offsetsInitializer.isInstanceOf[SnapshotOffsetsInitializer]) {
-      throw new UnsupportedOperationException("Upsert scan only support FULL startup mode.")
-    }
-    offsetsInitializer
-  }
-
-  override val stoppingOffsetsInitializer: OffsetsInitializer = {
-    FlussOffsetInitializers.stoppingOffsetsInitializer(true, options, flussConfig)
-  }
-
-  private val bucketOffsetsRetriever = new BucketOffsetsRetrieverImpl(admin, tablePath)
-
-  override def planInputPartitions(): Array[InputPartition] = {
-    if (tableInfo.isPartitioned) {
-      val matching =
-        SparkPartitionPredicate.filterPartitions(
-          tableInfo,
-          partitionInfos.asScala.toSeq,
-          partitionPredicate)
-      matching.flatMap {
-        partitionInfo =>
-          val partitionName = partitionInfo.getPartitionName
-          val kvSnapshots =
-            admin.getLatestKvSnapshots(tablePath, partitionName).get()
-          createUpsertPartitions(partitionName, kvSnapshots, bucketOffsetsRetriever)
-      }.toArray
+  override def createReaderFactory(): PartitionReaderFactory = {
+    if (upsertPlanner.hasLakeSnapshot) {
+      new FlussLakePartitionReaderFactory(
+        tableInfo.getProperties.toMap,
+        tablePath,
+        projection,
+        pushedPredicate,
+        upsertPlanner.logTailPredicate,
+        limit,
+        flussConfig)
     } else {
-      val kvSnapshots = admin.getLatestKvSnapshots(tablePath).get()
-      createUpsertPartitions(null, kvSnapshots, bucketOffsetsRetriever)
+      new FlussUpsertPartitionReaderFactory(tablePath, projection, limit, options, flussConfig)
     }
   }
-
-  override def createReaderFactory(): PartitionReaderFactory = {
-    new FlussUpsertPartitionReaderFactory(tablePath, projection, limit, options, flussConfig)
-  }
 }