Spark-3.4: Harmonize RewriteDataFilesSparkAction

[ajantha-bhat]

some code style fixups to RewriteDeleteFileSparkAction that can be applied as well to RewriteDataFileSparkAction #7565 (review)

Fixes #7625

[aokolnychyi]

Can we also add EMPTY_RESULT static variable like in the action for position deletes? There are two return statements in this method (not this line).

[ajantha-bhat]

ok. Added.

[aokolnychyi]

This should continue to use table.spec() and we don't need coercePartition. The old behavior of grouping all files from an old spec should be preserved. The two actions behave differently w.r.t. output partitioning.

[ajantha-bhat]

Thats strange. So, we don't have enough testcases to catch this behaviour?

I have reverted it back now. Also, I didn't use transform in the end of planFileGroups because it needs a special check of fileGroups.size() > 0 which is not handled in transform due to generics

[szehon-ho]

lgtm

[szehon-ho]

I guess i'd personally put the variable back, as in the RewriteDeleteFilesSparkAction case it was just a call without arguments, but here there are arguments so probably is more clearer if we assign it separately, but not a big deal

[ajantha-bhat]

I too thought about it during the change.
Since it is in if-else, I don't think it will make a huge difference.

I will leave it as it is (because the style looks similar to RewriteDeleteFilesSparkAction)

[dramaticlly]

LGTM

[szehon-ho]

Actually I took another pass, and realize we can still do the refactoring while preserving the logic as @aokolnychyi suggested

[szehon-ho]

Can't comment on the exact line , but on line 196-304 below , initially you had done some refactoring there. I see it was reverted as @aokolnychyi said we can't coerce partition to change the behavior in this discussion: #7630 (comment)

That is correct, but we can still do some refactoring without coerce. I took a brief look and came up with this to maintain the logic:

First add the methods:

private StructLikeMap<List<FileScanTask>> groupByPartition(
      StructType partitionType, Iterable<FileScanTask> tasks) {
    StructLikeMap<List<FileScanTask>> filesByPartition = StructLikeMap.create(partitionType);
    StructLike emptyStruct = GenericRecord.create(partitionType);

    for (FileScanTask task : tasks) {
      // If a task uses an incompatible partition spec the data inside could contain values
      // which belong to multiple partitions in the current spec. Treating all such files as
      // un-partitioned and grouping them together helps to minimize new files made.
      StructLike taskPartition =
          task.file().specId() == table.spec().specId()
              ? task.file().partition()
              : emptyStruct;

      List<FileScanTask> files = filesByPartition.get(taskPartition);
      if (files == null) {
        files = Lists.newArrayList();
      }

      files.add(task);
      filesByPartition.put(taskPartition, files);
    }
    return filesByPartition;
  }

  private StructLikeMap<List<List<FileScanTask>>> fileGroupsByPartition(StructLikeMap<List<FileScanTask>> filesByPartition) {
    return filesByPartition.transformValues(this::planFileGroups);
  }

  private List<List<FileScanTask>> planFileGroups(List<FileScanTask> tasks) {
    return ImmutableList.copyOf(rewriter.planFileGroups(tasks));
  }

and then we can make this method like:

  Map<StructLike, List<List<FileScanTask>>> planFileGroups(long startingSnapshotId) {
    CloseableIterable<FileScanTask> fileScanTasks =
        table
            .newScan()
            .useSnapshot(startingSnapshotId)
            .filter(filter)
            .ignoreResiduals()
            .planFiles();

    try {
      StructType partitionType = table.spec().partitionType();
      StructLikeMap<List<FileScanTask>> filesByPartition = groupByPartition(partitionType, fileScanTasks);
      return fileGroupsByPartition(filesByPartition);
    } finally {
      try {
        fileScanTasks.close();
      } catch (IOException io) {
        LOG.error("Cannot properly close file iterable while planning for rewrite", io);
      }
    }
  }

[ajantha-bhat]

I think I tried it before. Have you checked this comment?
#7630 (comment)

If I use transformValues, it cannot check the existing fileGroups.size() > 0 check due to generics. Are we ok with that? That's why I didn't do

[szehon-ho]

I think I tried it before. Have you checked this comment?
#7630 (comment)

Yea I saw that, so hence my suggestion did not use coerce, but rather the table's latest partitionType. It should be exactly the same as the existing code now?

If I use transformValues, it cannot check the existing fileGroups.size() > 0 check due to generics. Are we ok with that? That's why I didn't do

Yea , in RewritePositionDeleteFilesSparkAction, we've moved that check as part of toGroupStream, which I think you also copied here.

[ajantha-bhat]

we've moved that check as part of toGroupStream

I can see this now. SGTM.
I have addressed the suggestion.

[szehon-ho]

Merged, thanks @ajantha-bhat

[ajantha-bhat]

@szehon-ho: Thanks for the review. I will backport to other spark versions later today.

Also, I can see that new changes like RewritePositionDeleteFilesSparkAction.java is not backported. Do we need to handle this before 1.3 release?