Core: fix empty rest table scan

Author: geruh

core/src/main/java/org/apache/iceberg/rest/CatalogHandlers.java

@@ -26,6 +26,7 @@
 import java.time.OffsetDateTime;
 import java.time.ZoneOffset;
 import java.util.Collections;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -672,20 +673,30 @@ public static PlanTableScanResponse planTableScan(
     }
 
     String planId = "sync-" + UUID.randomUUID();
-    planFilesFor(tableScan, planId, tasksPerPlanTask.applyAsInt(tableScan));
-    Pair<List<FileScanTask>, String> initial = IN_MEMORY_PLANNING_STATE.initialScanTasksFor(planId);
-    return PlanTableScanResponse.builder()
-        .withPlanStatus(PlanStatus.COMPLETED)
-        .withPlanId(planId)
-        .withPlanTasks(IN_MEMORY_PLANNING_STATE.nextPlanTask(initial.second()))
-        .withFileScanTasks(initial.first())
-        .withDeleteFiles(
-            initial.first().stream()
-                .flatMap(task -> task.deletes().stream())
-                .distinct()
-                .collect(Collectors.toList()))
-        .withSpecsById(table.specs())
-        .build();
+    Pair<List<FileScanTask>, String> initial =
+        planFilesFor(tableScan, planId, tasksPerPlanTask.applyAsInt(tableScan));
+    List<String> nextPlanTasks =
+        initial.second() == null
+            ? Collections.emptyList()
+            : IN_MEMORY_PLANNING_STATE.nextPlanTask(initial.second());
+
+    PlanTableScanResponse.Builder builder =
+        PlanTableScanResponse.builder()
+            .withPlanStatus(PlanStatus.COMPLETED)
+            .withPlanId(planId)
+            .withFileScanTasks(initial.first())
+            .withDeleteFiles(
+                initial.first().stream()
+                    .flatMap(task -> task.deletes().stream())
+                    .distinct()
+                    .collect(Collectors.toList()))
+            .withSpecsById(table.specs());
+
+    if (!nextPlanTasks.isEmpty()) {
+      builder.withPlanTasks(nextPlanTasks);
+    }
+
+    return builder.build();
   }
 
   /**
@@ -763,28 +774,45 @@ static void clearPlanningState() {
    * @param tableScan the table scan to plan
    * @param planId the unique identifier for this plan
    * @param tasksPerPlanTask number of file scan tasks to group per plan task
+   * @return the initial file scan tasks and the first plan task key
    */
-  private static void planFilesFor(TableScan tableScan, String planId, int tasksPerPlanTask) {
+  private static Pair<List<FileScanTask>, String> planFilesFor(
+      TableScan tableScan, String planId, int tasksPerPlanTask) {
     Iterable<List<FileScanTask>> taskGroupings =
         Iterables.partition(tableScan.planFiles(), tasksPerPlanTask);
+    Iterator<List<FileScanTask>> taskIterator = taskGroupings.iterator();
+    String planTaskKeyPrefix = planId + "-" + tableScan.table().uuid() + "-";
+
+    if (!taskIterator.hasNext()) {
+      // Handle empty table scans
+      String emptyPlanTaskKey = planTaskKeyPrefix + 0;
+      IN_MEMORY_PLANNING_STATE.addPlanTask(emptyPlanTaskKey, Collections.emptyList());
+      return Pair.of(Collections.emptyList(), null);
+    }
+
     int planTaskSequence = 0;
     String previousPlanTask = null;
-    String planTaskKeyPrefix = planId + "-" + tableScan.table().uuid() + "-";
-    for (List<FileScanTask> taskGrouping : taskGroupings) {
+    String firstPlanTaskKey = null;
+    List<FileScanTask> initialFileScanTasks = null;
+
+    do {
+      List<FileScanTask> taskGrouping = taskIterator.next();
       String planTaskKey = planTaskKeyPrefix + planTaskSequence++;
       IN_MEMORY_PLANNING_STATE.addPlanTask(planTaskKey, taskGrouping);
+
+      if (firstPlanTaskKey == null) {
+        firstPlanTaskKey = planTaskKey;
+        initialFileScanTasks = taskGrouping;
+      }
+
       if (previousPlanTask != null) {
         IN_MEMORY_PLANNING_STATE.addNextPlanTask(previousPlanTask, planTaskKey);
       }
 
       previousPlanTask = planTaskKey;
-    }
+    } while (taskIterator.hasNext());
 
-    // If a scan produces no file tasks, create a single empty plan task
-    if (planTaskSequence == 0) {
-      String planTaskKey = planTaskKeyPrefix + planTaskSequence;
-      IN_MEMORY_PLANNING_STATE.addPlanTask(planTaskKey, Collections.emptyList());
-    }
+    return Pair.of(initialFileScanTasks, firstPlanTaskKey);
   }
 
   @SuppressWarnings("FutureReturnValueIgnored")