Core: Adds SortRewriteStrategy (#2609)

A rewrite strategy for data files which aims to reorder data with data files to optimally lay them out
in relation to a column. For example, if the Sort strategy is used on a set of files which is ordered
by column x and original has files File A (x: 0 - 50), File B ( x: 10 - 40) and File C ( x: 30 - 60),
this Strategy will attempt to rewrite those files into File A' (x: 0-20), File B' (x: 21 - 40),
File C' (x: 41 - 60).

Author: RussellSpitzer

api/src/main/java/org/apache/iceberg/SortOrder.java

@@ -286,7 +286,7 @@ public SortOrder build() {
     }
   }
 
-  static void checkCompatibility(SortOrder sortOrder, Schema schema) {
+  public static void checkCompatibility(SortOrder sortOrder, Schema schema) {
     for (SortField field : sortOrder.fields) {
       Type sourceType = schema.findType(field.sourceId());
       ValidationException.check(

core/src/main/java/org/apache/iceberg/actions/SortStrategy.java

@@ -0,0 +1,124 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.iceberg.actions;
+
+import java.util.Map;
+import java.util.Set;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.SortOrder;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableSet;
+import org.apache.iceberg.util.PropertyUtil;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A rewrite strategy for data files which aims to reorder data with data files to optimally lay them out
+ * in relation to a column. For example, if the Sort strategy is used on a set of files which is ordered
+ * by column x and original has files File A (x: 0 - 50), File B ( x: 10 - 40) and File C ( x: 30 - 60),
+ * this Strategy will attempt to rewrite those files into File A' (x: 0-20), File B' (x: 21 - 40),
+ * File C' (x: 41 - 60).
+ * <p>
+ * Currently the there is no file overlap detection and we will rewrite all files if {@link SortStrategy#REWRITE_ALL}
+ * is true (default: false). If this property is disabled any files that would be chosen by
+ * {@link BinPackStrategy} will be rewrite candidates.
+ * <p>
+ * In the future other algorithms for determining files to rewrite will be provided.
+ */
+public abstract class SortStrategy extends BinPackStrategy {
+  private static final Logger LOG = LoggerFactory.getLogger(SortStrategy.class);
+
+  /**
+   * Rewrites all files, regardless of their size. Defaults to false, rewriting only mis-sized
+   * files;
+   */
+  public static final String REWRITE_ALL = "rewrite-all";
+  public static final boolean REWRITE_ALL_DEFAULT = false;
+
+
+  private static final Set<String> validOptions = ImmutableSet.of(
+      REWRITE_ALL
+  );
+
+  private boolean rewriteAll;
+  private SortOrder sortOrder;
+
+  /**
+   * Sets the sort order to be used in this strategy when rewriting files
+   * @param order the order to use
+   * @return this for method chaining
+   */
+  public SortStrategy sortOrder(SortOrder order) {
+    this.sortOrder = order;
+    return this;
+  }
+
+  protected SortOrder sortOrder() {
+    return sortOrder;
+  }
+
+  @Override
+  public String name() {
+    return "SORT";
+  }
+
+  @Override
+  public Set<String> validOptions() {
+    return ImmutableSet.<String>builder()
+        .addAll(super.validOptions())
+        .addAll(validOptions)
+        .build();
+  }
+
+  @Override
+  public RewriteStrategy options(Map<String, String> options) {
+    super.options(options); // Also checks validity of BinPack options
+
+    rewriteAll = PropertyUtil.propertyAsBoolean(options,
+        REWRITE_ALL,
+        REWRITE_ALL_DEFAULT);
+
+    if (sortOrder == null) {
+      sortOrder = table().sortOrder();
+    }
+
+    validateOptions();
+    return this;
+  }
+
+  @Override
+  public Iterable<FileScanTask> selectFilesToRewrite(Iterable<FileScanTask> dataFiles) {
+    if (rewriteAll) {
+      LOG.info("Sort Strategy for table {} set to rewrite all data files", table().name());
+      return dataFiles;
+    } else {
+      return super.selectFilesToRewrite(dataFiles);
+    }
+  }
+
+  protected void validateOptions() {
+    Preconditions.checkArgument(!sortOrder.isUnsorted(),
+        "Can't use %s when there is no sort order, either define table %s's sort order or set sort" +
+            "order in the action",
+        name(), table().name());
+
+    SortOrder.checkCompatibility(sortOrder, table().schema());
+  }
+}

core/src/test/java/org/apache/iceberg/MockFileScanTask.java

@@ -19,6 +19,8 @@
 
 package org.apache.iceberg;
 
+import org.mockito.Mockito;
+
 public class MockFileScanTask extends BaseFileScanTask {
 
   private final long length;
@@ -28,6 +30,18 @@ public MockFileScanTask(long length) {
     this.length = length;
   }
 
+  public MockFileScanTask(DataFile file) {
+    super(file, null, null, null, null);
+    this.length = file.fileSizeInBytes();
+  }
+
+  public static MockFileScanTask mockTask(long length, int sortOrderId) {
+    DataFile mockFile = Mockito.mock(DataFile.class);
+    Mockito.when(mockFile.fileSizeInBytes()).thenReturn(length);
+    Mockito.when(mockFile.sortOrderId()).thenReturn(sortOrderId);
+    return new MockFileScanTask(mockFile);
+  }
+
   @Override
   public long length() {
     return length;

core/src/test/java/org/apache/iceberg/actions/TestSortStrategy.java

@@ -0,0 +1,142 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.iceberg.actions;
+
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.IntStream;
+import org.apache.iceberg.AssertHelpers;
+import org.apache.iceberg.DataFile;
+import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.MockFileScanTask;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.SortOrder;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableTestBase;
+import org.apache.iceberg.exceptions.ValidationException;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+import org.apache.iceberg.types.Types;
+import org.junit.Assert;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+@RunWith(Parameterized.class)
+public class TestSortStrategy extends TableTestBase {
+
+  @Parameterized.Parameters(name = "formatVersion = {0}")
+  public static Object[] parameters() {
+    return new Object[] {2}; // We don't actually use the format version since everything is mock
+  }
+
+  @Override
+  public void setupTable() throws Exception {
+    super.setupTable();
+    table.replaceSortOrder().asc("data").commit();
+  }
+
+  private static final long MB = 1024 * 1024;
+
+  public TestSortStrategy(int formatVersion) {
+    super(formatVersion);
+  }
+
+  class TestSortStrategyImpl extends SortStrategy {
+
+    @Override
+    public Table table() {
+      return table;
+    }
+
+    @Override
+    public Set<DataFile> rewriteFiles(List<FileScanTask> filesToRewrite) {
+      throw new UnsupportedOperationException();
+    }
+  }
+
+  private SortStrategy defaultSort() {
+    return (SortStrategy) new TestSortStrategyImpl().options(Collections.emptyMap());
+  }
+
+  private List<FileScanTask> tasksForSortOrder(int sortOrderId, int... fileSizesMB) {
+    ImmutableList.Builder<FileScanTask> files = ImmutableList.builder();
+    IntStream.of(fileSizesMB).forEach(length -> files.add(MockFileScanTask.mockTask(length * MB, sortOrderId)));
+    return files.build();
+  }
+
+  @Test
+  public void testInvalidSortOrder() {
+    AssertHelpers.assertThrows("Should not allow an unsorted Sort order", IllegalArgumentException.class,
+        () -> defaultSort().sortOrder(SortOrder.unsorted()).options(Collections.emptyMap()));
+
+    AssertHelpers.assertThrows("Should not allow a Sort order with bad columns", ValidationException.class,
+        () -> {
+          Schema badSchema = new Schema(
+              ImmutableList.of(Types.NestedField.required(0, "nonexistant", Types.IntegerType.get())));
+
+          defaultSort()
+              .sortOrder(SortOrder.builderFor(badSchema).asc("nonexistant").build())
+              .options(Collections.emptyMap());
+        });
+  }
+
+  @Test
+  public void testSelectAll() {
+    List<FileScanTask> invalid = ImmutableList.<FileScanTask>builder()
+        .addAll(tasksForSortOrder(-1, 500, 500, 500, 500))
+        .addAll(tasksForSortOrder(table.sortOrder().orderId(), 10, 10, 2000, 10))
+        .build();
+
+    List<FileScanTask> expected = ImmutableList.<FileScanTask>builder()
+        .addAll(invalid)
+        .addAll(tasksForSortOrder(table.sortOrder().orderId(), 500, 490, 520))
+        .build();
+
+    RewriteStrategy strategy = defaultSort().options(ImmutableMap.of(SortStrategy.REWRITE_ALL, "true"));
+    List<FileScanTask> actual = ImmutableList.copyOf(strategy.selectFilesToRewrite(expected));
+
+    Assert.assertEquals("Should mark all files for rewrite",
+        expected, actual);
+  }
+
+  @Test
+  public void testUseSizeOptions() {
+    List<FileScanTask> expected = ImmutableList.<FileScanTask>builder()
+        .addAll(tasksForSortOrder(table.sortOrder().orderId(), 498, 551))
+        .build();
+
+    List<FileScanTask> fileScanTasks = ImmutableList.<FileScanTask>builder()
+        .addAll(expected)
+        .addAll(tasksForSortOrder(table.sortOrder().orderId(), 500, 500))
+        .build();
+
+    RewriteStrategy strategy = defaultSort().options(ImmutableMap.of(
+        SortStrategy.MAX_FILE_SIZE_BYTES, Long.toString(550 * MB),
+        SortStrategy.MIN_FILE_SIZE_BYTES, Long.toString(499 * MB)));
+
+    List<FileScanTask> actual = ImmutableList.copyOf(strategy.selectFilesToRewrite(fileScanTasks));
+
+    Assert.assertEquals("Should mark files for rewrite with adjusted min and max size",
+        expected, actual);
+  }
+}