lasdf1234 commented on code in PR #12194:
URL: https://github.com/apache/gravitino/pull/12194#discussion_r3663625367


##########
iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/CatalogWrapperForREST.java:
##########
@@ -484,6 +479,232 @@ public PlanTableScanResponse planTableScan(
     }
   }
 
+  /**
+   * Fetch the scan tasks covered by a {@code plan-task} token previously 
handed out by {@link
+   * #planTableScan}, completing the second step of the Iceberg REST scan 
planning protocol.
+   *
+   * <p>Tokens are self-describing (see {@link PlanTaskToken}): each one 
carries the scan request it
+   * was planned from, with the snapshot pinned at planning time, plus the 
range of file scan tasks
+   * it stands for. The plan is reproduced here from the {@linkplain 
ScanPlanCache scan plan cache}
+   * when it is still cached and re-planned against the pinned snapshot 
otherwise, then the token's
+   * range is returned. Because no state is kept between the two calls, a 
token remains redeemable
+   * after a server restart and on any Gravitino instance serving the same 
catalog.
+   *
+   * @param tableIdentifier the table the plan task belongs to.
+   * @param request the request carrying the {@code plan-task} token.
+   * @return the file scan tasks the token covers.
+   * @throws org.apache.iceberg.exceptions.NoSuchTableException if the table 
doesn't exist.
+   * @throws NoSuchPlanTaskException if the token was not issued for this 
table, or the plan it
+   *     refers to can no longer be reproduced (for example its snapshot has 
expired).
+   */
+  @SuppressWarnings("deprecation")
+  public FetchScanTasksResponse fetchScanTasks(
+      TableIdentifier tableIdentifier, FetchScanTasksRequest request) {
+    Optional<PlanTaskToken> decodedToken = 
PlanTaskToken.decode(request.planTask());
+
+    // Validate the table exists first, so a bad table reports 404 for the 
table rather than
+    // masking it as an unknown plan task. Consistent with planTableScan 
behavior.
+    Table table = getCatalog().loadTable(tableIdentifier);
+
+    if (!decodedToken.isPresent() || 
!decodedToken.get().matchesTable(tableIdentifier)) {
+      LOG.info(
+          "Rejecting unknown plan task '{}' for table {}", request.planTask(), 
tableIdentifier);
+      throw new NoSuchPlanTaskException(
+          "Plan task %s was not issued for table %s", request.planTask(), 
tableIdentifier);
+    }
+
+    PlanTaskToken token = decodedToken.get();
+    PlanTableScanResponse fullPlan;
+    try {
+      fullPlan = planFullScan(tableIdentifier, table, token.scanRequest());
+    } catch (IllegalArgumentException e) {
+      // The pinned snapshot is gone (expired or rolled back), so the plan the 
token refers to can
+      // no longer be reproduced. That is a stale plan task, not a bad request.
+      LOG.info(
+          "Plan task '{}' for table {} can no longer be planned: {}",
+          request.planTask(),
+          tableIdentifier,
+          e.getMessage());
+      throw new NoSuchPlanTaskException(
+          "Plan task %s is no longer available for table %s: %s",
+          request.planTask(), tableIdentifier, e.getMessage());
+    }
+
+    List<FileScanTask> allTasks = fullPlan.fileScanTasks();
+    int taskCount = allTasks == null ? 0 : allTasks.size();
+    if (token.offset() >= taskCount) {
+      LOG.info(
+          "Plan task '{}' for table {} covers tasks from offset {}, but the 
plan has {} tasks",
+          request.planTask(),
+          tableIdentifier,
+          token.offset(),
+          taskCount);
+      throw new NoSuchPlanTaskException(
+          "Plan task %s is no longer available for table %s", 
request.planTask(), tableIdentifier);
+    }
+
+    List<FileScanTask> batch =
+        ImmutableList.copyOf(
+            allTasks.subList(token.offset(), Math.min(token.offset() + 
token.limit(), taskCount)));
+    LOG.info(
+        "Returning {} file scan tasks for plan task of table {} at offset {}",
+        batch.size(),
+        tableIdentifier,
+        token.offset());
+
+    FetchScanTasksResponse.Builder builder =
+        FetchScanTasksResponse.builder()
+            .withFileScanTasks(batch)
+            .withSpecsById(fullPlan.specsById());
+    List<DeleteFile> deleteFiles = referencedDeleteFiles(batch);
+    if (!deleteFiles.isEmpty()) {
+      builder.withDeleteFiles(deleteFiles);
+    }
+    return builder.build();
+  }
+
+  /**
+   * Plans the whole scan and returns every file scan task inline, serving the 
{@linkplain
+   * ScanPlanCache scan plan cache} when the same plan was computed before.
+   *
+   * <p>Tasks are ordered deterministically so that a plan-task token, which 
addresses tasks by
+   * position, resolves to the same tasks on a later re-plan of the same 
snapshot.
+   */
+  private PlanTableScanResponse planFullScan(
+      TableIdentifier tableIdentifier, Table table, PlanTableScanRequest 
scanRequest) {
+    ScanPlanCacheKey cacheKey = ScanPlanCacheKey.create(tableIdentifier, 
table, scanRequest);
+    Optional<PlanTableScanResponse> cachedResponse = 
scanPlanCache.get(cacheKey);
+    if (cachedResponse.isPresent()) {
+      LOG.info("Using cached scan plan for table: {}", tableIdentifier);
+      return cachedResponse.get();
+    }
+
+    List<FileScanTask> fileScanTasks = new ArrayList<>();
+    try (CloseableIterable<FileScanTask> scanTasks =
+        createFilePlanScanTasks(table, tableIdentifier, scanRequest)) {
+      for (FileScanTask fileScanTask : scanTasks) {
+        fileScanTasks.add(fileScanTask);
+      }
+    } catch (IOException e) {
+      LOG.error("Failed to close scan task iterator for table: {}", 
tableIdentifier, e);
+      throw new RuntimeException("Failed to plan scan tasks: " + 
e.getMessage(), e);
+    }
+
+    if (fileScanTasks.isEmpty()) {
+      LOG.info(
+          "Scan planning returned no tasks for table: {}. Table may be empty 
or fully filtered.",
+          tableIdentifier);
+    }
+
+    // Iceberg plans manifests in parallel, so the order tasks come back in is 
not reproducible.
+    // Sort them so positions stay stable across re-plans of the same snapshot.
+    fileScanTasks.sort(FILE_SCAN_TASK_ORDER);
+
+    PlanTableScanResponse response;
+    try {
+      response = buildCompletedPlanTableScanResponse(table, fileScanTasks);
+    } catch (Exception e) {
+      LOG.error("Failed to build scan plan response for table: {}", 
tableIdentifier, e);
+      throw new RuntimeException(
+          String.format(
+              "Failed to build scan plan response for table: %s. Error: %s",
+              tableIdentifier, e.getMessage()),
+          e);
+    }
+
+    scanPlanCache.put(cacheKey, response);
+    return response;
+  }
+
+  /**
+   * Keeps the first {@link IcebergConfig#SCAN_PLAN_TASK_BATCH_SIZE} file scan 
tasks of {@code
+   * fullPlan} inline and turns the remaining tasks into {@code plan-tasks} 
tokens, so a single
+   * response never carries an unbounded plan.
+   *
+   * <p>Returns {@code fullPlan} unchanged when batching is disabled or the 
plan already fits in one
+   * batch, which is the common case and keeps a plan a client can consume 
without a second call.
+   */
+  @SuppressWarnings("deprecation")
+  private PlanTableScanResponse splitIntoPlanTasks(
+      TableIdentifier tableIdentifier,
+      PlanTableScanRequest scanRequest,
+      PlanTableScanResponse fullPlan) {
+    List<FileScanTask> allTasks = fullPlan.fileScanTasks();
+    if (scanPlanTaskBatchSize <= 0
+        || allTasks == null
+        || allTasks.size() <= scanPlanTaskBatchSize) {
+      return fullPlan;
+    }
+
+    List<String> planTasks = new ArrayList<>();
+    for (int offset = scanPlanTaskBatchSize;
+        offset < allTasks.size();
+        offset += scanPlanTaskBatchSize) {
+      planTasks.add(
+          PlanTaskToken.encode(tableIdentifier, scanRequest, offset, 
scanPlanTaskBatchSize));
+    }
+
+    List<FileScanTask> firstBatch =
+        ImmutableList.copyOf(allTasks.subList(0, scanPlanTaskBatchSize));
+    LOG.info(
+        "Split scan plan of table {} into {} inline file scan tasks and {} 
plan tasks",
+        tableIdentifier,
+        firstBatch.size(),
+        planTasks.size());
+
+    PlanTableScanResponse.Builder builder =
+        PlanTableScanResponse.builder()
+            .withPlanStatus(PlanStatus.COMPLETED)
+            .withFileScanTasks(firstBatch)
+            .withPlanTasks(planTasks)
+            .withSpecsById(fullPlan.specsById());
+    List<DeleteFile> deleteFiles = referencedDeleteFiles(firstBatch);
+    if (!deleteFiles.isEmpty()) {
+      builder.withDeleteFiles(deleteFiles);
+    }
+    return builder.build();
+  }
+
+  /**
+   * Pins the snapshot a scan request plans against, so that plan-task tokens 
issued for the
+   * resulting plan keep resolving to the same snapshot after the table 
changes.
+   *
+   * <p>Requests that already name a snapshot, incremental requests (which pin 
a snapshot range) and
+   * scans of a table without a current snapshot are returned unchanged.
+   */
+  private static PlanTableScanRequest pinSnapshot(Table table, 
PlanTableScanRequest scanRequest) {
+    boolean isIncremental =

Review Comment:
   This method can be omitted (as it was only used once), and the logic can be 
incorporated into the original method.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to