[iceberg] branch master updated: Add TableMetadata.updateSortOrder. (#1551)

[dw...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dweeks pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iceberg.git


The following commit(s) were added to refs/heads/master by this push:
     new 97568ce  Add TableMetadata.updateSortOrder. (#1551)
97568ce is described below

commit 97568cebfa03d39079308a1c495967c2240829d3
Author: Ryan Blue <bl...@apache.org>
AuthorDate: Tue Oct 6 16:57:30 2020 -0700

    Add TableMetadata.updateSortOrder. (#1551)
---
 .../java/org/apache/iceberg/TableMetadata.java     | 37 +++++++++++++++
 .../java/org/apache/iceberg/TestTableMetadata.java | 53 ++++++++++++++++++++++
 2 files changed, 90 insertions(+)

diff --git a/core/src/main/java/org/apache/iceberg/TableMetadata.java b/core/src/main/java/org/apache/iceberg/TableMetadata.java
index 566f953..d00e2a8 100644
--- a/core/src/main/java/org/apache/iceberg/TableMetadata.java
+++ b/core/src/main/java/org/apache/iceberg/TableMetadata.java
@@ -491,6 +491,43 @@ public class TableMetadata implements Serializable {
         currentSnapshotId, snapshots, snapshotLog, addPreviousFile(file, lastUpdatedMillis));
   }
 
+  public TableMetadata updateSortOrder(SortOrder newOrder) {
+    SortOrder.checkCompatibility(newOrder, schema);
+
+    // determine the next order id
+    int newOrderId = INITIAL_SORT_ORDER_ID;
+    for (SortOrder order : sortOrders) {
+      if (order.sameOrder(newOrder)) {
+        newOrderId = order.orderId();
+        break;
+      } else if (newOrderId <= order.orderId()) {
+        newOrderId = order.orderId() + 1;
+      }
+    }
+
+    if (newOrderId == defaultSortOrderId) {
+      return this;
+    }
+
+    ImmutableList.Builder<SortOrder> builder = ImmutableList.builder();
+    builder.addAll(sortOrders);
+
+    if (!sortOrdersById.containsKey(newOrderId)) {
+      if (newOrder.isUnsorted()) {
+        newOrderId = SortOrder.unsorted().orderId();
+        builder.add(SortOrder.unsorted());
+      } else {
+        // rebuild the sort order using new column ids
+        builder.add(freshSortOrder(newOrderId, schema, newOrder));
+      }
+    }
+
+    return new TableMetadata(null, formatVersion, uuid, location,
+        lastSequenceNumber, System.currentTimeMillis(), lastColumnId, schema, defaultSpecId, specs,
+        newOrderId, builder.build(), properties, currentSnapshotId, snapshots, snapshotLog,
+        addPreviousFile(file, lastUpdatedMillis));
+  }
+
   public TableMetadata addStagedSnapshot(Snapshot snapshot) {
     ValidationException.check(formatVersion == 1 || snapshot.sequenceNumber() > lastSequenceNumber,
         "Cannot add snapshot with sequence number %s older than last sequence number %s",
diff --git a/core/src/test/java/org/apache/iceberg/TestTableMetadata.java b/core/src/test/java/org/apache/iceberg/TestTableMetadata.java
index 38512de..6d188fa 100644
--- a/core/src/test/java/org/apache/iceberg/TestTableMetadata.java
+++ b/core/src/test/java/org/apache/iceberg/TestTableMetadata.java
@@ -550,4 +550,57 @@ public class TestTableMetadata {
         ValidationException.class, "Spec does not use sequential IDs that are required in v1",
         () -> metadata.updatePartitionSpec(spec));
   }
+
+  @Test
+  public void testSortOrder() {
+    Schema schema = new Schema(
+        Types.NestedField.required(10, "x", Types.StringType.get())
+    );
+
+    TableMetadata meta = TableMetadata.newTableMetadata(
+        schema, PartitionSpec.unpartitioned(), null, ImmutableMap.of());
+    Assert.assertTrue("Should default to unsorted order", meta.sortOrder().isUnsorted());
+    Assert.assertSame("Should detect identical unsorted order", meta, meta.updateSortOrder(SortOrder.unsorted()));
+  }
+
+  @Test
+  public void testUpdateSortOrder() {
+    Schema schema = new Schema(
+        Types.NestedField.required(10, "x", Types.StringType.get())
+    );
+
+    SortOrder order = SortOrder.builderFor(schema).asc("x").build();
+
+    TableMetadata sortedByX = TableMetadata.newTableMetadata(
+        schema, PartitionSpec.unpartitioned(), order, null, ImmutableMap.of());
+    Assert.assertEquals("Should have 1 sort order", 1, sortedByX.sortOrders().size());
+    Assert.assertEquals("Should use orderId 1", 1, sortedByX.sortOrder().orderId());
+    Assert.assertEquals("Should be sorted by one field", 1, sortedByX.sortOrder().fields().size());
+    Assert.assertEquals("Should use the table's field ids", 1, sortedByX.sortOrder().fields().get(0).sourceId());
+    Assert.assertEquals("Should be ascending",
+        SortDirection.ASC, sortedByX.sortOrder().fields().get(0).direction());
+    Assert.assertEquals("Should be nulls first",
+        NullOrder.NULLS_FIRST, sortedByX.sortOrder().fields().get(0).nullOrder());
+
+    // build an equivalent order with the correct schema
+    SortOrder newOrder = SortOrder.builderFor(sortedByX.schema()).asc("x").build();
+
+    TableMetadata alsoSortedByX = sortedByX.updateSortOrder(newOrder);
+    Assert.assertSame("Should detect current sortOrder and not update", alsoSortedByX, sortedByX);
+
+    TableMetadata unsorted = alsoSortedByX.updateSortOrder(SortOrder.unsorted());
+    Assert.assertEquals("Should have 2 sort orders", 2, unsorted.sortOrders().size());
+    Assert.assertEquals("Should use orderId 0", 0, unsorted.sortOrder().orderId());
+    Assert.assertTrue("Should be unsorted", unsorted.sortOrder().isUnsorted());
+
+    TableMetadata sortedByXDesc = unsorted.updateSortOrder(SortOrder.builderFor(unsorted.schema()).desc("x").build());
+    Assert.assertEquals("Should have 3 sort orders", 3, sortedByXDesc.sortOrders().size());
+    Assert.assertEquals("Should use orderId 2", 2, sortedByXDesc.sortOrder().orderId());
+    Assert.assertEquals("Should be sorted by one field", 1, sortedByXDesc.sortOrder().fields().size());
+    Assert.assertEquals("Should use the table's field ids", 1, sortedByXDesc.sortOrder().fields().get(0).sourceId());
+    Assert.assertEquals("Should be ascending",
+        SortDirection.DESC, sortedByXDesc.sortOrder().fields().get(0).direction());
+    Assert.assertEquals("Should be nulls first",
+        NullOrder.NULLS_FIRST, sortedByX.sortOrder().fields().get(0).nullOrder());
+  }
 }