Add mechanism for 'safe' memory reads for complex types

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchMergeComplexMetricSerde.java

@@ -28,6 +28,7 @@
 import org.apache.druid.segment.column.ColumnBuilder;
 import org.apache.druid.segment.data.GenericIndexed;
 import org.apache.druid.segment.data.ObjectStrategy;
+import org.apache.druid.segment.data.SafeWritableMemory;
 import org.apache.druid.segment.serde.ComplexColumnPartSupplier;
 import org.apache.druid.segment.serde.ComplexMetricExtractor;
 import org.apache.druid.segment.serde.ComplexMetricSerde;
@@ -70,7 +71,7 @@ public HllSketch extractValue(final InputRow inputRow, final String metricName)
         if (object == null) {
           return null;
         }
-        return deserializeSketch(object);
+        return deserializeSketchSafe(object);
       }
     };
   }
@@ -98,6 +99,18 @@ static HllSketch deserializeSketch(final Object object)
     throw new IAE("Object is not of a type that can be deserialized to an HllSketch:" + object.getClass().getName());
   }
 
+  static HllSketch deserializeSketchSafe(final Object object)
+  {
+    if (object instanceof String) {
+      return HllSketch.wrap(SafeWritableMemory.wrap(StringUtils.decodeBase64(((String) object).getBytes(StandardCharsets.UTF_8))));
+    } else if (object instanceof byte[]) {
+      return HllSketch.wrap(SafeWritableMemory.wrap((byte[]) object));
+    } else if (object instanceof HllSketch) {
+      return (HllSketch) object;
+    }
+    throw new IAE("Object is not of a type that can be deserialized to an HllSketch:" + object.getClass().getName());
+  }
+
   // support large columns
   @Override
   public GenericColumnSerializer getSerializer(final SegmentWriteOutMedium segmentWriteOutMedium, final String column)

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchObjectStrategy.java

@@ -22,7 +22,9 @@
 import org.apache.datasketches.hll.HllSketch;
 import org.apache.datasketches.memory.Memory;
 import org.apache.druid.segment.data.ObjectStrategy;
+import org.apache.druid.segment.data.SafeWritableMemory;
 
+import javax.annotation.Nullable;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 
@@ -55,4 +57,12 @@ public byte[] toBytes(final HllSketch sketch)
     return sketch.toCompactByteArray();
   }
 
+  @Nullable
+  @Override
+  public HllSketch fromByteBufferSafe(ByteBuffer buffer, int numBytes)
+  {
+    return HllSketch.wrap(
+        SafeWritableMemory.wrap(buffer, ByteOrder.LITTLE_ENDIAN).region(buffer.position(), numBytes)
+    );
+  }
 }

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/kll/KllDoublesSketchComplexMetricSerde.java

@@ -91,7 +91,7 @@ public Object extractValue(final InputRow inputRow, final String metricName)
         if (object == null || object instanceof KllDoublesSketch || object instanceof Memory) {
           return object;
         }
-        return KllDoublesSketchOperations.deserialize(object);
+        return KllDoublesSketchOperations.deserializeSafe(object);
       }
     };
   }

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/kll/KllDoublesSketchObjectStrategy.java

@@ -23,7 +23,9 @@
 import org.apache.datasketches.kll.KllDoublesSketch;
 import org.apache.datasketches.memory.Memory;
 import org.apache.druid.segment.data.ObjectStrategy;
+import org.apache.druid.segment.data.SafeWritableMemory;
 
+import javax.annotation.Nullable;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 
@@ -60,4 +62,15 @@ public byte[] toBytes(final KllDoublesSketch sketch)
     return sketch.toByteArray();
   }
 
+  @Nullable
+  @Override
+  public KllDoublesSketch fromByteBufferSafe(ByteBuffer buffer, int numBytes)
+  {
+    if (numBytes == 0) {
+      return KllDoublesSketchOperations.EMPTY_SKETCH;
+    }
+    return KllDoublesSketch.wrap(
+        SafeWritableMemory.wrap(buffer, ByteOrder.LITTLE_ENDIAN).region(buffer.position(), numBytes)
+    );
+  }
 }

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/kll/KllDoublesSketchOperations.java

@@ -23,6 +23,7 @@
 import org.apache.datasketches.memory.Memory;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.segment.data.SafeWritableMemory;
 
 import java.nio.charset.StandardCharsets;
 
@@ -46,6 +47,16 @@ public static KllDoublesSketch deserialize(final Object serializedSketch)
     );
   }
 
+  public static KllDoublesSketch deserializeSafe(final Object serializedSketch)
+  {
+    if (serializedSketch instanceof String) {
+      return deserializeFromBase64EncodedStringSafe((String) serializedSketch);
+    } else if (serializedSketch instanceof byte[]) {
+      return deserializeFromByteArraySafe((byte[]) serializedSketch);
+    }
+    return deserialize(serializedSketch);
+  }
+
   public static KllDoublesSketch deserializeFromBase64EncodedString(final String str)
   {
     return deserializeFromByteArray(StringUtils.decodeBase64(str.getBytes(StandardCharsets.UTF_8)));
@@ -56,4 +67,14 @@ public static KllDoublesSketch deserializeFromByteArray(final byte[] data)
     return KllDoublesSketch.wrap(Memory.wrap(data));
   }
 
+  public static KllDoublesSketch deserializeFromBase64EncodedStringSafe(final String str)
+  {
+    return deserializeFromByteArraySafe(StringUtils.decodeBase64(str.getBytes(StandardCharsets.UTF_8)));
+  }
+
+  public static KllDoublesSketch deserializeFromByteArraySafe(final byte[] data)
+  {
+    return KllDoublesSketch.wrap(SafeWritableMemory.wrap(data));
+  }
+
 }

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/kll/KllFloatsSketchComplexMetricSerde.java

@@ -91,7 +91,7 @@ public Object extractValue(final InputRow inputRow, final String metricName)
         if (object == null || object instanceof KllFloatsSketch || object instanceof Memory) {
           return object;
         }
-        return KllFloatsSketchOperations.deserialize(object);
+        return KllFloatsSketchOperations.deserializeSafe(object);
       }
     };
   }

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/kll/KllFloatsSketchObjectStrategy.java

@@ -23,7 +23,9 @@
 import org.apache.datasketches.kll.KllFloatsSketch;
 import org.apache.datasketches.memory.Memory;
 import org.apache.druid.segment.data.ObjectStrategy;
+import org.apache.druid.segment.data.SafeWritableMemory;
 
+import javax.annotation.Nullable;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 
@@ -60,4 +62,15 @@ public byte[] toBytes(final KllFloatsSketch sketch)
     return sketch.toByteArray();
   }
 
+  @Nullable
+  @Override
+  public KllFloatsSketch fromByteBufferSafe(ByteBuffer buffer, int numBytes)
+  {
+    if (numBytes == 0) {
+      return KllFloatsSketchOperations.EMPTY_SKETCH;
+    }
+    return KllFloatsSketch.wrap(
+        SafeWritableMemory.wrap(buffer, ByteOrder.LITTLE_ENDIAN).region(buffer.position(), numBytes)
+    );
+  }
 }

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/kll/KllFloatsSketchOperations.java

@@ -23,6 +23,7 @@
 import org.apache.datasketches.memory.Memory;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.segment.data.SafeWritableMemory;
 
 import java.nio.charset.StandardCharsets;
 
@@ -46,6 +47,16 @@ public static KllFloatsSketch deserialize(final Object serializedSketch)
     );
   }
 
+  public static KllFloatsSketch deserializeSafe(final Object serializedSketch)
+  {
+    if (serializedSketch instanceof String) {
+      return deserializeFromBase64EncodedStringSafe((String) serializedSketch);
+    } else if (serializedSketch instanceof byte[]) {
+      return deserializeFromByteArraySafe((byte[]) serializedSketch);
+    }
+    return deserialize(serializedSketch);
+  }
+
   public static KllFloatsSketch deserializeFromBase64EncodedString(final String str)
   {
     return deserializeFromByteArray(StringUtils.decodeBase64(str.getBytes(StandardCharsets.UTF_8)));
@@ -56,4 +67,14 @@ public static KllFloatsSketch deserializeFromByteArray(final byte[] data)
     return KllFloatsSketch.wrap(Memory.wrap(data));
   }
 
+  public static KllFloatsSketch deserializeFromBase64EncodedStringSafe(final String str)
+  {
+    return deserializeFromByteArraySafe(StringUtils.decodeBase64(str.getBytes(StandardCharsets.UTF_8)));
+  }
+
+  public static KllFloatsSketch deserializeFromByteArraySafe(final byte[] data)
+  {
+    return KllFloatsSketch.wrap(SafeWritableMemory.wrap(data));
+  }
+
 }

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/quantiles/DoublesSketchComplexMetricSerde.java

@@ -92,7 +92,7 @@ public Object extractValue(final InputRow inputRow, final String metricName)
         if (object == null || object instanceof DoublesSketch || object instanceof Memory) {
           return object;
         }
-        return DoublesSketchOperations.deserialize(object);
+        return DoublesSketchOperations.deserializeSafe(object);
       }
     };
   }

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/quantiles/DoublesSketchObjectStrategy.java

@@ -23,7 +23,9 @@
 import org.apache.datasketches.memory.Memory;
 import org.apache.datasketches.quantiles.DoublesSketch;
 import org.apache.druid.segment.data.ObjectStrategy;
+import org.apache.druid.segment.data.SafeWritableMemory;
 
+import javax.annotation.Nullable;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 
@@ -60,4 +62,15 @@ public byte[] toBytes(final DoublesSketch sketch)
     return sketch.toByteArray(true);
   }
 
+  @Nullable
+  @Override
+  public DoublesSketch fromByteBufferSafe(ByteBuffer buffer, int numBytes)
+  {
+    if (numBytes == 0) {
+      return DoublesSketchOperations.EMPTY_SKETCH;
+    }
+    return DoublesSketch.wrap(
+        SafeWritableMemory.wrap(buffer, ByteOrder.LITTLE_ENDIAN).region(buffer.position(), numBytes)
+    );
+  }
 }

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/quantiles/DoublesSketchOperations.java

@@ -23,6 +23,7 @@
 import org.apache.datasketches.quantiles.DoublesSketch;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.segment.data.SafeWritableMemory;
 
 import java.nio.charset.StandardCharsets;
 
@@ -46,6 +47,16 @@ public static DoublesSketch deserialize(final Object serializedSketch)
     );
   }
 
+  public static DoublesSketch deserializeSafe(final Object serializedSketch)
+  {
+    if (serializedSketch instanceof String) {
+      return deserializeFromBase64EncodedStringSafe((String) serializedSketch);
+    } else if (serializedSketch instanceof byte[]) {
+      return deserializeFromByteArraySafe((byte[]) serializedSketch);
+    }
+    return deserialize(serializedSketch);
+  }
+
   public static DoublesSketch deserializeFromBase64EncodedString(final String str)
   {
     return deserializeFromByteArray(StringUtils.decodeBase64(str.getBytes(StandardCharsets.UTF_8)));
@@ -56,4 +67,13 @@ public static DoublesSketch deserializeFromByteArray(final byte[] data)
     return DoublesSketch.wrap(Memory.wrap(data));
   }
 
+  public static DoublesSketch deserializeFromBase64EncodedStringSafe(final String str)
+  {
+    return deserializeFromByteArraySafe(StringUtils.decodeBase64(str.getBytes(StandardCharsets.UTF_8)));
+  }
+
+  public static DoublesSketch deserializeFromByteArraySafe(final byte[] data)
+  {
+    return DoublesSketch.wrap(SafeWritableMemory.wrap(data));
+  }
 }

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/theta/SketchConstantPostAggregator.java

@@ -51,7 +51,7 @@ public SketchConstantPostAggregator(@JsonProperty("name") String name, @JsonProp
     Preconditions.checkArgument(value != null && !value.isEmpty(),
         "Constant value cannot be null or empty, expecting base64 encoded sketch string");
     this.value = value;
-    this.sketchValue = SketchHolder.deserialize(value);
+    this.sketchValue = SketchHolder.deserializeSafe(value);
   }
 
   @Override

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/theta/SketchHolder.java

@@ -34,6 +34,7 @@
 import org.apache.druid.java.util.common.IAE;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.segment.data.SafeWritableMemory;
 
 import javax.annotation.Nullable;
 
@@ -224,6 +225,17 @@ public static SketchHolder deserialize(Object serializedSketch)
     );
   }
 
+  public static SketchHolder deserializeSafe(Object serializedSketch)
+  {
+    if (serializedSketch instanceof String) {
+      return SketchHolder.of(deserializeFromBase64EncodedStringSafe((String) serializedSketch));
+    } else if (serializedSketch instanceof byte[]) {
+      return SketchHolder.of(deserializeFromByteArraySafe((byte[]) serializedSketch));
+    }
+
+    return deserialize(serializedSketch);
+  }
+
   private static Sketch deserializeFromBase64EncodedString(String str)
   {
     return deserializeFromByteArray(StringUtils.decodeBase64(StringUtils.toUtf8(str)));
@@ -234,6 +246,16 @@ private static Sketch deserializeFromByteArray(byte[] data)
     return deserializeFromMemory(Memory.wrap(data));
   }
 
+  private static Sketch deserializeFromBase64EncodedStringSafe(String str)
+  {
+    return deserializeFromByteArraySafe(StringUtils.decodeBase64(StringUtils.toUtf8(str)));
+  }
+
+  private static Sketch deserializeFromByteArraySafe(byte[] data)
+  {
+    return deserializeFromMemory(SafeWritableMemory.wrap(data));
+  }
+
   private static Sketch deserializeFromMemory(Memory mem)
   {
     if (Sketch.getSerializationVersion(mem) < 3) {

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/theta/SketchHolderObjectStrategy.java

@@ -23,6 +23,7 @@
 import org.apache.datasketches.memory.Memory;
 import org.apache.datasketches.theta.Sketch;
 import org.apache.druid.segment.data.ObjectStrategy;
+import org.apache.druid.segment.data.SafeWritableMemory;
 
 import javax.annotation.Nullable;
 import java.nio.ByteBuffer;
@@ -66,4 +67,17 @@ public byte[] toBytes(@Nullable SketchHolder obj)
       return ByteArrays.EMPTY_ARRAY;
     }
   }
+
+  @Nullable
+  @Override
+  public SketchHolder fromByteBufferSafe(ByteBuffer buffer, int numBytes)
+  {
+    if (numBytes == 0) {
+      return SketchHolder.EMPTY;
+    }
+
+    return SketchHolder.of(
+        SafeWritableMemory.wrap(buffer, ByteOrder.LITTLE_ENDIAN).region(buffer.position(), numBytes)
+    );
+  }
 }

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/theta/SketchMergeComplexMetricSerde.java

@@ -59,7 +59,7 @@ public Class<SketchHolder> extractedClass()
       public SketchHolder extractValue(InputRow inputRow, String metricName)
       {
         final Object object = inputRow.getRaw(metricName);
-        return object == null ? null : SketchHolder.deserialize(object);
+        return object == null ? null : SketchHolder.deserializeSafe(object);
       }
     };
   }

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/tuple/ArrayOfDoublesSketchMergeComplexMetricSerde.java

@@ -60,7 +60,7 @@ public Object extractValue(final InputRow inputRow, final String metricName)
         if (object == null || object instanceof ArrayOfDoublesSketch) {
           return object;
         }
-        return ArrayOfDoublesSketchOperations.deserialize(object);
+        return ArrayOfDoublesSketchOperations.deserializeSafe(object);
       }
     };
   }