Spark: Refactor building write distribution and ordering (apache#3720)

hf200012 · Dec 13, 2021 · 53557ac · 53557ac
1 parent 4c48df2
commit 53557ac
Show file tree

Hide file tree

Showing 6 changed files with 365 additions and 140 deletions.
diff --git a/.baseline/checkstyle/checkstyle.xml b/.baseline/checkstyle/checkstyle.xml
@@ -94,6 +94,7 @@
                 org.apache.commons.lang3.Validate.*,
                 org.apache.iceberg.expressions.Expressions.*,
                 org.apache.iceberg.expressions.Expression.Operation.*,
+                org.apache.iceberg.DistributionMode.*,
                 org.apache.iceberg.IsolationLevel.*,
                 org.apache.iceberg.NullOrder.*,
                 org.apache.iceberg.MetadataTableType.*,

diff --git a/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/OrderField.java b/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/OrderField.java
diff --git a/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/SortOrderToSpark.java b/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/SortOrderToSpark.java
@@ -24,8 +24,11 @@
 import org.apache.iceberg.Schema;
 import org.apache.iceberg.SortDirection;
 import org.apache.iceberg.transforms.SortOrderVisitor;
+import org.apache.spark.sql.connector.expressions.Expressions;
+import org.apache.spark.sql.connector.expressions.NullOrdering;
+import org.apache.spark.sql.connector.expressions.SortOrder;
 
-class SortOrderToSpark implements SortOrderVisitor<OrderField> {
+class SortOrderToSpark implements SortOrderVisitor<SortOrder> {
 
   private final Map<Integer, String> quotedNameById;
 
@@ -34,42 +37,56 @@ class SortOrderToSpark implements SortOrderVisitor<OrderField> {
   }
 
   @Override
-  public OrderField field(String sourceName, int id, SortDirection direction, NullOrder nullOrder) {
-    return OrderField.column(quotedName(id), direction, nullOrder);
+  public SortOrder field(String sourceName, int id, SortDirection direction, NullOrder nullOrder) {
+    return Expressions.sort(Expressions.column(quotedName(id)), toSpark(direction), toSpark(nullOrder));
   }
 
   @Override
-  public OrderField bucket(String sourceName, int id, int width, SortDirection direction, NullOrder nullOrder) {
-    return OrderField.bucket(quotedName(id), width, direction, nullOrder);
+  public SortOrder bucket(String sourceName, int id, int width, SortDirection direction, NullOrder nullOrder) {
+    return Expressions.sort(Expressions.bucket(width, quotedName(id)), toSpark(direction), toSpark(nullOrder));
   }
 
   @Override
-  public OrderField truncate(String sourceName, int id, int width, SortDirection direction, NullOrder nullOrder) {
-    return OrderField.truncate(quotedName(id), width, direction, nullOrder);
+  public SortOrder truncate(String sourceName, int id, int width, SortDirection direction, NullOrder nullOrder) {
+    return Expressions.sort(Expressions.apply(
+        "truncate", Expressions.column(quotedName(id)), Expressions.literal(width)),
+        toSpark(direction), toSpark(nullOrder));
   }
 
   @Override
-  public OrderField year(String sourceName, int id, SortDirection direction, NullOrder nullOrder) {
-    return OrderField.year(quotedName(id), direction, nullOrder);
+  public SortOrder year(String sourceName, int id, SortDirection direction, NullOrder nullOrder) {
+    return Expressions.sort(Expressions.years(quotedName(id)), toSpark(direction), toSpark(nullOrder));
   }
 
   @Override
-  public OrderField month(String sourceName, int id, SortDirection direction, NullOrder nullOrder) {
-    return OrderField.month(quotedName(id), direction, nullOrder);
+  public SortOrder month(String sourceName, int id, SortDirection direction, NullOrder nullOrder) {
+    return Expressions.sort(Expressions.months(quotedName(id)), toSpark(direction), toSpark(nullOrder));
   }
 
   @Override
-  public OrderField day(String sourceName, int id, SortDirection direction, NullOrder nullOrder) {
-    return OrderField.day(quotedName(id), direction, nullOrder);
+  public SortOrder day(String sourceName, int id, SortDirection direction, NullOrder nullOrder) {
+    return Expressions.sort(Expressions.days(quotedName(id)), toSpark(direction), toSpark(nullOrder));
   }
 
   @Override
-  public OrderField hour(String sourceName, int id, SortDirection direction, NullOrder nullOrder) {
-    return OrderField.hour(quotedName(id), direction, nullOrder);
+  public SortOrder hour(String sourceName, int id, SortDirection direction, NullOrder nullOrder) {
+    return Expressions.sort(Expressions.hours(quotedName(id)), toSpark(direction), toSpark(nullOrder));
   }
 
   private String quotedName(int id) {
     return quotedNameById.get(id);
   }
+
+  private org.apache.spark.sql.connector.expressions.SortDirection toSpark(SortDirection direction) {
+    if (direction == SortDirection.ASC) {
+      return org.apache.spark.sql.connector.expressions.SortDirection.ASCENDING;
+    } else {
+      return org.apache.spark.sql.connector.expressions.SortDirection.DESCENDING;
+    }
+  }
+
+  private NullOrdering toSpark(NullOrder nullOrder) {
+    return nullOrder == NullOrder.NULLS_FIRST ? NullOrdering.NULLS_FIRST : NullOrdering.NULLS_LAST;
+  }
 }
 
diff --git a/...k/v3.2/spark/src/main/java/org/apache/iceberg/spark/SparkDistributionAndOrderingUtil.java b/...k/v3.2/spark/src/main/java/org/apache/iceberg/spark/SparkDistributionAndOrderingUtil.java
@@ -40,19 +40,11 @@ public static Distribution buildRequiredDistribution(Table table, DistributionMo
         return Distributions.unspecified();
 
       case HASH:
-        if (table.spec().isUnpartitioned()) {
-          return Distributions.unspecified();
-        } else {
-          return Distributions.clustered(Spark3Util.toTransforms(table.spec()));
-        }
+        return Distributions.clustered(Spark3Util.toTransforms(table.spec()));
 
       case RANGE:
-        if (table.spec().isUnpartitioned() && table.sortOrder().isUnsorted()) {
-          return Distributions.unspecified();
-        } else {
-          org.apache.iceberg.SortOrder requiredSortOrder = SortOrderUtil.buildSortOrder(table);
-          return Distributions.ordered(convert(requiredSortOrder));
-        }
+        org.apache.iceberg.SortOrder requiredSortOrder = SortOrderUtil.buildSortOrder(table);
+        return Distributions.ordered(convert(requiredSortOrder));
 
       default:
         throw new IllegalArgumentException("Unsupported distribution mode: " + distributionMode);
@@ -71,7 +63,7 @@ public static SortOrder[] buildRequiredOrdering(Table table, Distribution distri
   }
 
   public static SortOrder[] convert(org.apache.iceberg.SortOrder sortOrder) {
-    List<OrderField> converted = SortOrderVisitor.visit(sortOrder, new SortOrderToSpark(sortOrder.schema()));
-    return converted.toArray(new OrderField[0]);
+    List<SortOrder> converted = SortOrderVisitor.visit(sortOrder, new SortOrderToSpark(sortOrder.schema()));
+    return converted.toArray(new SortOrder[0]);
   }
 }
diff --git a/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/SparkWriteConf.java b/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/SparkWriteConf.java
@@ -30,6 +30,10 @@
 import org.apache.spark.sql.RuntimeConfig;
 import org.apache.spark.sql.SparkSession;
 
+import static org.apache.iceberg.DistributionMode.HASH;
+import static org.apache.iceberg.DistributionMode.NONE;
+import static org.apache.iceberg.DistributionMode.RANGE;
+
 /**
  * A class for common Iceberg configs for Spark writes.
  * <p>
@@ -155,19 +159,23 @@ public String rewrittenFileSetId() {
   }
 
   public DistributionMode distributionMode() {
-    String defaultValue;
-    if (table.sortOrder().isSorted()) {
-      defaultValue = TableProperties.WRITE_DISTRIBUTION_MODE_RANGE;
-    } else {
-      defaultValue = TableProperties.WRITE_DISTRIBUTION_MODE_DEFAULT;
-    }
-
     String modeName = confParser.stringConf()
         .option(SparkWriteOptions.DISTRIBUTION_MODE)
         .tableProperty(TableProperties.WRITE_DISTRIBUTION_MODE)
-        .defaultValue(defaultValue)
-        .parse();
-    return DistributionMode.fromName(modeName);
+        .parseOptional();
+
+    if (modeName != null) {
+      DistributionMode mode = DistributionMode.fromName(modeName);
+      if (mode == RANGE && table.spec().isUnpartitioned() && table.sortOrder().isUnsorted()) {
+        return NONE;
+      } else if (mode == HASH && table.spec().isUnpartitioned()) {
+        return NONE;
+      } else {
+        return mode;
+      }
+    } else {
+      return table.sortOrder().isSorted() ? RANGE : NONE;
+    }
   }
 
   public DistributionMode deleteDistributionMode() {