Support Trigger.AvailableNow

client/src/main/scala/io/delta/sharing/spark/DeltaSharingSource.scala

@@ -26,7 +26,7 @@ import org.apache.spark.delta.sharing.{CachedTableManager, TableRefreshResult}
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.{DataFrame, DeltaSharingScanUtils, SparkSession}
 import org.apache.spark.sql.connector.read.streaming
-import org.apache.spark.sql.connector.read.streaming.{ReadAllAvailable, ReadLimit, ReadMaxFiles, SupportsAdmissionControl}
+import org.apache.spark.sql.connector.read.streaming.{ReadAllAvailable, ReadLimit, ReadMaxFiles, SupportsAdmissionControl, SupportsTriggerAvailableNow}
 import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelation}
 import org.apache.spark.sql.execution.streaming._
 import org.apache.spark.sql.types.StructType
@@ -95,6 +95,7 @@ case class DeltaSharingSource(
   deltaLog: RemoteDeltaLog,
   options: DeltaSharingOptions) extends Source
   with SupportsAdmissionControl
+  with SupportsTriggerAvailableNow
   with Logging {
 
   // This is to ensure that the request sent from the client contains the http header for streaming.
@@ -120,6 +121,33 @@ case class DeltaSharingSource(
   // This is checked before creating DeltaSharingSource
   assert(schema.nonEmpty, "schema cannot be empty in DeltaSharingSource.")
 
+  private var isTriggerAvailableNow = false
+
+  private var isLastOffsetForTriggerAvailableNowInitialized = false
+
+  private var lastOffsetForTriggerAvailableNow: Option[DeltaSharingSourceOffset] = None
+
+  override def prepareForTriggerAvailableNow(): Unit = {
+    logInfo(s"The streaming query reports to use Trigger.AvailableNow")
+    isTriggerAvailableNow = true
+  }
+
+  private def initForTriggerAvailableNowIfNeeded(): Unit = {
+    if (isTriggerAvailableNow && !isLastOffsetForTriggerAvailableNowInitialized) {
+      isLastOffsetForTriggerAvailableNowInitialized = true
+      initLastOffsetForTriggerAvailableNow()
+    }
+  }
+
+  private def initLastOffsetForTriggerAvailableNow(): Unit = {
+    val offset = latestOffsetInternal(ReadLimit.allAvailable())
+    if (offset != null) {
+      lastOffsetForTriggerAvailableNow = Some(DeltaSharingSourceOffset(tableId, offset))
+      logInfo("lastOffset for Trigger.AvailableNow has set to " +
+        s"${lastOffsetForTriggerAvailableNow.get.json}")
+    }
+  }
+
   /** A check on the source table that skips commits that contain removes from the set of files. */
   private val skipChangeCommits = options.skipChangeCommits
 
@@ -286,6 +314,13 @@ case class DeltaSharingSource(
     } else {
       getCDFFileChanges(fromVersion, fromIndex, endingVersionForQuery)
     }
+
+    lastOffsetForTriggerAvailableNow.foreach { bound =>
+      sortedFetchedFiles = sortedFetchedFiles.filter {
+        case IndexedFile(version, index, _, _, _, _, _) =>
+          version < bound.tableVersion || (version == bound.tableVersion && index <= bound.index)
+      }
+    }
   }
 
   private def resetGlobalTimestamp(): Unit = {
@@ -996,6 +1031,11 @@ case class DeltaSharingSource(
   }
 
   override def latestOffset(startOffset: streaming.Offset, limit: ReadLimit): streaming.Offset = {
+    initForTriggerAvailableNowIfNeeded()
+    latestOffsetInternal(limit)
+  }
+
+  private def latestOffsetInternal(limit: ReadLimit): streaming.Offset = {
     val limits = AdmissionLimits(limit)
 
     val currentOffset = if (previousOffset == null) {
@@ -1017,6 +1057,8 @@ case class DeltaSharingSource(
       getTableInfoForLogging)
     val endOffset = DeltaSharingSourceOffset(tableId, end)
 
+    initForTriggerAvailableNowIfNeeded()
+
     val (startVersion, startIndex, isStartingVersion, startSourceVersion) = if (
       startOffsetOption.isEmpty) {
       getStartingVersion match {