Change *VectorAggregator to support pair serdes

apache · Oct 12, 2023 · f5a1c50 · f5a1c50
1 parent 97d1dac
commit f5a1c50
Show file tree

Hide file tree

Showing 48 changed files with 1,352 additions and 470 deletions.
diff --git a/...he/druid/query/aggregation/AbstractSerializablePairLongObjectDeltaEncodedStagedSerde.java b/...he/druid/query/aggregation/AbstractSerializablePairLongObjectDeltaEncodedStagedSerde.java
@@ -22,6 +22,7 @@
 import com.google.common.base.Preconditions;
 import com.google.common.primitives.Ints;
 import org.apache.druid.collections.SerializablePair;
+import org.apache.druid.common.config.NullHandling;
 import org.apache.druid.segment.serde.cell.StagedSerde;
 import org.apache.druid.segment.serde.cell.StorableBuffer;
 
@@ -75,13 +76,16 @@ public void store(ByteBuffer byteBuffer)
         }
 
         if (rhsObject != null) {
+          byteBuffer.put(NullHandling.IS_NOT_NULL_BYTE);
           if (pairClass.isAssignableFrom(SerializablePairLongLong.class)) {
             byteBuffer.putLong((long) rhsObject);
           } else if (pairClass.isAssignableFrom(SerializablePairLongDouble.class)) {
             byteBuffer.putDouble((double) rhsObject);
           } else if (pairClass.isAssignableFrom(SerializablePairLongFloat.class)) {
             byteBuffer.putFloat((float) rhsObject);
           }
+        } else {
+          byteBuffer.put(NullHandling.IS_NULL_BYTE);
         }
       }
 
@@ -100,7 +104,7 @@ public int getSerializedSize()
           }
         }
 
-        return (useIntegerDelta ? Integer.BYTES : Long.BYTES) + rhsBytes;
+        return (useIntegerDelta ? Integer.BYTES : Long.BYTES) + Byte.BYTES + rhsBytes;
       }
     };
   }

diff --git a/...g/apache/druid/query/aggregation/AbstractSerializablePairLongObjectSimpleStagedSerde.java b/...g/apache/druid/query/aggregation/AbstractSerializablePairLongObjectSimpleStagedSerde.java
@@ -21,6 +21,7 @@
 
 import com.google.common.base.Preconditions;
 import org.apache.druid.collections.SerializablePair;
+import org.apache.druid.common.config.NullHandling;
 import org.apache.druid.segment.serde.cell.StagedSerde;
 import org.apache.druid.segment.serde.cell.StorableBuffer;
 
@@ -57,13 +58,16 @@ public void store(ByteBuffer byteBuffer)
         Preconditions.checkNotNull(value.getLhs(), String.format(Locale.ENGLISH, "Long in %s must be non-null", pairCLass.getSimpleName()));
         byteBuffer.putLong(value.getLhs());
         if (rhsObject != null) {
+          byteBuffer.put(NullHandling.IS_NOT_NULL_BYTE);
           if (pairCLass.isAssignableFrom(SerializablePairLongLong.class)) {
             byteBuffer.putLong((long) rhsObject);
           } else if (pairCLass.isAssignableFrom(SerializablePairLongDouble.class)) {
             byteBuffer.putDouble((double) rhsObject);
           } else if (pairCLass.isAssignableFrom(SerializablePairLongFloat.class)) {
             byteBuffer.putFloat((float) rhsObject);
           }
+        } else {
+          byteBuffer.put(NullHandling.IS_NULL_BYTE);
         }
       }
 
@@ -81,7 +85,7 @@ public int getSerializedSize()
             rhsBytes = Float.BYTES;
           }
         }
-        return Long.BYTES + rhsBytes;
+        return Long.BYTES + Byte.BYTES + rhsBytes;
       }
     };
   }

diff --git a/...java/org/apache/druid/query/aggregation/SerializablePairLongDoubleComplexMetricSerde.java b/...java/org/apache/druid/query/aggregation/SerializablePairLongDoubleComplexMetricSerde.java
@@ -20,7 +20,6 @@
 package org.apache.druid.query.aggregation;
 
 import org.apache.druid.collections.SerializablePair;
-import org.apache.druid.common.config.NullHandling;
 import org.apache.druid.segment.GenericColumnSerializer;
 import org.apache.druid.segment.column.ColumnBuilder;
 import org.apache.druid.segment.data.ObjectStrategy;
@@ -35,6 +34,8 @@ public class SerializablePairLongDoubleComplexMetricSerde extends AbstractSerial
 {
   public static final String TYPE_NAME = "serializablePairLongDouble";
 
+  private static final SerializablePairLongDoubleSimpleStagedSerde SERDE = new SerializablePairLongDoubleSimpleStagedSerde();
+
   private static final Comparator<SerializablePair<Long, Double>> COMPARATOR = SerializablePair.createNullHandlingComparator(
       Double::compare,
       true
@@ -90,32 +91,17 @@ public Class<? extends SerializablePairLongDouble> getClazz()
       @Override
       public SerializablePairLongDouble fromByteBuffer(ByteBuffer buffer, int numBytes)
       {
-        final ByteBuffer readOnlyBuffer = buffer.asReadOnlyBuffer();
-        long lhs = readOnlyBuffer.getLong();
-        boolean isNotNull = readOnlyBuffer.get() == NullHandling.IS_NOT_NULL_BYTE;
-        if (isNotNull) {
-          return new SerializablePairLongDouble(lhs, readOnlyBuffer.getDouble());
-        } else {
-          return new SerializablePairLongDouble(lhs, null);
-        }
+        ByteBuffer readOnlyByteBuffer = buffer.asReadOnlyBuffer().order(buffer.order());
+
+        readOnlyByteBuffer.limit(buffer.position() + numBytes);
+
+        return SERDE.deserialize(readOnlyByteBuffer);
       }
 
       @Override
       public byte[] toBytes(@Nullable SerializablePairLongDouble inPair)
       {
-        if (inPair == null) {
-          return new byte[]{};
-        }
-
-        ByteBuffer bbuf = ByteBuffer.allocate(Long.BYTES + Byte.BYTES + Double.BYTES);
-        bbuf.putLong(inPair.lhs);
-        if (inPair.rhs == null) {
-          bbuf.put(NullHandling.IS_NULL_BYTE);
-        } else {
-          bbuf.put(NullHandling.IS_NOT_NULL_BYTE);
-          bbuf.putDouble(inPair.rhs);
-        }
-        return bbuf.array();
+        return SERDE.serialize(inPair);
       }
     };
   }

diff --git a/...org/apache/druid/query/aggregation/SerializablePairLongDoubleDeltaEncodedStagedSerde.java b/...org/apache/druid/query/aggregation/SerializablePairLongDoubleDeltaEncodedStagedSerde.java
@@ -19,6 +19,8 @@
 
 package org.apache.druid.query.aggregation;
 
+import org.apache.druid.common.config.NullHandling;
+
 import javax.annotation.Nullable;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
@@ -51,7 +53,7 @@ public SerializablePairLongDouble deserialize(ByteBuffer byteBuffer)
     lhs += minValue;
 
     Double rhs = null;
-    if (readOnlyBuffer.hasRemaining()) {
+    if (readOnlyBuffer.get() == NullHandling.IS_NOT_NULL_BYTE) {
       rhs = readOnlyBuffer.getDouble();
     }
 

diff --git a/.../java/org/apache/druid/query/aggregation/SerializablePairLongDoubleSimpleStagedSerde.java b/.../java/org/apache/druid/query/aggregation/SerializablePairLongDoubleSimpleStagedSerde.java
@@ -19,6 +19,8 @@
 
 package org.apache.druid.query.aggregation;
 
+import org.apache.druid.common.config.NullHandling;
+
 import javax.annotation.Nullable;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
@@ -42,7 +44,7 @@ public SerializablePairLongDouble deserialize(ByteBuffer byteBuffer)
     long lhs = readOnlyBuffer.getLong();
 
     Double rhs = null;
-    if (readOnlyBuffer.hasRemaining()) {
+    if (readOnlyBuffer.get() == NullHandling.IS_NOT_NULL_BYTE) {
       rhs = readOnlyBuffer.getDouble();
     }
 

diff --git a/.../java/org/apache/druid/query/aggregation/SerializablePairLongFloatComplexMetricSerde.java b/.../java/org/apache/druid/query/aggregation/SerializablePairLongFloatComplexMetricSerde.java
@@ -20,7 +20,6 @@
 package org.apache.druid.query.aggregation;
 
 import org.apache.druid.collections.SerializablePair;
-import org.apache.druid.common.config.NullHandling;
 import org.apache.druid.segment.GenericColumnSerializer;
 import org.apache.druid.segment.column.ColumnBuilder;
 import org.apache.druid.segment.data.ObjectStrategy;
@@ -35,6 +34,8 @@ public class SerializablePairLongFloatComplexMetricSerde extends AbstractSeriali
 {
   public static final String TYPE_NAME = "serializablePairLongFloat";
 
+  private static final SerializablePairLongFloatSimpleStagedSerde SERDE = new SerializablePairLongFloatSimpleStagedSerde();
+
   private static final Comparator<SerializablePair<Long, Float>> COMPARATOR = SerializablePair.createNullHandlingComparator(
       Float::compare,
       true
@@ -91,32 +92,17 @@ public Class<? extends SerializablePairLongFloat> getClazz()
       @Override
       public SerializablePairLongFloat fromByteBuffer(ByteBuffer buffer, int numBytes)
       {
-        final ByteBuffer readOnlyBuffer = buffer.asReadOnlyBuffer();
-        long lhs = readOnlyBuffer.getLong();
-        boolean isNotNull = readOnlyBuffer.get() == NullHandling.IS_NOT_NULL_BYTE;
-        if (isNotNull) {
-          return new SerializablePairLongFloat(lhs, readOnlyBuffer.getFloat());
-        } else {
-          return new SerializablePairLongFloat(lhs, null);
-        }
+        ByteBuffer readOnlyByteBuffer = buffer.asReadOnlyBuffer().order(buffer.order());
+
+        readOnlyByteBuffer.limit(buffer.position() + numBytes);
+
+        return SERDE.deserialize(readOnlyByteBuffer);
       }
 
       @Override
       public byte[] toBytes(@Nullable SerializablePairLongFloat inPair)
       {
-        if (inPair == null) {
-          return new byte[]{};
-        }
-
-        ByteBuffer bbuf = ByteBuffer.allocate(Long.BYTES + Byte.BYTES + Float.BYTES);
-        bbuf.putLong(inPair.lhs);
-        if (inPair.rhs == null) {
-          bbuf.put(NullHandling.IS_NULL_BYTE);
-        } else {
-          bbuf.put(NullHandling.IS_NOT_NULL_BYTE);
-          bbuf.putFloat(inPair.rhs);
-        }
-        return bbuf.array();
+        return SERDE.serialize(inPair);
       }
     };
   }

diff --git a/.../org/apache/druid/query/aggregation/SerializablePairLongFloatDeltaEncodedStagedSerde.java b/.../org/apache/druid/query/aggregation/SerializablePairLongFloatDeltaEncodedStagedSerde.java
@@ -19,6 +19,8 @@
 
 package org.apache.druid.query.aggregation;
 
+import org.apache.druid.common.config.NullHandling;
+
 import javax.annotation.Nullable;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
@@ -51,7 +53,7 @@ public SerializablePairLongFloat deserialize(ByteBuffer byteBuffer)
     lhs += minValue;
 
     Float rhs = null;
-    if (readOnlyBuffer.hasRemaining()) {
+    if (readOnlyBuffer.get() == NullHandling.IS_NOT_NULL_BYTE) {
       rhs = readOnlyBuffer.getFloat();
     }
 

diff --git a/...n/java/org/apache/druid/query/aggregation/SerializablePairLongFloatSimpleStagedSerde.java b/...n/java/org/apache/druid/query/aggregation/SerializablePairLongFloatSimpleStagedSerde.java
@@ -19,6 +19,8 @@
 
 package org.apache.druid.query.aggregation;
 
+import org.apache.druid.common.config.NullHandling;
+
 import javax.annotation.Nullable;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
@@ -42,7 +44,7 @@ public SerializablePairLongFloat deserialize(ByteBuffer byteBuffer)
     long lhs = readOnlyBuffer.getLong();
 
     Float rhs = null;
-    if (readOnlyBuffer.hasRemaining()) {
+    if (readOnlyBuffer.get() == NullHandling.IS_NOT_NULL_BYTE) {
       rhs = readOnlyBuffer.getFloat();
     }
 

diff --git a/...n/java/org/apache/druid/query/aggregation/SerializablePairLongLongComplexMetricSerde.java b/...n/java/org/apache/druid/query/aggregation/SerializablePairLongLongComplexMetricSerde.java
@@ -20,7 +20,6 @@
 package org.apache.druid.query.aggregation;
 
 import org.apache.druid.collections.SerializablePair;
-import org.apache.druid.common.config.NullHandling;
 import org.apache.druid.segment.GenericColumnSerializer;
 import org.apache.druid.segment.column.ColumnBuilder;
 import org.apache.druid.segment.data.ObjectStrategy;
@@ -35,6 +34,8 @@ public class SerializablePairLongLongComplexMetricSerde extends AbstractSerializ
 {
   public static final String TYPE_NAME = "serializablePairLongLong";
 
+  private static final SerializablePairLongLongSimpleStagedSerde SERDE = new SerializablePairLongLongSimpleStagedSerde();
+
   private static final Comparator<SerializablePair<Long, Long>> COMPARATOR = SerializablePair.createNullHandlingComparator(
       Long::compare,
       true
@@ -90,32 +91,17 @@ public Class<? extends SerializablePairLongLong> getClazz()
       @Override
       public SerializablePairLongLong fromByteBuffer(ByteBuffer buffer, int numBytes)
       {
-        final ByteBuffer readOnlyBuffer = buffer.asReadOnlyBuffer();
-        long lhs = readOnlyBuffer.getLong();
-        boolean isNotNull = readOnlyBuffer.get() == NullHandling.IS_NOT_NULL_BYTE;
-        if (isNotNull) {
-          return new SerializablePairLongLong(lhs, readOnlyBuffer.getLong());
-        } else {
-          return new SerializablePairLongLong(lhs, null);
-        }
+        ByteBuffer readOnlyByteBuffer = buffer.asReadOnlyBuffer().order(buffer.order());
+
+        readOnlyByteBuffer.limit(buffer.position() + numBytes);
+
+        return SERDE.deserialize(readOnlyByteBuffer);
       }
 
       @Override
       public byte[] toBytes(@Nullable SerializablePairLongLong inPair)
       {
-        if (inPair == null) {
-          return new byte[]{};
-        }
-
-        ByteBuffer bbuf = ByteBuffer.allocate(Long.BYTES + Byte.BYTES + Long.BYTES);
-        bbuf.putLong(inPair.lhs);
-        if (inPair.rhs == null) {
-          bbuf.put(NullHandling.IS_NULL_BYTE);
-        } else {
-          bbuf.put(NullHandling.IS_NOT_NULL_BYTE);
-          bbuf.putLong(inPair.rhs);
-        }
-        return bbuf.array();
+        return SERDE.serialize(inPair);
       }
     };
   }

diff --git a/...a/org/apache/druid/query/aggregation/SerializablePairLongLongDeltaEncodedStagedSerde.java b/...a/org/apache/druid/query/aggregation/SerializablePairLongLongDeltaEncodedStagedSerde.java
@@ -19,6 +19,8 @@
 
 package org.apache.druid.query.aggregation;
 
+import org.apache.druid.common.config.NullHandling;
+
 import javax.annotation.Nullable;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
@@ -51,7 +53,7 @@ public SerializablePairLongLong deserialize(ByteBuffer byteBuffer)
     lhs += minValue;
 
     Long rhs = null;
-    if (readOnlyBuffer.hasRemaining()) {
+    if (readOnlyBuffer.get() == NullHandling.IS_NOT_NULL_BYTE) {
       rhs = readOnlyBuffer.getLong();
     }
 

diff --git a/...in/java/org/apache/druid/query/aggregation/SerializablePairLongLongSimpleStagedSerde.java b/...in/java/org/apache/druid/query/aggregation/SerializablePairLongLongSimpleStagedSerde.java
@@ -19,6 +19,8 @@
 
 package org.apache.druid.query.aggregation;
 
+import org.apache.druid.common.config.NullHandling;
+
 import javax.annotation.Nullable;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
@@ -42,7 +44,7 @@ public SerializablePairLongLong deserialize(ByteBuffer byteBuffer)
     long lhs = readOnlyBuffer.getLong();
 
     Long rhs = null;
-    if (readOnlyBuffer.hasRemaining()) {
+    if (readOnlyBuffer.get() == NullHandling.IS_NOT_NULL_BYTE) {
       rhs = readOnlyBuffer.getLong();
     }
 

diff --git a/processing/src/main/java/org/apache/druid/query/aggregation/any/NilVectorAggregator.java b/processing/src/main/java/org/apache/druid/query/aggregation/any/NilVectorAggregator.java
@@ -19,8 +19,10 @@
 
 package org.apache.druid.query.aggregation.any;
 
-import org.apache.druid.collections.SerializablePair;
 import org.apache.druid.common.config.NullHandling;
+import org.apache.druid.query.aggregation.SerializablePairLongDouble;
+import org.apache.druid.query.aggregation.SerializablePairLongFloat;
+import org.apache.druid.query.aggregation.SerializablePairLongLong;
 import org.apache.druid.query.aggregation.VectorAggregator;
 
 import javax.annotation.Nullable;
@@ -43,9 +45,9 @@ public class NilVectorAggregator implements VectorAggregator
       NullHandling.defaultLongValue()
   );
 
-  public static final SerializablePair<Long, Double> DOUBLE_NIL_PAIR = new SerializablePair<>(0L, NullHandling.defaultDoubleValue());
-  public static final SerializablePair<Long, Long> LONG_NIL_PAIR = new SerializablePair<>(0L, NullHandling.defaultLongValue());
-  public static final SerializablePair<Long, Float> FLOAT_NIL_PAIR = new SerializablePair<>(0L, NullHandling.defaultFloatValue());
+  public static final SerializablePairLongDouble DOUBLE_NIL_PAIR = new SerializablePairLongDouble(0L, NullHandling.defaultDoubleValue());
+  public static final SerializablePairLongLong LONG_NIL_PAIR = new SerializablePairLongLong(0L, NullHandling.defaultLongValue());
+  public static final SerializablePairLongFloat FLOAT_NIL_PAIR = new SerializablePairLongFloat(0L, NullHandling.defaultFloatValue());
 
   /**
    * @return A vectorized aggregator that returns the default double value.