make StructInternalRow null safe

apache · Nov 5, 2020 · b0e3e45 · b0e3e45
1 parent 7cfe754
commit b0e3e45
Showing 1 changed file with 30 additions and 6 deletions.
diff --git a/spark/src/main/java/org/apache/iceberg/spark/source/StructInternalRow.java b/spark/src/main/java/org/apache/iceberg/spark/source/StructInternalRow.java
@@ -135,17 +135,29 @@ public double getDouble(int ordinal) {
 
   @Override
   public Decimal getDecimal(int ordinal, int precision, int scale) {
+    return isNullAt(ordinal) ? null : getDecimalInternal(ordinal, precision, scale);
+  }
+
+  private Decimal getDecimalInternal(int ordinal, int precision, int scale) {
     return Decimal.apply(struct.get(ordinal, BigDecimal.class));
   }
 
   @Override
   public UTF8String getUTF8String(int ordinal) {
+    return isNullAt(ordinal) ? null : getUTF8StringInternal(ordinal);
+  }
+
+  private UTF8String getUTF8StringInternal(int ordinal) {
     CharSequence seq = struct.get(ordinal, CharSequence.class);
     return UTF8String.fromString(seq.toString());
   }
 
   @Override
   public byte[] getBinary(int ordinal) {
+    return isNullAt(ordinal) ? null : getBinaryInternal(ordinal);
+  }
+
+  private byte[] getBinaryInternal(int ordinal) {
     Object bytes = struct.get(ordinal, Object.class);
 
     // should only be either ByteBuffer or byte[]
@@ -165,20 +177,32 @@ public CalendarInterval getInterval(int ordinal) {
 
   @Override
   public InternalRow getStruct(int ordinal, int numFields) {
+    return isNullAt(ordinal) ? null : getStructInternal(ordinal, numFields);
+  }
+
+  private InternalRow getStructInternal(int ordinal, int numFields) {
     return new StructInternalRow(
         type.fields().get(ordinal).type().asStructType(),
         struct.get(ordinal, StructLike.class));
   }
 
   @Override
   public ArrayData getArray(int ordinal) {
+    return isNullAt(ordinal) ? null : getArrayInternal(ordinal);
+  }
+
+  private ArrayData getArrayInternal(int ordinal) {
     return collectionToArrayData(
         type.fields().get(ordinal).type().asListType().elementType(),
         struct.get(ordinal, Collection.class));
   }
 
   @Override
   public MapData getMap(int ordinal) {
+    return isNullAt(ordinal) ? null : getMapInternal(ordinal);
+  }
+
+  private MapData getMapInternal(int ordinal) {
     return mapToMapData(type.fields().get(ordinal).type().asMapType(), struct.get(ordinal, Map.class));
   }
 
@@ -194,22 +218,22 @@ public Object get(int ordinal, DataType dataType) {
     } else if (dataType instanceof LongType) {
       return getLong(ordinal);
     } else if (dataType instanceof StringType) {
-      return getUTF8String(ordinal);
+      return getUTF8StringInternal(ordinal);
     } else if (dataType instanceof FloatType) {
       return getFloat(ordinal);
     } else if (dataType instanceof DoubleType) {
       return getDouble(ordinal);
     } else if (dataType instanceof DecimalType) {
       DecimalType decimalType = (DecimalType) dataType;
-      return getDecimal(ordinal, decimalType.precision(), decimalType.scale());
+      return getDecimalInternal(ordinal, decimalType.precision(), decimalType.scale());
     } else if (dataType instanceof BinaryType) {
-      return getBinary(ordinal);
+      return getBinaryInternal(ordinal);
     } else if (dataType instanceof StructType) {
-      return getStruct(ordinal, ((StructType) dataType).size());
+      return getStructInternal(ordinal, ((StructType) dataType).size());
     } else if (dataType instanceof ArrayType) {
-      return getArray(ordinal);
+      return getArrayInternal(ordinal);
     } else if (dataType instanceof MapType) {
-      return getMap(ordinal);
+      return getMapInternal(ordinal);
     } else if (dataType instanceof BooleanType) {
       return getBoolean(ordinal);
     } else if (dataType instanceof ByteType) {