Refactor bucket functions.

Author: zhongyujiang

paimon-common/src/main/java/org/apache/paimon/data/serializer/InternalRowSerializer.java

@@ -24,14 +24,20 @@
 import org.apache.paimon.data.BinaryRowWriter;
 import org.apache.paimon.data.BinaryWriter;
 import org.apache.paimon.data.BinaryWriter.ValueSetter;
+import org.apache.paimon.data.Decimal;
 import org.apache.paimon.data.GenericRow;
 import org.apache.paimon.data.InternalRow;
 import org.apache.paimon.data.InternalRow.FieldGetter;
 import org.apache.paimon.data.NestedRow;
+import org.apache.paimon.data.Timestamp;
 import org.apache.paimon.io.DataInputView;
 import org.apache.paimon.io.DataOutputView;
 import org.apache.paimon.types.DataType;
+import org.apache.paimon.types.DataTypeChecks;
+import org.apache.paimon.types.DecimalType;
+import org.apache.paimon.types.LocalZonedTimestampType;
 import org.apache.paimon.types.RowType;
+import org.apache.paimon.types.TimestampType;
 
 import java.io.IOException;
 import java.util.Arrays;
@@ -46,10 +52,20 @@ public class InternalRowSerializer extends AbstractRowDataSerializer<InternalRow
     private final Serializer[] fieldSerializers;
     private final FieldGetter[] fieldGetters;
     private final ValueSetter[] valueSetters;
+    private final boolean[] writeNulls;
 
     private transient BinaryRow reuseRow;
     private transient BinaryRowWriter reuseWriter;
 
+    public InternalRowSerializer(RowType rowType, boolean isBucketKeySerializer) {
+        this(
+                rowType.getFieldTypes().toArray(new DataType[0]),
+                rowType.getFieldTypes().stream()
+                        .map(InternalSerializers::create)
+                        .toArray(Serializer[]::new),
+                isBucketKeySerializer);
+    }
+
     public InternalRowSerializer(RowType rowType) {
         this(
                 rowType.getFieldTypes().toArray(new DataType[0]),
@@ -65,16 +81,34 @@ public InternalRowSerializer(DataType... types) {
     }
 
     public InternalRowSerializer(DataType[] types, Serializer<?>[] fieldSerializers) {
+        this(types, fieldSerializers, false);
+    }
+
+    private InternalRowSerializer(
+            DataType[] types, Serializer<?>[] fieldSerializers, boolean isBucketKeySerializer) {
         this.types = types;
         this.fieldSerializers = fieldSerializers;
         this.binarySerializer = new BinaryRowSerializer(types.length);
         this.fieldGetters = new FieldGetter[types.length];
         this.valueSetters = new ValueSetter[types.length];
+        this.writeNulls = new boolean[types.length];
+        Arrays.fill(writeNulls, true);
         for (int i = 0; i < types.length; i++) {
             DataType type = types[i];
             fieldGetters[i] = InternalRow.createFieldGetter(type, i);
             // pass serializer to avoid infinite loop
             valueSetters[i] = BinaryWriter.createValueSetter(type, fieldSerializers[i]);
+
+            // BucketKey projection use different write logic for non-compact type nulls,
+            // reference: org.apache.paimon.codegen.GenerateUtils.binaryWriterWriteNull
+            if (isBucketKeySerializer) {
+                if (type instanceof DecimalType) {
+                    writeNulls[i] = Decimal.isCompact(DataTypeChecks.getPrecision(type));
+                } else if (type instanceof TimestampType
+                        || type instanceof LocalZonedTimestampType) {
+                    writeNulls[i] = Timestamp.isCompact(DataTypeChecks.getPrecision(type));
+                }
+            }
         }
     }
 
@@ -157,7 +191,7 @@ public BinaryRow toBinaryRow(InternalRow row) {
         reuseWriter.writeRowKind(row.getRowKind());
         for (int i = 0; i < types.length; i++) {
             Object field = fieldGetters[i].getFieldOrNull(row);
-            if (field == null) {
+            if (field == null && writeNulls[i]) {
                 reuseWriter.setNullAt(i);
             } else {
                 valueSetters[i].setValue(reuseWriter, i, field);

paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkInternalRowWrapper.java

@@ -40,12 +40,13 @@
 import org.apache.spark.sql.types.TimestampNTZType;
 import org.apache.spark.sql.types.TimestampType;
 
+import java.io.Serializable;
 import java.math.BigDecimal;
 
 /** Wrapper to fetch value from the spark internal row. */
-public class SparkInternalRowWrapper implements InternalRow {
+public class SparkInternalRowWrapper implements InternalRow, Serializable {
 
-    private org.apache.spark.sql.catalyst.InternalRow internalRow;
+    private transient org.apache.spark.sql.catalyst.InternalRow internalRow;
     private final int length;
     private final int rowKindIdx;
     private final StructType structType;