Spark 3.2: Align formatting in bucket and truncate functions (apache#…

…5574)
yabola · Aug 19, 2022 · 58368f5 · 58368f5
1 parent 34b8349
commit 58368f5
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 18 deletions.
diff --git a/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/functions/BucketFunction.java b/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/functions/BucketFunction.java
@@ -53,8 +53,10 @@
  * results.
  */
 public class BucketFunction implements UnboundFunction {
+
   private static final int NUM_BUCKETS_ORDINAL = 0;
   private static final int VALUE_ORDINAL = 1;
+
   private static final Set<DataType> SUPPORTED_NUM_BUCKETS_TYPES =
       ImmutableSet.of(DataTypes.ByteType, DataTypes.ShortType, DataTypes.IntegerType);
 
@@ -156,9 +158,11 @@ public String canonicalName() {
     @Override
     public Integer produceResult(InternalRow input) {
       // return null for null input to match what Spark does in the code-generated versions.
-      return input.isNullAt(NUM_BUCKETS_ORDINAL) || input.isNullAt(VALUE_ORDINAL)
-          ? null
-          : invoke(input.getInt(NUM_BUCKETS_ORDINAL), input.getInt(VALUE_ORDINAL));
+      if (input.isNullAt(NUM_BUCKETS_ORDINAL) || input.isNullAt(VALUE_ORDINAL)) {
+        return null;
+      } else {
+        return invoke(input.getInt(NUM_BUCKETS_ORDINAL), input.getInt(VALUE_ORDINAL));
+      }
     }
   }
 
@@ -192,9 +196,11 @@ public String canonicalName() {
 
     @Override
     public Integer produceResult(InternalRow input) {
-      return input.isNullAt(NUM_BUCKETS_ORDINAL) || input.isNullAt(VALUE_ORDINAL)
-          ? null
-          : invoke(input.getInt(NUM_BUCKETS_ORDINAL), input.getLong(VALUE_ORDINAL));
+      if (input.isNullAt(NUM_BUCKETS_ORDINAL) || input.isNullAt(VALUE_ORDINAL)) {
+        return null;
+      } else {
+        return invoke(input.getInt(NUM_BUCKETS_ORDINAL), input.getLong(VALUE_ORDINAL));
+      }
     }
   }
 
@@ -226,9 +232,11 @@ public String canonicalName() {
 
     @Override
     public Integer produceResult(InternalRow input) {
-      return input.isNullAt(NUM_BUCKETS_ORDINAL) || input.isNullAt(VALUE_ORDINAL)
-          ? null
-          : invoke(input.getInt(NUM_BUCKETS_ORDINAL), input.getUTF8String(VALUE_ORDINAL));
+      if (input.isNullAt(NUM_BUCKETS_ORDINAL) || input.isNullAt(VALUE_ORDINAL)) {
+        return null;
+      } else {
+        return invoke(input.getInt(NUM_BUCKETS_ORDINAL), input.getUTF8String(VALUE_ORDINAL));
+      }
     }
   }
 
@@ -253,9 +261,11 @@ public DataType[] inputTypes() {
 
     @Override
     public Integer produceResult(InternalRow input) {
-      return input.isNullAt(NUM_BUCKETS_ORDINAL) || input.isNullAt(VALUE_ORDINAL)
-          ? null
-          : invoke(input.getInt(NUM_BUCKETS_ORDINAL), input.getBinary(VALUE_ORDINAL));
+      if (input.isNullAt(NUM_BUCKETS_ORDINAL) || input.isNullAt(VALUE_ORDINAL)) {
+        return null;
+      } else {
+        return invoke(input.getInt(NUM_BUCKETS_ORDINAL), input.getBinary(VALUE_ORDINAL));
+      }
     }
 
     @Override
@@ -296,10 +306,13 @@ public DataType[] inputTypes() {
 
     @Override
     public Integer produceResult(InternalRow input) {
-      return input.isNullAt(NUM_BUCKETS_ORDINAL) || input.isNullAt(VALUE_ORDINAL)
-          ? null
-          : invoke(
-              input.getInt(NUM_BUCKETS_ORDINAL), input.getDecimal(VALUE_ORDINAL, precision, scale));
+      if (input.isNullAt(NUM_BUCKETS_ORDINAL) || input.isNullAt(VALUE_ORDINAL)) {
+        return null;
+      } else {
+        int numBuckets = input.getInt(NUM_BUCKETS_ORDINAL);
+        Decimal value = input.getDecimal(VALUE_ORDINAL, precision, scale);
+        return invoke(numBuckets, value);
+      }
     }
 
     @Override

diff --git a/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/functions/TruncateFunction.java b/spark/v3.2/spark/src/main/java/org/apache/iceberg/spark/functions/TruncateFunction.java
@@ -348,8 +348,9 @@ public Decimal produceResult(InternalRow input) {
       if (input.isNullAt(WIDTH_ORDINAL) || input.isNullAt(VALUE_ORDINAL)) {
         return null;
       } else {
-        return invoke(
-            input.getInt(WIDTH_ORDINAL), input.getDecimal(VALUE_ORDINAL, precision, scale));
+        int width = input.getInt(WIDTH_ORDINAL);
+        Decimal value = input.getDecimal(VALUE_ORDINAL, precision, scale);
+        return invoke(width, value);
       }
     }
   }