[fix] fix get seatunnel row size npe issue (#6681)

apache · Apr 11, 2024 · 4f4fd7b · 4f4fd7b
1 parent 6296d0c
commit 4f4fd7b
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 35 deletions.
diff --git a/seatunnel-api/src/main/java/org/apache/seatunnel/api/table/type/SeaTunnelRow.java b/seatunnel-api/src/main/java/org/apache/seatunnel/api/table/type/SeaTunnelRow.java
@@ -171,29 +171,39 @@ private int getBytesForArray(Object v, BasicType<?> dataType) {
             case STRING:
                 int s = 0;
                 for (String i : ((String[]) v)) {
-                    s += i.length();
+                    s += i == null ? 0 : i.length();
                 }
                 return s;
             case BOOLEAN:
-                return ((Boolean[]) v).length;
+                return getArrayNotNullSize((Boolean[]) v);
             case TINYINT:
-                return ((Byte[]) v).length;
+                return getArrayNotNullSize((Byte[]) v);
             case SMALLINT:
-                return ((Short[]) v).length * 2;
+                return getArrayNotNullSize((Short[]) v) * 2;
             case INT:
-                return ((Integer[]) v).length * 4;
+                return getArrayNotNullSize((Integer[]) v) * 4;
             case FLOAT:
-                return ((Float[]) v).length * 4;
+                return getArrayNotNullSize((Float[]) v) * 4;
             case BIGINT:
-                return ((Long[]) v).length * 8;
+                return getArrayNotNullSize((Long[]) v) * 8;
             case DOUBLE:
-                return ((Double[]) v).length * 8;
+                return getArrayNotNullSize((Double[]) v) * 8;
             case NULL:
             default:
                 return 0;
         }
     }
 
+    private int getArrayNotNullSize(Object[] values) {
+        int c = 0;
+        for (Object value : values) {
+            if (value != null) {
+                c++;
+            }
+        }
+        return c;
+    }
+
     public int getBytesSize() {
         if (size == 0) {
             int s = 0;
@@ -235,25 +245,21 @@ private int getBytesForValue(Object v) {
             case "LocalDateTime":
                 return 48;
             case "String[]":
-                int s = 0;
-                for (String i : ((String[]) v)) {
-                    s += i.length();
-                }
-                return s;
+                return getBytesForArray(v, BasicType.STRING_TYPE);
             case "Boolean[]":
-                return ((Boolean[]) v).length;
+                return getBytesForArray(v, BasicType.BOOLEAN_TYPE);
             case "Byte[]":
-                return ((Byte[]) v).length;
+                return getBytesForArray(v, BasicType.BYTE_TYPE);
             case "Short[]":
-                return ((Short[]) v).length * 2;
+                return getBytesForArray(v, BasicType.SHORT_TYPE);
             case "Integer[]":
-                return ((Integer[]) v).length * 4;
+                return getBytesForArray(v, BasicType.INT_TYPE);
             case "Long[]":
-                return ((Long[]) v).length * 8;
+                return getBytesForArray(v, BasicType.LONG_TYPE);
             case "Float[]":
-                return ((Float[]) v).length * 4;
+                return getBytesForArray(v, BasicType.FLOAT_TYPE);
             case "Double[]":
-                return ((Double[]) v).length * 8;
+                return getBytesForArray(v, BasicType.DOUBLE_TYPE);
             case "HashMap":
             case "LinkedHashMap":
                 int size = 0;

diff --git a/seatunnel-api/src/test/java/org/apache/seatunnel/api/table/type/SeaTunnelRowTest.java b/seatunnel-api/src/test/java/org/apache/seatunnel/api/table/type/SeaTunnelRowTest.java
@@ -51,12 +51,40 @@ void testForRowSize() {
                             1L,
                             map,
                             new BigDecimal("3333.333"),
-                            new String[] {"test2", "test", "3333.333"}
+                            new String[] {"test2", "test", "3333.333"},
+                            new Integer[] {1, 2, 3},
+                            new Long[] {1L, 2L, 3L},
+                            new Double[] {1D, 2D},
+                            new Float[] {1F, 2F},
+                            new Boolean[] {Boolean.TRUE, Boolean.FALSE},
+                            new Byte[] {1, 2, 3, 4},
+                            new Short[] {Short.parseShort("1")}
+                        });
+
+        SeaTunnelRow row2 =
+                new SeaTunnelRow(
+                        new Object[] {
+                            1,
+                            "test",
+                            1L,
+                            map,
+                            new BigDecimal("3333.333"),
+                            new String[] {"test2", "test", "3333.333", null},
+                            new Integer[] {1, 2, 3, null},
+                            new Long[] {1L, 2L, 3L, null},
+                            new Double[] {1D, 2D, null},
+                            new Float[] {1F, 2F, null},
+                            new Boolean[] {Boolean.TRUE, Boolean.FALSE, null},
+                            new Byte[] {1, 2, 3, 4, null},
+                            new Short[] {Short.parseShort("1"), null}
                         });
 
         SeaTunnelRowType rowType =
                 new SeaTunnelRowType(
-                        new String[] {"f0", "f1", "f2", "f3", "f4", "f5"},
+                        new String[] {
+                            "f0", "f1", "f2", "f3", "f4", "f5", "f6", "f7", "f8", "f9", "f10",
+                            "f11", "f12"
+                        },
                         new SeaTunnelDataType<?>[] {
                             BasicType.INT_TYPE,
                             BasicType.STRING_TYPE,
@@ -72,22 +100,21 @@ void testForRowSize() {
                                                 new DecimalType(10, 3)
                                             })),
                             new DecimalType(10, 3),
-                            ArrayType.STRING_ARRAY_TYPE
+                            ArrayType.STRING_ARRAY_TYPE,
+                            ArrayType.INT_ARRAY_TYPE,
+                            ArrayType.LONG_ARRAY_TYPE,
+                            ArrayType.DOUBLE_ARRAY_TYPE,
+                            ArrayType.FLOAT_ARRAY_TYPE,
+                            ArrayType.BOOLEAN_ARRAY_TYPE,
+                            ArrayType.BYTE_ARRAY_TYPE,
+                            ArrayType.SHORT_ARRAY_TYPE
                         });
 
-        Assertions.assertEquals(181, row.getBytesSize(rowType));
+        Assertions.assertEquals(249, row.getBytesSize(rowType));
+        Assertions.assertEquals(249, row.getBytesSize());
 
-        SeaTunnelRow row2 =
-                new SeaTunnelRow(
-                        new Object[] {
-                            1,
-                            "test",
-                            1L,
-                            map,
-                            new BigDecimal("3333.333"),
-                            new String[] {"test2", "test", "3333.333"}
-                        });
-        Assertions.assertEquals(181, row2.getBytesSize());
+        Assertions.assertEquals(249, row2.getBytesSize(rowType));
+        Assertions.assertEquals(249, row2.getBytesSize());
     }
 
     @Test