Adding the null check in combine and fold in doublesSketch (#14568)

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

@@ -240,8 +240,12 @@ public class DoublesSketchAggregatorFactory extends AggregatorFactory
   public Object combine(final Object lhs, final Object rhs)
   {
     final DoublesUnion union = DoublesUnion.builder().setMaxK(k).build();
-    union.union((DoublesSketch) lhs);
-    union.union((DoublesSketch) rhs);
+    if (lhs != null) {
+      union.union((DoublesSketch) lhs);
+    }
+    if (rhs != null) {
+      union.union((DoublesSketch) rhs);
+    }
     return union.getResultAndReset();
   }
 
@@ -263,7 +267,9 @@ public class DoublesSketchAggregatorFactory extends AggregatorFactory
       public void fold(final ColumnValueSelector selector)
       {
         final DoublesSketch sketch = (DoublesSketch) selector.getObject();
-        union.union(sketch);
+        if (sketch != null) {
+          union.union(sketch);
+        }
       }
 
       @Nullable

extensions-core/datasketches/src/test/java/org/apache/druid/query/aggregation/datasketches/quantiles/DoublesSketchAggregatorFactoryTest.java

@@ -25,8 +25,11 @@ import nl.jqno.equalsverifier.EqualsVerifier;
 import org.apache.druid.jackson.DefaultObjectMapper;
 import org.apache.druid.java.util.common.granularity.Granularities;
 import org.apache.druid.query.Druids;
+import org.apache.druid.query.aggregation.AggregateCombiner;
+import org.apache.druid.query.aggregation.Aggregator;
 import org.apache.druid.query.aggregation.AggregatorFactory;
 import org.apache.druid.query.aggregation.CountAggregatorFactory;
+import org.apache.druid.query.aggregation.TestDoubleColumnSelectorImpl;
 import org.apache.druid.query.aggregation.post.FieldAccessPostAggregator;
 import org.apache.druid.query.aggregation.post.FinalizingFieldAccessPostAggregator;
 import org.apache.druid.query.timeseries.TimeseriesQuery;
@@ -172,4 +175,24 @@ public class DoublesSketchAggregatorFactoryTest
     Assert.assertEquals(factory, factory.withName("myFactory"));
     Assert.assertEquals("newTest", factory.withName("newTest").getName());
   }
+
+  @Test
+  public void testNullSketches()
+  {
+    final DoublesSketchAggregatorFactory factory = new DoublesSketchAggregatorFactory(
+        "myFactory",
+        "myField",
+        1024,
+        1000L,
+        null
+    );
+    final double[] values = new double[]{1, 2, 3, 4, 5, 6};
+    final TestDoubleColumnSelectorImpl selector = new TestDoubleColumnSelectorImpl(values);
+    final Aggregator agg1 = new DoublesSketchBuildAggregator(selector, 8);
+    Assert.assertNotNull(factory.combine(null, agg1.get()));
+    Assert.assertNotNull(factory.combine(agg1.get(), null));
+    AggregateCombiner ac = factory.makeAggregateCombiner();
+    ac.fold(new TestDoublesSketchColumnValueSelector());
+    Assert.assertNotNull(ac.getObject());
+  }
 }

extensions-core/datasketches/src/test/java/org/apache/druid/query/aggregation/datasketches/quantiles/TestDoublesSketchColumnValueSelector.java

@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.query.aggregation.datasketches.quantiles;
+
+import org.apache.datasketches.quantiles.DoublesSketch;
+import org.apache.druid.query.monomorphicprocessing.RuntimeShapeInspector;
+import org.apache.druid.segment.ColumnValueSelector;
+
+import javax.annotation.Nullable;
+
+public class TestDoublesSketchColumnValueSelector implements ColumnValueSelector<DoublesSketch>
+{
+  @Override
+  public void inspectRuntimeShape(RuntimeShapeInspector inspector)
+  {
+  }
+
+  @Override
+  public double getDouble()
+  {
+    return 99;
+  }
+
+  @Override
+  public float getFloat()
+  {
+    return 99;
+  }
+
+  @Override
+  public long getLong()
+  {
+    return 99;
+  }
+
+  @Override
+  public boolean isNull()
+  {
+    return false;
+  }
+
+  @Nullable
+  @Override
+  public DoublesSketch getObject()
+  {
+    return DoublesSketchOperations.EMPTY_SKETCH;
+  }
+
+  @Override
+  public Class<? extends DoublesSketch> classOfObject()
+  {
+    return DoublesSketchOperations.EMPTY_SKETCH.getClass();
+  }
+}