[GitHub] [iceberg] aokolnychyi commented on a diff in pull request #2276: Core: Add a util method to combine tasks by partition

[GitBox <gi...@apache.org>]

aokolnychyi commented on code in PR #2276:
URL: https://github.com/apache/iceberg/pull/2276#discussion_r1019834880


##########
core/src/main/java/org/apache/iceberg/util/TableScanUtil.java:
##########
@@ -128,6 +136,61 @@ public static <T extends ScanTask> CloseableIterable<ScanTaskGroup<T>> planTaskG
         combinedTasks -> new BaseScanTaskGroup<>(mergeTasks(combinedTasks)));
   }
 
+  @SuppressWarnings("unchecked")
+  public static <T extends PartitionScanTask> List<ScanTaskGroup<T>> planTaskGroups(
+      List<T> tasks,
+      long splitSize,
+      int lookback,
+      long openFileCost,
+      Types.StructType projectedPartitionType) {
+
+    Preconditions.checkArgument(splitSize > 0, "Invalid split size (negative or 0): %s", splitSize);
+    Preconditions.checkArgument(
+        lookback > 0, "Invalid split planning lookback (negative or 0): %s", lookback);
+    Preconditions.checkArgument(
+        openFileCost >= 0, "Invalid file open cost (negative): %s", openFileCost);
+
+    Function<T, Long> weightFunc =
+        task -> Math.max(task.sizeBytes(), task.filesCount() * openFileCost);
+
+    Map<Integer, StructProjection> projectionsBySpec = Maps.newHashMap();
+
+    // Group tasks by their partition values
+    StructLikeMap<List<T>> tasksByPartition = StructLikeMap.create(projectedPartitionType);
+
+    for (T task : tasks) {
+      PartitionSpec spec = task.spec();
+      StructProjection projectedStruct =
+          projectionsBySpec.computeIfAbsent(
+              spec.specId(),
+              specId -> StructProjection.create(spec.partitionType(), projectedPartitionType));
+      List<T> taskList =
+          tasksByPartition.computeIfAbsent(
+              projectedStruct.copyFor(task.partition()), k -> Lists.newArrayList());
+      if (task instanceof SplittableScanTask<?>) {
+        ((SplittableScanTask<? extends T>) task).split(splitSize).forEach(taskList::add);
+      } else {
+        taskList.add(task);
+      }
+    }
+
+    // Now apply task combining within each partition
+    return tasksByPartition.values().stream()
+        .flatMap(ts -> toTaskGroupStream(ts, splitSize, lookback, weightFunc))
+        .collect(Collectors.toList());
+  }
+
+  private static <T extends ScanTask> Stream<ScanTaskGroup<T>> toTaskGroupStream(
+      Iterable<T> tasks, long splitSize, int lookback, Function<T, Long> weightFunc) {
+    CloseableIterable<ScanTaskGroup<T>> taskGroups =

Review Comment:
   After a closer look, do we need `CloseableIterable`? I don't think there is anything to close. What about using plain iterables instead of the stream then?
   
   ```
   private static <T extends ScanTask> Iterable<ScanTaskGroup<T>> toTaskGroupIterable(
       Iterable<T> tasks, long splitSize, int lookback, Function<T, Long> weightFunc) {
   
     return Iterables.transform(
         new BinPacking.PackingIterable<>(tasks, splitSize, lookback, weightFunc, true),
         combinedTasks -> new BaseScanTaskGroup<>(mergeTasks(combinedTasks)));
   }
   ```
   
   We could call it like this:
   
   ```
   return FluentIterable.from(tasksByPartition.values())
       .transformAndConcat(ts -> toTaskGroupIterable(ts, splitSize, lookback, weightFunc))
       .toList();
   ```
   
   I feel `transformAndConcat` in `FluentIterable` would be best.
   
   



##########
api/src/main/java/org/apache/iceberg/util/StructProjection.java:
##########
@@ -171,6 +178,11 @@ public StructProjection wrap(StructLike newStruct) {
     return this;
   }
 
+  public StructProjection copyFor(StructLike newStruct) {
+    return new StructProjection(this.type, this.positionMap, this.nestedProjections)

Review Comment:
   nit: We usually use explicit `this.` only when setting values. I think it would fit on one line if we drop `this.` from these three variables.



-- 
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: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org