[Kernel] [Pagination] Add PaginationScanImpl Class (delta-io#4880)

Author: mmmyr

kernel/kernel-api/src/main/java/io/delta/kernel/internal/PaginatedScanImpl.java

@@ -0,0 +1,67 @@
+/*
+ * Copyright (2025) The Delta Lake Project Authors.
+ *
+ * Licensed 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 io.delta.kernel.internal;
+
+import io.delta.kernel.PaginatedScan;
+import io.delta.kernel.PaginatedScanFilesIterator;
+import io.delta.kernel.data.FilteredColumnarBatch;
+import io.delta.kernel.data.Row;
+import io.delta.kernel.engine.Engine;
+import io.delta.kernel.expressions.Predicate;
+import io.delta.kernel.internal.replay.PageToken;
+import io.delta.kernel.internal.replay.PaginatedScanFilesIteratorImpl;
+import io.delta.kernel.internal.replay.PaginationContext;
+import io.delta.kernel.utils.CloseableIterator;
+import java.util.Optional;
+
+/** Implementation of {@link PaginatedScan} */
+public class PaginatedScanImpl implements PaginatedScan {
+  private final long pageSize;
+  private final Optional<PageToken> pageTokenOpt;
+  private final ScanImpl baseScan;
+
+  public PaginatedScanImpl(ScanImpl baseScan, Optional<Row> pageTokenRowOpt, long pageSize) {
+    this.baseScan = baseScan;
+    this.pageTokenOpt = pageTokenRowOpt.map(PageToken::fromRow);
+    this.pageSize = pageSize;
+  }
+
+  @Override
+  public Optional<Predicate> getRemainingFilter() {
+    return baseScan.getRemainingFilter();
+  }
+
+  @Override
+  public Row getScanState(Engine engine) {
+    return baseScan.getScanState(engine);
+  }
+
+  @Override
+  public PaginatedScanFilesIterator getScanFiles(Engine engine) {
+    return this.getScanFiles(engine, false /* include stats */);
+  }
+
+  public PaginatedScanFilesIterator getScanFiles(Engine engine, boolean includeStates) {
+    PaginationContext paginationContext =
+        pageTokenOpt
+            .map(token -> PaginationContext.forPageWithPageToken(pageSize, token))
+            .orElseGet(() -> PaginationContext.forFirstPage(pageSize));
+    CloseableIterator<FilteredColumnarBatch> filteredScanFilesIter =
+        baseScan.getScanFiles(engine, includeStates, Optional.of(paginationContext));
+    return new PaginatedScanFilesIteratorImpl(filteredScanFilesIter, paginationContext);
+  }
+}

kernel/kernel-api/src/main/java/io/delta/kernel/internal/ScanBuilderImpl.java

@@ -17,7 +17,6 @@
 package io.delta.kernel.internal;
 
 import io.delta.kernel.PaginatedScan;
-import io.delta.kernel.Scan;
 import io.delta.kernel.ScanBuilder;
 import io.delta.kernel.data.Row;
 import io.delta.kernel.expressions.Predicate;
@@ -76,7 +75,7 @@ public ScanBuilder withReadSchema(StructType readSchema) {
   }
 
   @Override
-  public Scan build() {
+  public ScanImpl build() {
     return new ScanImpl(
         snapshotSchema,
         readSchema,
@@ -89,7 +88,8 @@ public Scan build() {
   }
 
   @Override
-  public PaginatedScan buildPaginated(long pageSize, Optional<Row> pageToken) {
-    throw new UnsupportedOperationException("not implemented");
+  public PaginatedScan buildPaginated(long pageSize, Optional<Row> pageTokenRowOpt) {
+    ScanImpl baseScan = this.build();
+    return new PaginatedScanImpl(baseScan, pageTokenRowOpt, pageSize);
   }
 }

kernel/kernel-api/src/main/java/io/delta/kernel/internal/ScanImpl.java

@@ -34,6 +34,7 @@
 import io.delta.kernel.internal.metrics.ScanReportImpl;
 import io.delta.kernel.internal.metrics.Timer;
 import io.delta.kernel.internal.replay.LogReplay;
+import io.delta.kernel.internal.replay.PaginationContext;
 import io.delta.kernel.internal.skipping.DataSkippingPredicate;
 import io.delta.kernel.internal.skipping.DataSkippingUtils;
 import io.delta.kernel.internal.util.*;
@@ -102,7 +103,17 @@ public ScanImpl(
    */
   @Override
   public CloseableIterator<FilteredColumnarBatch> getScanFiles(Engine engine) {
-    return getScanFiles(engine, false);
+    return getScanFiles(engine, false /* includeStats */);
+  }
+
+  /**
+   * Get an iterator of data files in this version of scan that survived the predicate pruning.
+   *
+   * @return data in {@link ColumnarBatch} batch format. Each row correspond to one survived file.
+   */
+  public CloseableIterator<FilteredColumnarBatch> getScanFiles(
+      Engine engine, boolean includeStats) {
+    return getScanFiles(engine, includeStats, Optional.empty() /* paginationContextOpt */);
   }
 
   /**
@@ -115,10 +126,11 @@ public CloseableIterator<FilteredColumnarBatch> getScanFiles(Engine engine) {
    *
    * @param engine the {@link Engine} instance to use
    * @param includeStats whether to read and include the JSON statistics
+   * @param paginationContextOpt pagination context if present
    * @return the surviving scan files as {@link FilteredColumnarBatch}s
    */
-  public CloseableIterator<FilteredColumnarBatch> getScanFiles(
-      Engine engine, boolean includeStats) {
+  protected CloseableIterator<FilteredColumnarBatch> getScanFiles(
+      Engine engine, boolean includeStats, Optional<PaginationContext> paginationContextOpt) {
     if (accessedScanFiles) {
       throw new IllegalStateException("Scan files are already fetched from this instance");
     }
@@ -165,7 +177,8 @@ public CloseableIterator<FilteredColumnarBatch> getScanFiles(
                       predicate ->
                           rewritePartitionPredicateOnCheckpointFileSchema(
                               predicate, partitionColToStructFieldMap.get())),
-              scanMetrics);
+              scanMetrics,
+              paginationContextOpt);
 
       // Apply partition pruning
       scanFileIter = applyPartitionPruning(engine, scanFileIter);

kernel/kernel-api/src/main/java/io/delta/kernel/internal/replay/ActionsIterator.java

@@ -84,22 +84,31 @@ public class ActionsIterator implements CloseableIterator<ActionWrapper> {
    */
   private Optional<CloseableIterator<ActionWrapper>> actionsIter;
 
+  private final Optional<PaginationContext> paginationContextOpt;
+
   private boolean closed;
 
   public ActionsIterator(
       Engine engine,
       List<FileStatus> files,
       StructType deltaReadSchema,
       Optional<Predicate> checkpointPredicate) {
-    this(engine, files, deltaReadSchema, deltaReadSchema, checkpointPredicate);
+    this(
+        engine,
+        files,
+        deltaReadSchema,
+        deltaReadSchema,
+        checkpointPredicate,
+        Optional.empty() /* paginationContextOpt */);
   }
 
   public ActionsIterator(
       Engine engine,
       List<FileStatus> files,
       StructType deltaReadSchema,
       StructType checkpointReadSchema,
-      Optional<Predicate> checkpointPredicate) {
+      Optional<Predicate> checkpointPredicate,
+      Optional<PaginationContext> paginationContextOpt) {
     this.engine = engine;
     this.checkpointPredicate = checkpointPredicate;
     this.filesList = new LinkedList<>();
@@ -109,6 +118,7 @@ public ActionsIterator(
     this.checkpointReadSchema = checkpointReadSchema;
     this.actionsIter = Optional.empty();
     this.schemaContainsAddOrRemoveFiles = LogReplay.containsAddOrRemoveFileActions(deltaReadSchema);
+    this.paginationContextOpt = paginationContextOpt;
   }
 
   @Override

kernel/kernel-api/src/main/java/io/delta/kernel/internal/replay/LogReplay.java

@@ -215,7 +215,8 @@ public CloseableIterator<FilteredColumnarBatch> getAddFilesAsColumnarBatches(
       Engine engine,
       boolean shouldReadStats,
       Optional<Predicate> checkpointPredicate,
-      ScanMetrics scanMetrics) {
+      ScanMetrics scanMetrics,
+      Optional<PaginationContext> paginationContextOpt) {
     // We do not need to look at any `remove` files from the checkpoints. Skip the column to save
     // I/O. Note that we are still going to process the row groups. Adds and removes are randomly
     // scattered through checkpoint part files, so row group push down is unlikely to be useful.
@@ -225,7 +226,8 @@ public CloseableIterator<FilteredColumnarBatch> getAddFilesAsColumnarBatches(
             getLogReplayFiles(getLogSegment()),
             getAddRemoveReadSchema(shouldReadStats),
             getAddReadSchema(shouldReadStats),
-            checkpointPredicate);
+            checkpointPredicate,
+            paginationContextOpt);
     return new ActiveAddFilesIterator(engine, addRemoveIter, dataPath, scanMetrics);
   }
 

kernel/kernel-api/src/main/java/io/delta/kernel/internal/replay/PaginatedScanFilesIteratorImpl.java

@@ -0,0 +1,66 @@
+/*
+ * Copyright (2025) The Delta Lake Project Authors.
+ *
+ * Licensed 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 io.delta.kernel.internal.replay;
+
+import io.delta.kernel.PaginatedScanFilesIterator;
+import io.delta.kernel.data.FilteredColumnarBatch;
+import io.delta.kernel.data.Row;
+import io.delta.kernel.utils.CloseableIterator;
+import java.io.IOException;
+
+/** Implementation of {@link PaginatedScanFilesIterator} */
+public class PaginatedScanFilesIteratorImpl implements PaginatedScanFilesIterator {
+
+  private final CloseableIterator<FilteredColumnarBatch> filteredScanFilesIter;
+  private final long pageSize;
+
+  /**
+   * Constructs a paginated iterator over scan files on top of a given filtered scan files iterator
+   * and pagination context.
+   *
+   * @param filteredScanFilesIter The underlying scan files iterator with data skipping and
+   *     partition pruning applied. This iterator serves as the source of filtered scan results for
+   *     pagination.
+   * @param paginationContext The pagination context that carries pagination-related information,
+   *     such as the maximum number of files to return in a page.
+   */
+  public PaginatedScanFilesIteratorImpl(
+      CloseableIterator<FilteredColumnarBatch> filteredScanFilesIter,
+      PaginationContext paginationContext) {
+    this.filteredScanFilesIter = filteredScanFilesIter;
+    this.pageSize = paginationContext.getPageSize();
+  }
+
+  @Override
+  public Row getCurrentPageToken() {
+    throw new UnsupportedOperationException("Not implemented");
+  }
+
+  @Override
+  public boolean hasNext() {
+    throw new UnsupportedOperationException("Not implemented");
+  }
+
+  @Override
+  public FilteredColumnarBatch next() {
+    throw new UnsupportedOperationException("Not implemented");
+  }
+
+  @Override
+  public void close() throws IOException {
+    throw new UnsupportedOperationException("Not implemented");
+  }
+}