diff --git a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/SparkTable.scala b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/SparkTable.scala index 0a9d9663f6..144db03aea 100644 --- a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/SparkTable.scala +++ b/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) } } } diff --git a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/catalog/WithFlussAdmin.scala b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/catalog/WithFlussAdmin.scala index 0579bca105..97f2c0ee45 100644 --- a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/catalog/WithFlussAdmin.scala +++ b/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 { diff --git a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussBatch.scala b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussBatch.scala index 43fc8bc37e..022723d9c1 100644 --- a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussBatch.scala +++ b/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,134 +66,33 @@ 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). */ @@ -249,47 +100,25 @@ 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) - } } diff --git a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussScan.scala b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussScan.scala index cce6c8c139..0a0d7def1a 100644 --- a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussScan.scala +++ b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussScan.scala @@ -21,7 +21,6 @@ import org.apache.fluss.config.Configuration import org.apache.fluss.metadata.{TableInfo, TablePath} import org.apache.fluss.predicate.{Predicate => FlussPredicate} import org.apache.fluss.spark.SparkConversions -import org.apache.fluss.spark.read.lake.{FlussLakeAppendBatch, FlussLakeUpsertBatch} import org.apache.spark.sql.connector.expressions.filter.Predicate import org.apache.spark.sql.connector.metric.CustomMetric @@ -70,7 +69,11 @@ trait FlussScan extends Scan { Array(FlussNumRowsReadMetric()) } -/** Fluss Append Scan. */ +/** + * Fluss Append (log-table) scan. Whether the underlying batch reads from Fluss only or unions Fluss + * with a lake snapshot is determined by the [[AppendSplitPlanner]] instance passed in from the + * ScanBuilder. Description reflects the planner category. + */ case class FlussAppendScan( tablePath: TablePath, tableInfo: TableInfo, @@ -80,10 +83,12 @@ case class FlussAppendScan( override val pushedSparkPredicates: Seq[Predicate], override val limit: Option[Int], options: CaseInsensitiveStringMap, - flussConfig: Configuration) + flussConfig: Configuration, + planner: AppendSplitPlanner) extends FlussScan { - override protected val scanType: String = "Append" + override protected val scanType: String = + if (planner.hasLakeSnapshot) "LakeAppend" else "Append" override def toBatch: Batch = { new FlussAppendBatch( @@ -91,48 +96,10 @@ case class FlussAppendScan( tableInfo, readSchema, pushedPredicate, - partitionPredicate, limit, options, - flussConfig) - } - - override def toMicroBatchStream(checkpointLocation: String): MicroBatchStream = { - new FlussAppendMicroBatchStream( - tablePath, - tableInfo, - readSchema, - options, flussConfig, - checkpointLocation) - } -} - -/** Fluss Lake Append Scan. */ -case class FlussLakeAppendScan( - tablePath: TablePath, - tableInfo: TableInfo, - requiredSchema: Option[StructType], - pushedPredicate: Option[FlussPredicate], - override val partitionPredicate: Option[FlussPredicate], - override val pushedSparkPredicates: Seq[Predicate], - override val limit: Option[Int], - options: CaseInsensitiveStringMap, - flussConfig: Configuration) - extends FlussScan { - - override protected val scanType: String = "LakeAppend" - - override def toBatch: Batch = { - new FlussLakeAppendBatch( - tablePath, - tableInfo, - readSchema, - pushedPredicate, - partitionPredicate, - limit, - options, - flussConfig) + planner) } override def toMicroBatchStream(checkpointLocation: String): MicroBatchStream = { @@ -146,43 +113,12 @@ case class FlussLakeAppendScan( } } -/** Fluss Upsert Scan. */ +/** + * Fluss Upsert (primary-key table) scan. Whether the underlying batch reads from Fluss only or + * unions Fluss with a lake snapshot is determined by the [[UpsertSplitPlanner]] instance passed in + * from the ScanBuilder. + */ case class FlussUpsertScan( - tablePath: TablePath, - 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, - limit, - options, - flussConfig) - } - - override def toMicroBatchStream(checkpointLocation: String): MicroBatchStream = { - new FlussUpsertMicroBatchStream( - tablePath, - tableInfo, - readSchema, - options, - flussConfig, - checkpointLocation) - } -} - -/** Fluss Lake Upsert Scan for lake-enabled primary key tables. */ -case class FlussLakeUpsertScan( tablePath: TablePath, tableInfo: TableInfo, requiredSchema: Option[StructType], @@ -191,21 +127,23 @@ case class FlussLakeUpsertScan( override val pushedSparkPredicates: Seq[Predicate], override val limit: Option[Int], options: CaseInsensitiveStringMap, - flussConfig: Configuration) + flussConfig: Configuration, + planner: UpsertSplitPlanner) extends FlussScan { - override protected val scanType: String = "LakeUpsert" + override protected val scanType: String = + if (planner.hasLakeSnapshot) "LakeUpsert" else "Upsert" override def toBatch: Batch = { - new FlussLakeUpsertBatch( + new FlussUpsertBatch( tablePath, tableInfo, readSchema, pushedPredicate, - partitionPredicate, limit, options, - flussConfig) + flussConfig, + planner) } override def toMicroBatchStream(checkpointLocation: String): MicroBatchStream = { diff --git a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussScanBuilder.scala b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussScanBuilder.scala index a31e5e218f..fe7d60e860 100644 --- a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussScanBuilder.scala +++ b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussScanBuilder.scala @@ -20,7 +20,7 @@ package org.apache.fluss.spark.read import org.apache.fluss.config.{Configuration => FlussConfiguration} import org.apache.fluss.metadata.{LogFormat, TableInfo, TablePath} import org.apache.fluss.predicate.{Predicate => FlussPredicate} -import org.apache.fluss.spark.read.lake.{FlussLakeBatch, FlussLakeUtils} +import org.apache.fluss.spark.read.lake.FlussLakeUtils import org.apache.fluss.spark.utils.{SparkPartitionPredicate, SparkPredicateConverter} import org.apache.spark.sql.connector.expressions.filter.Predicate @@ -72,92 +72,80 @@ trait FlussSupportsPushDownPartitionFilters override def pushedPredicates(): Array[Predicate] = acceptedPredicates } -trait FlussSupportsPushDownV2Filters extends FlussSupportsPushDownPartitionFilters { - - protected def convertAndStorePredicates(predicates: Array[Predicate]): Unit = { - val (predicate, accepted) = - SparkPredicateConverter.convertPredicates(tableInfo.getRowType, predicates.toSeq) - pushedPredicate = predicate - acceptedPredicates = accepted.toArray - } - - override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { - val nonPartitionPredicates = super.pushPredicates(predicates) - if (!tableInfo.hasPrimaryKey && tableInfo.getTableConfig.getLogFormat == LogFormat.ARROW) { - // Server-side batch filter for log table only supports ARROW; other log formats reject it. - convertAndStorePredicates(nonPartitionPredicates) - } - nonPartitionPredicates - } -} - /** - * Lake reads push to the lake source regardless of log format. Each convertible predicate is - * offered to the lake source individually; only the lake-accepted subset is reported back to Spark - * and combined into the predicate handed to the scan. + * Data-predicate push-down. Two branches: + * 1. Non-PK log tables with ARROW format: converts predicates to server-side batch filters. + * 2. Lake-enabled tables (including PK tables): probes the lake source for which predicates it + * accepts, reports those back to Spark, and passes the accepted predicate to the planner which + * applies it on the lake-union path. + * + * This consolidates the former `FlussSupportsPushDownV2Filters` (branch 1) and + * `FlussLakeSupportsPushDownV2Filters` (branch 2) into a single trait. */ -trait FlussLakeSupportsPushDownV2Filters extends FlussSupportsPushDownPartitionFilters { +trait FlussSupportsPushDownV2Filters extends FlussSupportsPushDownPartitionFilters { def tablePath: TablePath def flussConfig: FlussConfiguration override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { - val nonPartitionPredicates = super.pushPredicates(predicates) - - // Pass ALL predicates to Lake Source (including partition predicates) for lake-side filtering - val pairs = - SparkPredicateConverter.convertPerPredicate(tableInfo.getRowType, predicates.toSeq) - val (acceptedSpark, acceptedFluss) = if (pairs.isEmpty) { - (Seq.empty[Predicate], Seq.empty[FlussPredicate]) - } else { - val lakeSource = - FlussLakeUtils.createLakeSource(flussConfig.toMap, tableInfo.getProperties.toMap, tablePath) - val result = FlussLakeBatch.applyLakeFilters(lakeSource, pairs.map(_._2).asJava) - // Identity-match: lake sources are expected to return the same instances they received. - val acceptedSet: JSet[FlussPredicate] = - Collections.newSetFromMap(new IdentityHashMap()) - acceptedSet.addAll(result.acceptedPredicates()) - pairs.collect { case (sp, fp) if acceptedSet.contains(fp) => (sp, fp) }.unzip + val nonPartition = super.pushPredicates(predicates) + if (!tableInfo.hasPrimaryKey && tableInfo.getTableConfig.getLogFormat == LogFormat.ARROW) { + // Server-side batch filter for log table only supports ARROW; other log formats reject it. + val (predicate, accepted) = + SparkPredicateConverter.convertPredicates(tableInfo.getRowType, nonPartition.toSeq) + pushedPredicate = predicate + acceptedPredicates = accepted.toArray + } else if (tableInfo.getTableConfig.isDataLakeEnabled) { + // Lake-enabled tables: probe the lake source for which predicates it accepts. All predicates + // (including partition) are offered because the lake source handles both partition pruning + // and data filtering internally. This recovers the former FlussLakeSupportsPushDownV2Filters + // behavior that was split into a separate trait before consolidation. + val pairs = + SparkPredicateConverter.convertPerPredicate(tableInfo.getRowType, predicates.toSeq) + if (pairs.nonEmpty) { + val lakeSource = FlussLakeUtils.createLakeSource( + flussConfig.toMap, + tableInfo.getProperties.toMap, + tablePath) + val result = FlussLakeUtils.applyLakeFilters(lakeSource, pairs.map(_._2).asJava) + val acceptedSet: JSet[FlussPredicate] = + Collections.newSetFromMap(new IdentityHashMap[FlussPredicate, java.lang.Boolean]()) + acceptedSet.addAll(result.acceptedPredicates()) + val (acceptedSpark, acceptedFluss) = + pairs.collect { case (sp, fp) if acceptedSet.contains(fp) => (sp, fp) }.unzip + pushedPredicate = SparkPredicateConverter.combineAnd(acceptedFluss) + acceptedPredicates = acceptedSpark.toArray + } } - pushedPredicate = SparkPredicateConverter.combineAnd(acceptedFluss) - acceptedPredicates = acceptedSpark.toArray - nonPartitionPredicates + nonPartition } } -/** Fluss Append Scan Builder. */ +/** + * Fluss Append (log-table) Scan Builder. The concrete [[AppendPlanner]] is materialized in + * [[build]] once pushdown/prune state is settled. The planner probes the readable lake snapshot + * itself at construction — the ScanBuilder is intentionally unaware of lake-union vs log-only + * routing. + */ class FlussAppendScanBuilder( - tablePath: TablePath, + val tablePath: TablePath, val tableInfo: TableInfo, options: CaseInsensitiveStringMap, val flussConfig: FlussConfiguration) extends FlussSupportsPushDownV2Filters { override def build(): Scan = { - FlussAppendScan( + val projection = FlussScanBuilder.projectionOf(tableInfo, requiredSchema) + val planner = new AppendPlanner( tablePath, tableInfo, - requiredSchema, - pushedPredicate, partitionPredicate, - acceptedPredicates.toSeq, - limit, + pushedPredicate, + projection, options, flussConfig) - } -} - -/** Fluss Lake Append Scan Builder. */ -class FlussLakeAppendScanBuilder( - val tablePath: TablePath, - val tableInfo: TableInfo, - options: CaseInsensitiveStringMap, - val flussConfig: FlussConfiguration) - extends FlussLakeSupportsPushDownV2Filters { - - override def build(): Scan = { - FlussLakeAppendScan( + FlussAppendScan( tablePath, tableInfo, requiredSchema, @@ -166,40 +154,34 @@ class FlussLakeAppendScanBuilder( acceptedPredicates.toSeq, limit, options, - flussConfig) + flussConfig, + planner) } } -/** Fluss Upsert Scan Builder. */ +/** + * Fluss Upsert (primary-key table) Scan Builder. The concrete [[UpsertPlanner]] is materialized in + * [[build]] once pushdown/prune state is settled. The planner probes the readable lake snapshot + * itself at construction. + */ class FlussUpsertScanBuilder( - tablePath: TablePath, + val tablePath: TablePath, val tableInfo: TableInfo, options: CaseInsensitiveStringMap, val flussConfig: FlussConfiguration) extends FlussSupportsPushDownV2Filters { override def build(): Scan = { - FlussUpsertScan( + val projection = FlussScanBuilder.projectionOf(tableInfo, requiredSchema) + val planner = new UpsertPlanner( tablePath, tableInfo, - requiredSchema, partitionPredicate, - limit, + pushedPredicate, + projection, options, flussConfig) - } -} - -/** Fluss Lake Upsert Scan Builder for lake-enabled primary key tables. */ -class FlussLakeUpsertScanBuilder( - val tablePath: TablePath, - val tableInfo: TableInfo, - options: CaseInsensitiveStringMap, - val flussConfig: FlussConfiguration) - extends FlussLakeSupportsPushDownV2Filters { - - override def build(): Scan = { - FlussLakeUpsertScan( + FlussUpsertScan( tablePath, tableInfo, requiredSchema, @@ -208,6 +190,27 @@ class FlussLakeUpsertScanBuilder( acceptedPredicates.toSeq, limit, options, - flussConfig) + flussConfig, + planner) + } +} + +object FlussScanBuilder { + + /** Convert a Spark required-schema projection back to Fluss column indices. */ + def projectionOf(tableInfo: TableInfo, requiredSchema: Option[StructType]): Array[Int] = { + val allFields = (0 until tableInfo.getRowType.getFieldCount).toArray + requiredSchema match { + case None => allFields + case Some(schema) => + val columnNameToIndex = + tableInfo.getSchema.getColumnNames.asScala.zipWithIndex.toMap + schema.fields.map { + f => + columnNameToIndex.getOrElse( + f.name, + throw new IllegalArgumentException(s"Invalid field name: ${f.name}")) + } + } } } diff --git a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/SplitPlanner.scala b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/SplitPlanner.scala new file mode 100644 index 0000000000..62e3d65ea0 --- /dev/null +++ b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/SplitPlanner.scala @@ -0,0 +1,790 @@ +/* + * 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.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} +import org.apache.fluss.client.metadata.{KvSnapshots, LakeSnapshot} +import org.apache.fluss.client.table.scanner.log.LogScanner +import org.apache.fluss.config.Configuration +import org.apache.fluss.exception.LakeTableSnapshotNotExistException +import org.apache.fluss.lake.source.{LakeSource, LakeSplit} +import org.apache.fluss.metadata.{LogFormat, PartitionInfo, ResolvedPartitionSpec, TableBucket, TableInfo, TablePath} +import org.apache.fluss.predicate.{Predicate => FlussPredicate} +import org.apache.fluss.spark.SparkFlussConf +import org.apache.fluss.spark.read.lake.{FlussLakeInputPartition, FlussLakeUpsertInputPartition, FlussLakeUtils} +import org.apache.fluss.spark.utils.SparkPartitionPredicate +import org.apache.fluss.utils.ExceptionUtils + +import org.apache.spark.sql.connector.read.InputPartition +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +import java.util + +import scala.collection.JavaConverters._ +import scala.collection.mutable + +/** + * Plans Spark [[InputPartition]]s for a Fluss batch scan. The planner owns a Fluss client + * connection for the lifetime of the scan and is closed by the enclosing batch. + * + * Concrete planners probe for a readable lake snapshot at construction time and branch inside + * `plan()`: + * - Lake-union branch (snapshot present): unions lake splits with a Fluss log-tail (append) or + * kv+log-tail (upsert). + * - Log-only branch (snapshot absent): the plan is derived exclusively from Fluss metadata — the + * full Fluss log (append) or Fluss kv snapshots + log tail (upsert). + * + * The probe is snapshot-isolated: it is performed exactly once per planner instance (i.e. once per + * Spark scan build). There is no runtime fallback path — presence/absence is decided + * deterministically at construction and `plan()` picks a branch accordingly. + */ +sealed trait SplitPlanner extends AutoCloseable { + + def tablePath: TablePath + + def tableInfo: TableInfo + + def flussConfig: Configuration + + def plan(): Array[InputPartition] + + /** + * Whether this planner will union its Fluss plan with a lake snapshot. Determined once at + * construction; deterministic across `plan()` invocations. + */ + def hasLakeSnapshot: Boolean + + /** + * Server-side batch filter to attach to the Fluss log-tail reader. Only ARROW-formatted log + * tables accept a server filter; other formats must return None to avoid a server-side reject. + * Upsert planners always return None because the reader must reconcile the full log tail with kv + * snapshots. + */ + def logTailPredicate: Option[FlussPredicate] +} + +/** Marker: planner yields partitions consumable by an append (log-table) reader factory. */ +sealed trait AppendSplitPlanner extends SplitPlanner + +/** Marker: planner yields partitions consumable by an upsert (primary-key) reader factory. */ +sealed trait UpsertSplitPlanner extends SplitPlanner + +/** + * Base implementation: owns the shared Fluss client Connection + Admin so concrete planners can + * issue metadata requests. [[close]] tears down the whole chain. + * + * Also centralizes the lake-snapshot probe: if data lake is enabled at the table level, try loading + * the readable snapshot; treat [[LakeTableSnapshotNotExistException]] as "no snapshot", all other + * exceptions propagate. + */ +abstract class AbstractSplitPlanner( + override val tablePath: TablePath, + override val tableInfo: TableInfo, + override val flussConfig: Configuration) + extends SplitPlanner { + + protected lazy val conn: Connection = ConnectionFactory.createConnection(flussConfig) + + protected lazy val admin: Admin = conn.getAdmin + + protected lazy val partitionInfos: util.List[PartitionInfo] = + admin.listPartitionInfos(tablePath).get() + + protected def stoppingOffsetsInitializer: OffsetsInitializer + + /** + * Probes the readable lake snapshot exactly once. Absent = data-lake disabled at the table level, + * OR the readable snapshot admin call reported no snapshot; Present = snapshot to union with the + * Fluss tail. Other exceptions propagate. + */ + protected def probeLakeSnapshot(): Option[LakeSnapshot] = { + if (!tableInfo.getTableConfig.isDataLakeEnabled) { + None + } else { + try { + Some(admin.getReadableLakeSnapshot(tablePath).get()) + } catch { + case e: Exception => + if ( + ExceptionUtils + .stripExecutionException(e) + .isInstanceOf[LakeTableSnapshotNotExistException] + ) { + None + } else { + throw e + } + } + } + } + + protected def getBucketOffsets( + initializer: OffsetsInitializer, + partitionName: String, + buckets: Seq[Int], + bucketOffsetsRetriever: BucketOffsetsRetrieverImpl): Map[Int, Long] = { + initializer + .getBucketOffsets(partitionName, buckets.map(Integer.valueOf).asJava, bucketOffsetsRetriever) + .asScala + .map(e => (e._1.intValue(), Long2long(e._2))) + .toMap + } + + override def close(): Unit = { + if (admin != null) admin.close() + if (conn != null) conn.close() + } +} + +/** + * Single append (log-table) planner. Probes a readable lake snapshot at construction; if present, + * the plan is a union of lake splits and the Fluss log-tail (from each bucket's snapshotLogOffset + * to committed). If absent, the plan is a pure Fluss log scan from earliest to committed + * (SCAN_START_UP_MODE deliberately ignored — see class scaladoc note below). + * + * Batch semantics note: start offset is hardcoded to [[OffsetsInitializer.full]] instead of + * consuming the user-facing SCAN_START_UP_MODE. Rationale — batch reads semantically mean "the full + * table". Letting SCAN_START_UP_MODE (a streaming concept) also gate batch planning creates two + * problems: (a) the same config key would carry different semantics on append vs upsert tables (KV + * snapshot has no partial-read semantics), which is confusing; (b) with mode=latest and no writes + * since planning time, start==stop==tail — an empty range that trips the reader-side + * `Invalid offset range` guard. Symmetric "batch = earliest → committed" closes both concerns and + * keeps append/upsert planners aligned. Time-range batch reads should be expressed via predicate + * pushdown on the timestamp column, not startup mode. + * + * `OffsetsInitializer.full()` is chosen over `OffsetsInitializer.earliest()` intentionally: for a + * log table the two are semantically equivalent (see OffsetsInitializer.full javadoc), but full() + * resolves each bucket's start offset to a concrete numeric value via RPC, whereas earliest() + * returns the `LogScanner.EARLIEST_OFFSET` (-2) sentinel — the split-by-max-records logic below + * requires concrete numeric bounds to compute range partitions. + */ +class AppendPlanner( + override val tablePath: TablePath, + override val tableInfo: TableInfo, + partitionPredicate: Option[FlussPredicate], + pushedPredicate: Option[FlussPredicate], + projection: Array[Int], + options: CaseInsensitiveStringMap, + override val flussConfig: Configuration) + extends AbstractSplitPlanner(tablePath, tableInfo, flussConfig) + with AppendSplitPlanner { + + private val readableLakeSnapshot: Option[LakeSnapshot] = probeLakeSnapshot() + + override val hasLakeSnapshot: Boolean = readableLakeSnapshot.isDefined + + private val startOffsetsInitializer: OffsetsInitializer = OffsetsInitializer.full() + + override protected val stoppingOffsetsInitializer: OffsetsInitializer = + FlussOffsetInitializers.stoppingOffsetsInitializer(true, options, flussConfig) + + // Server-side log filter requires ARROW format. Pushdown already gates this on the log-only + // path (never sets pushedPredicate for non-ARROW), but re-checking here keeps the planner + // self-consistent regardless of upstream pushdown behavior and applies uniformly to the + // lake-union path (whose log-tail component reads from the same Fluss reader). + override val logTailPredicate: Option[FlussPredicate] = + if (tableInfo.getTableConfig.getLogFormat == LogFormat.ARROW) pushedPredicate else None + + override def plan(): Array[InputPartition] = readableLakeSnapshot match { + case Some(snap) => planLakeUnion(snap) + case None => planLogOnly() + } + + // --------------------------------------------------------------------------------------------- + // Log-only branch: pure Fluss log scan from earliest → committed with optional range splitting. + // --------------------------------------------------------------------------------------------- + + private def planLogOnly(): 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 + + 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 + } 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) + } + } + + // --------------------------------------------------------------------------------------------- + // Lake-union branch: lake splits + Fluss log tail from each bucket's snapshotLogOffset onward. + // --------------------------------------------------------------------------------------------- + + private def planLakeUnion(snap: LakeSnapshot): Array[InputPartition] = { + val lakeSource = + FlussLakeUtils.createLakeSource(flussConfig.toMap, tableInfo.getProperties.toMap, tablePath) + lakeSource.withProject(FlussLakeUtils.lakeProjection(projection)) + pushedPredicate.foreach(FlussLakeUtils.applyLakeFilters(lakeSource, _)) + + val lakeSplits = lakeSource + .createPlanner(new LakeSource.PlannerContext { + override def snapshotId(): Long = snap.getSnapshotId + }) + .plan() + + val tableBucketsOffset = snap.getTableBucketsOffset + val buckets = (0 until tableInfo.getNumBuckets).toSeq + val bucketOffsetsRetriever = new BucketOffsetsRetrieverImpl(admin, tablePath) + + if (tableInfo.isPartitioned) { + planLakePartitionedTable( + lakeSplits.asScala.toSeq, + tableBucketsOffset, + buckets, + bucketOffsetsRetriever) + } else { + planLakeNonPartitionedTable( + lakeSplits.asScala.toSeq, + tableBucketsOffset, + buckets, + bucketOffsetsRetriever) + } + } + + private def planLakeNonPartitionedTable( + lakeSplits: Seq[LakeSplit], + tableBucketsOffset: java.util.Map[TableBucket, java.lang.Long], + buckets: Seq[Int], + bucketOffsetsRetriever: BucketOffsetsRetrieverImpl): Array[InputPartition] = { + val tableId = tableInfo.getTableId + + val lakePartitions = createLakePartitions(lakeSplits, tableId, partitionId = None) + + val stoppingOffsets = + getBucketOffsets(stoppingOffsetsInitializer, null, buckets, bucketOffsetsRetriever) + val logPartitions = buckets.flatMap { + bucketId => + val tableBucket = new TableBucket(tableId, bucketId) + createLogTailPartition(tableBucket, tableBucketsOffset, stoppingOffsets(bucketId)) + } + + (lakePartitions ++ logPartitions).toArray + } + + private def planLakePartitionedTable( + lakeSplits: Seq[LakeSplit], + tableBucketsOffset: java.util.Map[TableBucket, java.lang.Long], + buckets: Seq[Int], + bucketOffsetsRetriever: BucketOffsetsRetrieverImpl): Array[InputPartition] = { + val tableId = tableInfo.getTableId + + val filteredPartitionInfos = SparkPartitionPredicate.filterPartitions( + tableInfo, + partitionInfos.asScala.toSeq, + partitionPredicate) + + val flussPartitionIdByName = mutable.LinkedHashMap.empty[String, Long] + filteredPartitionInfos.foreach { + pi => flussPartitionIdByName(pi.getPartitionName) = pi.getPartitionId + } + + val lakeSplitsByPartition = groupLakeSplitsByPartitionBuffered(lakeSplits) + var lakeSplitPartitionId = -1L + + val lakeAndLogPartitions = lakeSplitsByPartition.flatMap { + case (partitionName, (partitionValues, splits)) => + flussPartitionIdByName.remove(partitionName) match { + case Some(partitionId) => + val lakePartitions = + createLakePartitions(splits.toSeq, tableId, Some(partitionId)) + + val stoppingOffsets = getBucketOffsets( + stoppingOffsetsInitializer, + partitionName, + buckets, + bucketOffsetsRetriever) + val logPartitions = buckets.flatMap { + bucketId => + val tableBucket = new TableBucket(tableId, partitionId, bucketId) + createLogTailPartition(tableBucket, tableBucketsOffset, stoppingOffsets(bucketId)) + } + + lakePartitions ++ logPartitions + + case None => + if ( + SparkPartitionPredicate + .matchesPartition(tableInfo, partitionValues, partitionPredicate) + ) { + val pid = lakeSplitPartitionId + lakeSplitPartitionId -= 1 + createLakePartitions(splits.toSeq, tableId, Some(pid)) + } else { + Seq.empty + } + } + }.toSeq + + val flussOnlyPartitions = flussPartitionIdByName.flatMap { + case (partitionName, partitionId) => + val stoppingOffsets = getBucketOffsets( + stoppingOffsetsInitializer, + partitionName, + buckets, + bucketOffsetsRetriever) + buckets.flatMap { + bucketId => + val stoppingOffset = stoppingOffsets(bucketId) + if (stoppingOffset > 0) { + val tableBucket = new TableBucket(tableId, partitionId, bucketId) + Some( + FlussAppendInputPartition( + tableBucket, + LogScanner.EARLIEST_OFFSET, + stoppingOffset): InputPartition) + } else { + None + } + } + }.toSeq + + (lakeAndLogPartitions ++ flussOnlyPartitions).toArray + } + + private def groupLakeSplitsByPartitionBuffered(lakeSplits: Seq[LakeSplit]) + : mutable.LinkedHashMap[String, (Seq[String], mutable.ArrayBuffer[LakeSplit])] = { + val grouped = + mutable.LinkedHashMap.empty[String, (Seq[String], mutable.ArrayBuffer[LakeSplit])] + lakeSplits.foreach { + split => + val partitionValues = + if (split.partition() == null) Seq.empty[String] else split.partition().asScala.toSeq + val partitionName = + if (partitionValues.isEmpty) "" + else partitionValues.mkString(ResolvedPartitionSpec.PARTITION_SPEC_SEPARATOR) + val (_, buf) = + grouped.getOrElseUpdate(partitionName, (partitionValues, mutable.ArrayBuffer.empty)) + buf += split + } + grouped + } + + private def createLakePartitions( + splits: Seq[LakeSplit], + tableId: Long, + partitionId: Option[Long]): Seq[InputPartition] = { + splits.map { + split => + val tableBucket = partitionId match { + case Some(pid) => new TableBucket(tableId, pid, split.bucket()) + case None => new TableBucket(tableId, split.bucket()) + } + FlussLakeInputPartition(tableBucket, split) + } + } + + private def createLogTailPartition( + tableBucket: TableBucket, + tableBucketsOffset: java.util.Map[TableBucket, java.lang.Long], + stoppingOffset: Long): Option[InputPartition] = { + val snapshotLogOffset = tableBucketsOffset.get(tableBucket) + if (snapshotLogOffset != null) { + if (snapshotLogOffset.longValue() < stoppingOffset) { + Some(FlussAppendInputPartition(tableBucket, snapshotLogOffset.longValue(), stoppingOffset)) + } else { + None + } + } else if (stoppingOffset > 0) { + Some(FlussAppendInputPartition(tableBucket, LogScanner.EARLIEST_OFFSET, stoppingOffset)) + } else { + None + } + } +} + +/** + * Single upsert (primary-key table) planner. Probes a readable lake snapshot at construction; if + * present, the plan is a union of lake splits and per-bucket Fluss (kv-snapshot + log-tail) + * partitions. If absent, the plan is a pure Fluss upsert scan derived from kv snapshots + log tail. + * + * Startup-mode gating has been removed: a batch upsert scan is always full-table regardless of the + * user-facing SCAN_START_UP_MODE setting — same rationale as [[AppendPlanner]]. + */ +class UpsertPlanner( + override val tablePath: TablePath, + override val tableInfo: TableInfo, + partitionPredicate: Option[FlussPredicate], + pushedPredicate: Option[FlussPredicate], + projection: Array[Int], + options: CaseInsensitiveStringMap, + override val flussConfig: Configuration) + extends AbstractSplitPlanner(tablePath, tableInfo, flussConfig) + with UpsertSplitPlanner { + + private val readableLakeSnapshot: Option[LakeSnapshot] = probeLakeSnapshot() + + override val hasLakeSnapshot: Boolean = readableLakeSnapshot.isDefined + + override protected val stoppingOffsetsInitializer: OffsetsInitializer = + FlussOffsetInitializers.stoppingOffsetsInitializer(true, options, flussConfig) + + // Upsert never pushes a server-side log filter (kv+log union semantics require full log tail + // to be reconciled with kv snapshots — see FlussUpsertPartitionReader). + override val logTailPredicate: Option[FlussPredicate] = None + + override def plan(): Array[InputPartition] = readableLakeSnapshot match { + case Some(snap) => planLakeUnion(snap) + case None => planLogOnly() + } + + // --------------------------------------------------------------------------------------------- + // Log-only branch: Fluss kv snapshots + log tail (no lake involvement). + // --------------------------------------------------------------------------------------------- + + private def planLogOnly(): Array[InputPartition] = { + val bucketOffsetsRetriever = new BucketOffsetsRetrieverImpl(admin, tablePath) + + 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 + } else { + val kvSnapshots = admin.getLatestKvSnapshots(tablePath).get() + createUpsertPartitions(null, kvSnapshots, bucketOffsetsRetriever) + } + } + + private 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 + } + + // --------------------------------------------------------------------------------------------- + // Lake-union branch: lake splits (upsert view) + Fluss log tail after snapshotLogOffset. + // --------------------------------------------------------------------------------------------- + + private def planLakeUnion(snap: LakeSnapshot): Array[InputPartition] = { + val lakeSource = + FlussLakeUtils.createLakeSource(flussConfig.toMap, tableInfo.getProperties.toMap, tablePath) + lakeSource.withProject(FlussLakeUtils.lakeProjection(projection)) + pushedPredicate.foreach(FlussLakeUtils.applyLakeFilters(lakeSource, _)) + + val lakeSplits = lakeSource + .createPlanner(new LakeSource.PlannerContext { + override def snapshotId(): Long = snap.getSnapshotId + }) + .plan() + + val tableBucketsOffset = snap.getTableBucketsOffset + val bucketOffsetsRetriever = new BucketOffsetsRetrieverImpl(admin, tablePath) + + if (tableInfo.isPartitioned) { + planLakePartitionedTable(lakeSplits.asScala.toSeq, tableBucketsOffset, bucketOffsetsRetriever) + } else { + planLakeNonPartitionedTable( + lakeSplits.asScala.toSeq, + tableBucketsOffset, + bucketOffsetsRetriever) + } + } + + private def planLakeNonPartitionedTable( + lakeSplits: Seq[LakeSplit], + tableBucketsOffset: java.util.Map[TableBucket, java.lang.Long], + bucketOffsetsRetriever: BucketOffsetsRetrieverImpl): Array[InputPartition] = { + val tableId = tableInfo.getTableId + val buckets = (0 until tableInfo.getNumBuckets).toSeq + + val stoppingOffsets = + getBucketOffsets(stoppingOffsetsInitializer, null, buckets, bucketOffsetsRetriever) + + val lakeSplitsByBucket = lakeSplits.groupBy(_.bucket()).mapValues(_.toSeq).toMap + + buckets.flatMap { + bucketId => + val tableBucket = new TableBucket(tableId, bucketId) + val snapshotLogOffset = tableBucketsOffset.get(tableBucket) + val stoppingOffset = stoppingOffsets(bucketId) + + createLakeUpsertPartition( + tableBucket, + lakeSplitsByBucket.get(bucketId), + snapshotLogOffset, + stoppingOffset) + }.toArray + } + + private def planLakePartitionedTable( + lakeSplits: Seq[LakeSplit], + tableBucketsOffset: java.util.Map[TableBucket, java.lang.Long], + bucketOffsetsRetriever: BucketOffsetsRetrieverImpl): Array[InputPartition] = { + val tableId = tableInfo.getTableId + val buckets = (0 until tableInfo.getNumBuckets).toSeq + + val filteredPartitionInfos = SparkPartitionPredicate.filterPartitions( + tableInfo, + partitionInfos.asScala.toSeq, + partitionPredicate) + + val flussPartitionIdByName = mutable.LinkedHashMap.empty[String, Long] + filteredPartitionInfos.foreach { + pi => flussPartitionIdByName(pi.getPartitionName) = pi.getPartitionId + } + + val lakeSplitsByPartition = groupLakeSplitsByPartition(lakeSplits) + + val lakePartitions = lakeSplitsByPartition.flatMap { + case (partitionName, (partitionValues, splitsByBucket)) => + flussPartitionIdByName.remove(partitionName) match { + case Some(partitionId) => + val stoppingOffsets = getBucketOffsets( + stoppingOffsetsInitializer, + partitionName, + buckets, + bucketOffsetsRetriever) + + buckets.flatMap { + bucketId => + val tableBucket = new TableBucket(tableId, partitionId, bucketId) + val snapshotLogOffset = tableBucketsOffset.get(tableBucket) + val stoppingOffset = stoppingOffsets(bucketId) + + createLakeUpsertPartition( + tableBucket, + splitsByBucket.get(bucketId), + snapshotLogOffset, + stoppingOffset) + } + + case None => + if ( + SparkPartitionPredicate + .matchesPartition(tableInfo, partitionValues, partitionPredicate) + ) { + buckets.flatMap { + bucketId => + val tableBucket = new TableBucket(tableId, -1, bucketId) + splitsByBucket.getOrElse(bucketId, Seq.empty).map { + lakeSplit => FlussLakeInputPartition(tableBucket, lakeSplit) + } + } + } else { + Seq.empty + } + } + } + + val flussOnlyPartitions = flussPartitionIdByName.flatMap { + case (partitionName, partitionId) => + val stoppingOffsets = getBucketOffsets( + stoppingOffsetsInitializer, + partitionName, + buckets, + bucketOffsetsRetriever) + + buckets.flatMap { + bucketId => + val tableBucket = new TableBucket(tableId, partitionId, bucketId) + val stoppingOffset = stoppingOffsets(bucketId) + + if (stoppingOffset > 0) { + Some( + FlussLakeUpsertInputPartition( + tableBucket, + null, + LogScanner.EARLIEST_OFFSET, + stoppingOffset + )) + } else { + None + } + } + } + + (lakePartitions ++ flussOnlyPartitions).toArray + } + + private def groupLakeSplitsByPartition( + lakeSplits: Seq[LakeSplit]): Map[String, (Seq[String], mutable.Map[Int, Seq[LakeSplit]])] = { + val grouped = + mutable.LinkedHashMap.empty[String, (Seq[String], mutable.Map[Int, Seq[LakeSplit]])] + lakeSplits.foreach { + split => + val partitionValues = + if (split.partition() == null) Seq.empty[String] else split.partition().asScala.toSeq + val partitionName = + if (partitionValues.isEmpty) "" + else partitionValues.mkString(ResolvedPartitionSpec.PARTITION_SPEC_SEPARATOR) + val (_, bucketMap) = + grouped.getOrElseUpdate(partitionName, (partitionValues, mutable.Map.empty)) + val bucketId = split.bucket() + val splits = bucketMap.getOrElse(bucketId, Seq.empty) + bucketMap(bucketId) = splits :+ split + } + grouped.toMap + } + + private def createLakeUpsertPartition( + tableBucket: TableBucket, + lakeSplits: Option[Seq[LakeSplit]], + snapshotLogOffset: java.lang.Long, + stoppingOffset: Long): Option[InputPartition] = { + val needLogSplit = if (snapshotLogOffset == null) { + stoppingOffset > 0 + } else { + snapshotLogOffset < stoppingOffset.longValue() + } + val needLakeSplit = lakeSplits.isDefined && lakeSplits.get.nonEmpty + if (!needLogSplit && !needLakeSplit) { + return None + } + + val lakeSplitList = + if (lakeSplits.isDefined && lakeSplits.get.nonEmpty) { + new java.util.ArrayList[LakeSplit](lakeSplits.get.asJava) + } else null + + val logStartingOffset = + if (snapshotLogOffset != null) snapshotLogOffset.longValue() + else LogScanner.EARLIEST_OFFSET + + Some( + FlussLakeUpsertInputPartition( + tableBucket, + lakeSplitList, + logStartingOffset, + stoppingOffset + )) + } +} diff --git a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/lake/FlussLakeAppendBatch.scala b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/lake/FlussLakeAppendBatch.scala deleted file mode 100644 index 50bfc7f1cd..0000000000 --- a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/lake/FlussLakeAppendBatch.scala +++ /dev/null @@ -1,320 +0,0 @@ -/* - * 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.fluss.spark.read.lake - -import org.apache.fluss.client.initializer.{BucketOffsetsRetrieverImpl, OffsetsInitializer} -import org.apache.fluss.client.table.scanner.log.LogScanner -import org.apache.fluss.config.Configuration -import org.apache.fluss.exception.LakeTableSnapshotNotExistException -import org.apache.fluss.lake.source.{LakeSource, LakeSplit} -import org.apache.fluss.metadata.{LogFormat, ResolvedPartitionSpec, TableBucket, TableInfo, TablePath} -import org.apache.fluss.predicate.{Predicate => FlussPredicate} -import org.apache.fluss.spark.read._ -import org.apache.fluss.spark.utils.SparkPartitionPredicate -import org.apache.fluss.utils.ExceptionUtils - -import org.apache.spark.sql.connector.read.{InputPartition, PartitionReaderFactory} -import org.apache.spark.sql.types.StructType -import org.apache.spark.sql.util.CaseInsensitiveStringMap - -import scala.collection.JavaConverters._ -import scala.collection.mutable - -/** Batch for reading lake-enabled log table (append-only table with datalake). */ -class FlussLakeAppendBatch( - tablePath: TablePath, - tableInfo: TableInfo, - readSchema: StructType, - pushedPredicate: Option[FlussPredicate], - partitionPredicate: Option[FlussPredicate], - limit: Option[Int], - options: CaseInsensitiveStringMap, - flussConfig: Configuration) - extends FlussLakeBatch(tablePath, tableInfo, readSchema, limit, options, flussConfig) { - - // Required by FlussLakeBatch but unused — lake snapshot determines start offsets. - override val startOffsetsInitializer: OffsetsInitializer = OffsetsInitializer.earliest() - - // Server-side log filter requires ARROW format. - private val logTailPredicate: Option[FlussPredicate] = - if (tableInfo.getTableConfig.getLogFormat == LogFormat.ARROW) pushedPredicate else None - - override def createReaderFactory(): PartitionReaderFactory = { - if (isFallback) { - new FlussAppendPartitionReaderFactory( - tablePath, - projection, - logTailPredicate, - limit, - options, - flussConfig) - } else { - new FlussLakePartitionReaderFactory( - tableInfo.getProperties.toMap, - tablePath, - projection, - pushedPredicate, - logTailPredicate, - limit, - flussConfig) - } - } - - override def doPlan(): (Array[InputPartition], Boolean) = { - val lakeSnapshot = - try { - admin.getReadableLakeSnapshot(tablePath).get() - } catch { - case e: Exception => - if ( - ExceptionUtils - .stripExecutionException(e) - .isInstanceOf[LakeTableSnapshotNotExistException] - ) { - return (planFallbackPartitions(), true) - } - throw e - } - - val lakeSource = - FlussLakeUtils.createLakeSource(flussConfig.toMap, tableInfo.getProperties.toMap, tablePath) - lakeSource.withProject(FlussLakeUtils.lakeProjection(projection)) - pushedPredicate.foreach(FlussLakeBatch.applyLakeFilters(lakeSource, _)) - - val lakeSplits = lakeSource - .createPlanner(new LakeSource.PlannerContext { - override def snapshotId(): Long = lakeSnapshot.getSnapshotId - }) - .plan() - - val tableBucketsOffset = lakeSnapshot.getTableBucketsOffset - val buckets = (0 until tableInfo.getNumBuckets).toSeq - val bucketOffsetsRetriever = new BucketOffsetsRetrieverImpl(admin, tablePath) - - val partitions = if (tableInfo.isPartitioned) { - planPartitionedTable( - lakeSplits.asScala.toSeq, - tableBucketsOffset, - buckets, - bucketOffsetsRetriever) - } else { - planNonPartitionedTable( - lakeSplits.asScala.toSeq, - tableBucketsOffset, - buckets, - bucketOffsetsRetriever) - } - - (partitions, false) - } - - private def planNonPartitionedTable( - lakeSplits: Seq[LakeSplit], - tableBucketsOffset: java.util.Map[TableBucket, java.lang.Long], - buckets: Seq[Int], - bucketOffsetsRetriever: BucketOffsetsRetrieverImpl): Array[InputPartition] = { - val tableId = tableInfo.getTableId - - val lakePartitions = - createLakePartitions(lakeSplits, tableId, partitionId = None) - - val stoppingOffsets = - getBucketOffsets(stoppingOffsetsInitializer, null, buckets, bucketOffsetsRetriever) - val logPartitions = buckets.flatMap { - bucketId => - val tableBucket = new TableBucket(tableId, bucketId) - createLogTailPartition(tableBucket, tableBucketsOffset, stoppingOffsets(bucketId)) - } - - (lakePartitions ++ logPartitions).toArray - } - - private def planPartitionedTable( - lakeSplits: Seq[LakeSplit], - tableBucketsOffset: java.util.Map[TableBucket, java.lang.Long], - buckets: Seq[Int], - bucketOffsetsRetriever: BucketOffsetsRetrieverImpl): Array[InputPartition] = { - val tableId = tableInfo.getTableId - - // Filter Fluss-known partitions using the partition predicate to skip non-matching ones - val filteredPartitionInfos = SparkPartitionPredicate.filterPartitions( - tableInfo, - partitionInfos.asScala.toSeq, - partitionPredicate) - - val flussPartitionIdByName = mutable.LinkedHashMap.empty[String, Long] - filteredPartitionInfos.foreach { - pi => flussPartitionIdByName(pi.getPartitionName) = pi.getPartitionId - } - - val lakeSplitsByPartition = groupLakeSplitsByPartition(lakeSplits) - var lakeSplitPartitionId = -1L - - val lakeAndLogPartitions = lakeSplitsByPartition.flatMap { - case (partitionName, (partitionValues, splits)) => - flussPartitionIdByName.remove(partitionName) match { - case Some(partitionId) => - // Partition in both lake and Fluss — lake splits + log tail - val lakePartitions = - createLakePartitions(splits.toSeq, tableId, Some(partitionId)) - - val stoppingOffsets = getBucketOffsets( - stoppingOffsetsInitializer, - partitionName, - buckets, - bucketOffsetsRetriever) - val logPartitions = buckets.flatMap { - bucketId => - val tableBucket = new TableBucket(tableId, partitionId, bucketId) - createLogTailPartition(tableBucket, tableBucketsOffset, stoppingOffsets(bucketId)) - } - - lakePartitions ++ logPartitions - - case None => - // Partition only in lake (expired in Fluss). Apply the partition predicate directly - // on the resolved partition values to avoid round-tripping through partition names. - if ( - SparkPartitionPredicate - .matchesPartition(tableInfo, partitionValues, partitionPredicate) - ) { - val pid = lakeSplitPartitionId - lakeSplitPartitionId -= 1 - createLakePartitions(splits.toSeq, tableId, Some(pid)) - } else { - Seq.empty - } - } - }.toSeq - - // Partitions only in Fluss (not yet tiered) — log from earliest - val flussOnlyPartitions = flussPartitionIdByName.flatMap { - case (partitionName, partitionId) => - val stoppingOffsets = getBucketOffsets( - stoppingOffsetsInitializer, - partitionName, - buckets, - bucketOffsetsRetriever) - buckets.flatMap { - bucketId => - val stoppingOffset = stoppingOffsets(bucketId) - if (stoppingOffset > 0) { - val tableBucket = new TableBucket(tableId, partitionId, bucketId) - Some( - FlussAppendInputPartition( - tableBucket, - LogScanner.EARLIEST_OFFSET, - stoppingOffset): InputPartition) - } else { - None - } - } - }.toSeq - - (lakeAndLogPartitions ++ flussOnlyPartitions).toArray - } - - private def groupLakeSplitsByPartition(lakeSplits: Seq[LakeSplit]) - : mutable.LinkedHashMap[String, (Seq[String], mutable.ArrayBuffer[LakeSplit])] = { - val grouped = - mutable.LinkedHashMap.empty[String, (Seq[String], mutable.ArrayBuffer[LakeSplit])] - lakeSplits.foreach { - split => - val partitionValues = - if (split.partition() == null) Seq.empty[String] else split.partition().asScala.toSeq - val partitionName = - if (partitionValues.isEmpty) "" - else partitionValues.mkString(ResolvedPartitionSpec.PARTITION_SPEC_SEPARATOR) - val (_, buf) = - grouped.getOrElseUpdate(partitionName, (partitionValues, mutable.ArrayBuffer.empty)) - buf += split - } - grouped - } - - private def createLakePartitions( - splits: Seq[LakeSplit], - tableId: Long, - partitionId: Option[Long]): Seq[InputPartition] = { - splits.map { - split => - val tableBucket = partitionId match { - case Some(pid) => new TableBucket(tableId, pid, split.bucket()) - case None => new TableBucket(tableId, split.bucket()) - } - FlussLakeInputPartition(tableBucket, split) - } - } - - private def createLogTailPartition( - tableBucket: TableBucket, - tableBucketsOffset: java.util.Map[TableBucket, java.lang.Long], - stoppingOffset: Long): Option[InputPartition] = { - val snapshotLogOffset = tableBucketsOffset.get(tableBucket) - if (snapshotLogOffset != null) { - if (snapshotLogOffset.longValue() < stoppingOffset) { - Some(FlussAppendInputPartition(tableBucket, snapshotLogOffset.longValue(), stoppingOffset)) - } else { - None - } - } else if (stoppingOffset > 0) { - Some(FlussAppendInputPartition(tableBucket, LogScanner.EARLIEST_OFFSET, stoppingOffset)) - } else { - None - } - } - - private def planFallbackPartitions(): Array[InputPartition] = { - val fallbackStartInit = FlussOffsetInitializers.startOffsetsInitializer(options, flussConfig) - val bucketOffsetsRetriever = new BucketOffsetsRetrieverImpl(admin, tablePath) - val buckets = (0 until tableInfo.getNumBuckets).toSeq - val tableId = tableInfo.getTableId - - def createPartitions( - partitionId: Option[Long], - partitionName: String): Array[InputPartition] = { - val startOffsets = - getBucketOffsets(fallbackStartInit, partitionName, buckets, bucketOffsetsRetriever) - val stoppingOffsets = - getBucketOffsets(stoppingOffsetsInitializer, partitionName, buckets, bucketOffsetsRetriever) - - buckets.map { - bucketId => - val tableBucket = partitionId match { - case Some(pid) => new TableBucket(tableId, pid, bucketId) - case None => new TableBucket(tableId, bucketId) - } - FlussAppendInputPartition( - tableBucket, - startOffsets(bucketId), - stoppingOffsets(bucketId) - ): InputPartition - }.toArray - } - - if (tableInfo.isPartitioned) { - val matching = SparkPartitionPredicate.filterPartitions( - tableInfo, - partitionInfos.asScala.toSeq, - partitionPredicate) - matching.flatMap(pi => createPartitions(Some(pi.getPartitionId), pi.getPartitionName)).toArray - } else { - createPartitions(None, null) - } - } -} diff --git a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/lake/FlussLakeBatch.scala b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/lake/FlussLakeBatch.scala deleted file mode 100644 index bed1017b87..0000000000 --- a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/lake/FlussLakeBatch.scala +++ /dev/null @@ -1,88 +0,0 @@ -/* - * 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.fluss.spark.read.lake - -import org.apache.fluss.client.initializer.{BucketOffsetsRetrieverImpl, OffsetsInitializer} -import org.apache.fluss.config.Configuration -import org.apache.fluss.lake.source.{LakeSource, LakeSplit} -import org.apache.fluss.metadata.{TableInfo, TablePath} -import org.apache.fluss.predicate.{Predicate => FlussPredicate} -import org.apache.fluss.spark.read._ - -import org.apache.spark.internal.Logging -import org.apache.spark.sql.connector.read.InputPartition -import org.apache.spark.sql.types.StructType -import org.apache.spark.sql.util.CaseInsensitiveStringMap - -import java.util.Collections - -import scala.collection.JavaConverters._ - -abstract class FlussLakeBatch( - tablePath: TablePath, - tableInfo: TableInfo, - readSchema: StructType, - limit: Option[Int], - options: CaseInsensitiveStringMap, - flussConfig: Configuration) - extends FlussBatch(tablePath, tableInfo, readSchema, limit, flussConfig) { - - override val stoppingOffsetsInitializer: OffsetsInitializer = { - FlussOffsetInitializers.stoppingOffsetsInitializer(true, options, flussConfig) - } - - lazy val (partitions, isFallback) = doPlan() - - override def planInputPartitions(): Array[InputPartition] = partitions - - /** - * Plans input partitions for reading. The returned isFallback flag is true when no lake snapshot - * exists and the plan falls back to reading directly from Fluss. - */ - def doPlan(): (Array[InputPartition], Boolean) - - def getBucketOffsets( - initializer: OffsetsInitializer, - partitionName: String, - buckets: Seq[Int], - bucketOffsetsRetriever: BucketOffsetsRetrieverImpl): Map[Int, Long] = { - initializer - .getBucketOffsets(partitionName, buckets.map(Integer.valueOf).asJava, bucketOffsetsRetriever) - .asScala - .map(e => (e._1.intValue(), Long2long(e._2))) - .toMap - } -} - -object FlussLakeBatch extends Logging { - - def applyLakeFilters( - lakeSource: LakeSource[LakeSplit], - predicates: java.util.List[FlussPredicate]): LakeSource.FilterPushDownResult = { - val result = lakeSource.withFilters(predicates) - logInfo( - s"Lake source accepted ${result.acceptedPredicates()}, " + - s"remaining ${result.remainingPredicates()}") - result - } - - def applyLakeFilters( - lakeSource: LakeSource[LakeSplit], - predicate: FlussPredicate): LakeSource.FilterPushDownResult = - applyLakeFilters(lakeSource, Collections.singletonList(predicate)) -} diff --git a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/lake/FlussLakePartitionReaderFactory.scala b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/lake/FlussLakePartitionReaderFactory.scala index d69b4c08da..fc17bb6d31 100644 --- a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/lake/FlussLakePartitionReaderFactory.scala +++ b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/lake/FlussLakePartitionReaderFactory.scala @@ -42,7 +42,7 @@ class FlussLakePartitionReaderFactory( @transient private lazy val lakeSource: LakeSource[LakeSplit] = { val source = FlussLakeUtils.createLakeSource(flussConfig.toMap, tableProperties, tablePath) source.withProject(FlussLakeUtils.lakeProjection(projection)) - flussPredicate.foreach(FlussLakeBatch.applyLakeFilters(source, _)) + flussPredicate.foreach(FlussLakeUtils.applyLakeFilters(source, _)) source } diff --git a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/lake/FlussLakeUpsertBatch.scala b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/lake/FlussLakeUpsertBatch.scala deleted file mode 100644 index c37a7ec483..0000000000 --- a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/lake/FlussLakeUpsertBatch.scala +++ /dev/null @@ -1,316 +0,0 @@ -/* - * 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.fluss.spark.read.lake - -import org.apache.fluss.client.initializer.{BucketOffsetsRetrieverImpl, OffsetsInitializer, SnapshotOffsetsInitializer} -import org.apache.fluss.client.table.scanner.log.LogScanner -import org.apache.fluss.config.Configuration -import org.apache.fluss.exception.LakeTableSnapshotNotExistException -import org.apache.fluss.lake.source.LakeSplit -import org.apache.fluss.metadata.{ResolvedPartitionSpec, TableBucket, TableInfo, TablePath} -import org.apache.fluss.predicate.{Predicate => FlussPredicate} -import org.apache.fluss.spark.read._ -import org.apache.fluss.spark.utils.SparkPartitionPredicate -import org.apache.fluss.utils.ExceptionUtils - -import org.apache.spark.sql.connector.read.{InputPartition, PartitionReaderFactory} -import org.apache.spark.sql.types.StructType -import org.apache.spark.sql.util.CaseInsensitiveStringMap - -import scala.collection.JavaConverters._ -import scala.collection.mutable - -/** - * Batch for reading lake-enabled primary key tables. Combines lake snapshot data with Fluss kv - * tail, merging them using sort-merge algorithm. - */ -class FlussLakeUpsertBatch( - tablePath: TablePath, - tableInfo: TableInfo, - readSchema: StructType, - pushedPredicate: Option[FlussPredicate], - partitionPredicate: Option[FlussPredicate], - limit: Option[Int], - options: CaseInsensitiveStringMap, - flussConfig: Configuration) - extends FlussLakeBatch(tablePath, tableInfo, readSchema, limit, options, flussConfig) { - - override val startOffsetsInitializer: OffsetsInitializer = { - val offsetsInitializer = FlussOffsetInitializers.startOffsetsInitializer(options, flussConfig) - if (!offsetsInitializer.isInstanceOf[SnapshotOffsetsInitializer]) { - throw new UnsupportedOperationException("Upsert scan only supports FULL startup mode.") - } - offsetsInitializer - } - - override def createReaderFactory(): PartitionReaderFactory = { - if (isFallback) { - new FlussUpsertPartitionReaderFactory(tablePath, projection, limit, options, flussConfig) - } else { - // PK kv-tail reader does not consume server-side log filters. - new FlussLakePartitionReaderFactory( - tableInfo.getProperties.toMap, - tablePath, - projection, - pushedPredicate, - None, - limit, - flussConfig) - } - } - - override def doPlan(): (Array[InputPartition], Boolean) = { - val lakeSnapshot = - try { - admin.getReadableLakeSnapshot(tablePath).get() - } catch { - case e: Exception => - if ( - ExceptionUtils - .stripExecutionException(e) - .isInstanceOf[LakeTableSnapshotNotExistException] - ) { - return (planFallbackPartitions(), true) - } - throw e - } - - val lakeSource = - FlussLakeUtils.createLakeSource(flussConfig.toMap, tableInfo.getProperties.toMap, tablePath) - lakeSource.withProject(FlussLakeUtils.lakeProjection(projection)) - pushedPredicate.foreach(FlussLakeBatch.applyLakeFilters(lakeSource, _)) - - val lakeSplits = lakeSource - .createPlanner(() => lakeSnapshot.getSnapshotId) - .plan() - - val tableBucketsOffset = lakeSnapshot.getTableBucketsOffset - val bucketOffsetsRetriever = new BucketOffsetsRetrieverImpl(admin, tablePath) - - val partitions = if (tableInfo.isPartitioned) { - planPartitionedTable(lakeSplits.asScala.toSeq, tableBucketsOffset, bucketOffsetsRetriever) - } else { - planNonPartitionedTable(lakeSplits.asScala.toSeq, tableBucketsOffset, bucketOffsetsRetriever) - } - - (partitions, false) - } - - private def planNonPartitionedTable( - lakeSplits: Seq[LakeSplit], - tableBucketsOffset: java.util.Map[TableBucket, java.lang.Long], - bucketOffsetsRetriever: BucketOffsetsRetrieverImpl): Array[InputPartition] = { - val tableId = tableInfo.getTableId - val buckets = (0 until tableInfo.getNumBuckets).toSeq - - val stoppingOffsets = - getBucketOffsets(stoppingOffsetsInitializer, null, buckets, bucketOffsetsRetriever) - - // Group lake splits by bucket - val lakeSplitsByBucket = lakeSplits.groupBy(_.bucket()).mapValues(_.toSeq).toMap - - buckets.flatMap { - bucketId => - val tableBucket = new TableBucket(tableId, bucketId) - val snapshotLogOffset = tableBucketsOffset.get(tableBucket) - val stoppingOffset = stoppingOffsets(bucketId) - - createLakeUpsertPartition( - tableBucket, - lakeSplitsByBucket.get(bucketId), - snapshotLogOffset, - stoppingOffset) - }.toArray - } - - private def planPartitionedTable( - lakeSplits: Seq[LakeSplit], - tableBucketsOffset: java.util.Map[TableBucket, java.lang.Long], - bucketOffsetsRetriever: BucketOffsetsRetrieverImpl): Array[InputPartition] = { - val tableId = tableInfo.getTableId - val buckets = (0 until tableInfo.getNumBuckets).toSeq - - // Filter Fluss-known partitions using the partition predicate to skip non-matching ones - val filteredPartitionInfos = SparkPartitionPredicate.filterPartitions( - tableInfo, - partitionInfos.asScala.toSeq, - partitionPredicate) - - val flussPartitionIdByName = mutable.LinkedHashMap.empty[String, Long] - filteredPartitionInfos.foreach { - pi => flussPartitionIdByName(pi.getPartitionName) = pi.getPartitionId - } - - val lakeSplitsByPartition = groupLakeSplitsByPartition(lakeSplits) - - val lakePartitions = lakeSplitsByPartition.flatMap { - case (partitionName, (partitionValues, splitsByBucket)) => - flussPartitionIdByName.remove(partitionName) match { - case Some(partitionId) => - // Partition in both lake and Fluss (already passed the predicate filter above) - val stoppingOffsets = getBucketOffsets( - stoppingOffsetsInitializer, - partitionName, - buckets, - bucketOffsetsRetriever) - - buckets.flatMap { - bucketId => - val tableBucket = new TableBucket(tableId, partitionId, bucketId) - val snapshotLogOffset = tableBucketsOffset.get(tableBucket) - val stoppingOffset = stoppingOffsets(bucketId) - - createLakeUpsertPartition( - tableBucket, - splitsByBucket.get(bucketId), - snapshotLogOffset, - stoppingOffset) - } - - case None => - // Partition only in lake (expired in Fluss). Apply the partition predicate directly - // on the resolved partition values to avoid round-tripping through partition names. - if ( - SparkPartitionPredicate - .matchesPartition(tableInfo, partitionValues, partitionPredicate) - ) { - buckets.flatMap { - bucketId => - val tableBucket = new TableBucket(tableId, -1, bucketId) - splitsByBucket.getOrElse(bucketId, Seq.empty).map { - lakeSplit => FlussLakeInputPartition(tableBucket, lakeSplit) - } - } - } else { - Seq.empty - } - } - } - - // Partitions only in Fluss (not yet tiered) - read from earliest - val flussOnlyPartitions = flussPartitionIdByName.flatMap { - case (partitionName, partitionId) => - val stoppingOffsets = getBucketOffsets( - stoppingOffsetsInitializer, - partitionName, - buckets, - bucketOffsetsRetriever) - - buckets.flatMap { - bucketId => - val tableBucket = new TableBucket(tableId, partitionId, bucketId) - val stoppingOffset = stoppingOffsets(bucketId) - - // No lake snapshot for this bucket, skip if there's no data to read - if (stoppingOffset > 0) { - Some( - FlussLakeUpsertInputPartition( - tableBucket, - null, // no lake splits - LogScanner.EARLIEST_OFFSET, - stoppingOffset - )) - } else { - None - } - } - } - - (lakePartitions ++ flussOnlyPartitions).toArray - } - - /** - * Group lake splits by partition. Each entry stores the resolved partition values along with - * splits keyed by bucket id, so callers can both look up Fluss partitions by name and evaluate - * the partition predicate directly on the values without re-parsing the joined name. - */ - private def groupLakeSplitsByPartition( - lakeSplits: Seq[LakeSplit]): Map[String, (Seq[String], mutable.Map[Int, Seq[LakeSplit]])] = { - val grouped = - mutable.LinkedHashMap.empty[String, (Seq[String], mutable.Map[Int, Seq[LakeSplit]])] - lakeSplits.foreach { - split => - val partitionValues = - if (split.partition() == null) Seq.empty[String] else split.partition().asScala.toSeq - val partitionName = - if (partitionValues.isEmpty) "" - else partitionValues.mkString(ResolvedPartitionSpec.PARTITION_SPEC_SEPARATOR) - val (_, bucketMap) = - grouped.getOrElseUpdate(partitionName, (partitionValues, mutable.Map.empty)) - val bucketId = split.bucket() - val splits = bucketMap.getOrElse(bucketId, Seq.empty) - bucketMap(bucketId) = splits :+ split - } - grouped.toMap - } - - private def createLakeUpsertPartition( - tableBucket: TableBucket, - lakeSplits: Option[Seq[LakeSplit]], - snapshotLogOffset: java.lang.Long, - stoppingOffset: Long): Option[InputPartition] = { - val needLogSplit = if (snapshotLogOffset == null) { - stoppingOffset > 0 - } else { - snapshotLogOffset < stoppingOffset.longValue() - } - val needLakeSplit = lakeSplits.isDefined && lakeSplits.get.nonEmpty - if (!needLogSplit && !needLakeSplit) { - return None - } - - val lakeSplitList = - if (lakeSplits.isDefined && lakeSplits.get.nonEmpty) { - new java.util.ArrayList[LakeSplit](lakeSplits.get.asJava) - } else null - - val logStartingOffset = - if (snapshotLogOffset != null) snapshotLogOffset.longValue() - else LogScanner.EARLIEST_OFFSET - - Some( - FlussLakeUpsertInputPartition( - tableBucket, - lakeSplitList, - logStartingOffset, - stoppingOffset - )) - } - - private def planFallbackPartitions(): Array[InputPartition] = { - val bucketOffsetsRetriever = new BucketOffsetsRetrieverImpl(admin, tablePath) - - if (tableInfo.isPartitioned) { - // Filter partitions using partition predicate early to skip non-matching partitions - val filteredPartitionInfos = SparkPartitionPredicate.filterPartitions( - tableInfo, - partitionInfos.asScala.toSeq, - partitionPredicate) - - filteredPartitionInfos.flatMap { - pi => - val partitionName = pi.getPartitionName - val kvSnapshots = admin.getLatestKvSnapshots(tablePath, partitionName).get() - createUpsertPartitions(partitionName, kvSnapshots, bucketOffsetsRetriever) - }.toArray - } else { - val kvSnapshots = admin.getLatestKvSnapshots(tablePath).get() - createUpsertPartitions(null, kvSnapshots, bucketOffsetsRetriever) - } - } -} diff --git a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/lake/FlussLakeUtils.scala b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/lake/FlussLakeUtils.scala index 0bdcb8721f..9e450ef04b 100644 --- a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/lake/FlussLakeUtils.scala +++ b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/lake/FlussLakeUtils.scala @@ -21,13 +21,17 @@ import org.apache.fluss.config.{ConfigOptions, Configuration} import org.apache.fluss.lake.lakestorage.LakeStoragePluginSetUp import org.apache.fluss.lake.source.{LakeSource, LakeSplit} import org.apache.fluss.metadata.TablePath +import org.apache.fluss.predicate.{Predicate => FlussPredicate} import org.apache.fluss.utils.PropertiesUtils +import org.apache.spark.internal.Logging + import java.util +import java.util.Collections import scala.collection.JavaConverters._ -object FlussLakeUtils { +object FlussLakeUtils extends Logging { /** * Paimon DLF supports multiple authentication schemes (see @@ -69,4 +73,19 @@ object FlussLakeUtils { def lakeProjection(projection: Array[Int]): Array[Array[Int]] = { projection.map(i => Array(i)) } + + def applyLakeFilters( + lakeSource: LakeSource[LakeSplit], + predicates: java.util.List[FlussPredicate]): LakeSource.FilterPushDownResult = { + val result = lakeSource.withFilters(predicates) + logInfo( + s"Lake source accepted ${result.acceptedPredicates()}, " + + s"remaining ${result.remainingPredicates()}") + result + } + + def applyLakeFilters( + lakeSource: LakeSource[LakeSplit], + predicate: FlussPredicate): LakeSource.FilterPushDownResult = + applyLakeFilters(lakeSource, Collections.singletonList(predicate)) } diff --git a/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/SparkPrimaryKeyTableReadTest.scala b/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/SparkPrimaryKeyTableReadTest.scala index 4e30d552d1..28d245efcf 100644 --- a/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/SparkPrimaryKeyTableReadTest.scala +++ b/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/SparkPrimaryKeyTableReadTest.scala @@ -135,13 +135,6 @@ class SparkPrimaryKeyTableReadTest extends FlussSparkTestBase { Row(800L, 23L, "addr3") :: Nil ) - - // Only support FULL startup mode. - withSQLConf( - s"${SparkFlussConf.SPARK_FLUSS_CONF_PREFIX}${SparkFlussConf.SCAN_START_UP_MODE.key()}" -> "latest") { - intercept[UnsupportedOperationException]( - sql(s"SELECT * FROM $DEFAULT_DATABASE.t ORDER BY orderId").show()) - } } } diff --git a/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/lake/SparkLakeTableReadTestBase.scala b/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/lake/SparkLakeTableReadTestBase.scala index 5c507372d4..21a7a65351 100644 --- a/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/lake/SparkLakeTableReadTestBase.scala +++ b/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/lake/SparkLakeTableReadTestBase.scala @@ -22,7 +22,7 @@ import org.apache.fluss.flink.tiering.LakeTieringJobBuilder import org.apache.fluss.flink.tiering.source.TieringSourceOptions import org.apache.fluss.metadata.{DataLakeFormat, TableBucket} import org.apache.fluss.spark.FlussSparkTestBase -import org.apache.fluss.spark.read.{FlussLakeAppendScan, FlussLakeUpsertScan, FlussScan} +import org.apache.fluss.spark.read.{FlussAppendScan, FlussScan, FlussUpsertScan} import org.apache.flink.api.common.RuntimeExecutionMode import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment @@ -166,8 +166,8 @@ abstract class SparkLakeTableReadTestBase extends FlussSparkTestBase { } scans .collect { - case upsert: FlussLakeUpsertScan => upsert - case append: FlussLakeAppendScan => append + case upsert: FlussUpsertScan => upsert + case append: FlussAppendScan => append } .flatMap(_.toBatch.planInputPartitions()) .toArray