DRILL-8434: Add Median Function (#2803)

contrib/udfs/src/main/java/org/apache/drill/exec/udfs/DistributionFunctions.java

@@ -29,6 +29,7 @@
 import org.apache.drill.exec.expr.holders.Float8Holder;
 import org.apache.drill.exec.expr.holders.IntHolder;
 
+@SuppressWarnings("unused")
 public class DistributionFunctions {
 
   @FunctionTemplate(names = {"width_bucket", "widthBucket"},
@@ -150,7 +151,6 @@ public void output() {
       tau.value = result;
     }
   }
-
   @FunctionTemplate(names = {"regr_slope", "regrSlope"},
       scope = FunctionScope.POINT_AGGREGATE,
       nulls = NullHandling.INTERNAL)

contrib/udfs/src/test/resources/test_data.csvh

@@ -1,5 +1,5 @@
-col1,col2
-2,25
-3,32
-4,49
-5,32
+col1,col2,col3,col4
+2,25,2.0,0
+3,32,3.0,0
+4,49,4.5,
+5,32,9,9

exec/java-exec/src/main/codegen/config.fmpp

@@ -36,6 +36,7 @@ data: {
     decimalaggrtypes3:        tdd(../data/DecimalAggrTypes3.tdd),
     aggrtypes2:               tdd(../data/AggrTypes2.tdd),
     aggrtypes3:               tdd(../data/AggrTypes3.tdd),
+    aggrtypes4:               tdd(../data/AggrTypes4.tdd),
     covarTypes:               tdd(../data/CovarTypes.tdd),
     corrTypes:                tdd(../data/CorrelationTypes.tdd),
     logicalTypes:             tdd(../data/AggrBitwiseLogicalTypes.tdd),

exec/java-exec/src/main/codegen/data/AggrTypes4.tdd

@@ -0,0 +1,33 @@
+#
+# 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.
+#
+
+{
+  aggrtypes: [
+    {className: "Median", funcName: "median",  aliasName: "", types: [
+      {inputType: "Int", outputType: "NullableBigInt", major: "Numeric", java: "Long", medianHelper: "StreamingIntMedianHelper"},
+      {inputType: "BigInt", outputType: "NullableBigInt",  major: "Numeric", java: "Long", medianHelper: "StreamingIntMedianHelper"},
+      {inputType: "NullableInt", outputType: "NullableBigInt", major: "Numeric", java: "Long", medianHelper: "StreamingIntMedianHelper"},
+      {inputType: "NullableBigInt", outputType: "NullableBigInt", major: "Numeric", java: "Long", medianHelper: "StreamingIntMedianHelper"},
+      {inputType: "Float4", outputType: "NullableFloat8", major: "Numeric", java: "Double", medianHelper: "StreamingDoubleMedianHelper"},
+      {inputType: "Float8", outputType: "NullableFloat8", major: "Numeric", java: "Double", medianHelper: "StreamingDoubleMedianHelper"},
+      {inputType: "NullableFloat4", outputType: "NullableFloat8", major: "Numeric", java: "Double", medianHelper: "StreamingDoubleMedianHelper"},
+      {inputType: "NullableFloat8", outputType: "NullableFloat8", major: "Numeric", java: "Double", medianHelper: "StreamingDoubleMedianHelper"},
+     ]
+   }
+  ]
+}

exec/java-exec/src/main/codegen/templates/MedianFunctions.java

@@ -0,0 +1,106 @@
+/*
+ * 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.
+ */
+<@pp.dropOutputFile />
+
+
+
+<#list aggrtypes4.aggrtypes as aggrtype>
+<@pp.changeOutputFile name="/org/apache/drill/exec/expr/fn/impl/gaggr/${aggrtype.className}Functions.java" />
+
+<#include "/@includes/license.ftl" />
+
+/*
+ * This class is generated using freemarker and the ${.template_name} template.
+ */
+
+<#-- A utility class that is used to generate java code for median functions -->
+
+package org.apache.drill.exec.expr.fn.impl.gaggr;
+
+import org.apache.drill.exec.expr.DrillAggFunc;
+import org.apache.drill.exec.expr.annotations.FunctionTemplate;
+import org.apache.drill.exec.expr.annotations.FunctionTemplate.NullHandling;
+import org.apache.drill.exec.expr.annotations.FunctionTemplate.FunctionScope;
+import org.apache.drill.exec.expr.annotations.Output;
+import org.apache.drill.exec.expr.annotations.Param;
+import org.apache.drill.exec.expr.annotations.Workspace;
+import org.apache.drill.exec.expr.holders.*;
+import org.apache.drill.exec.expr.fn.impl.StreamingMedianHelpers.StreamingMedianHelper;
+import org.apache.drill.exec.expr.fn.impl.StreamingMedianHelpers.StreamingIntMedianHelper;
+import org.apache.drill.exec.expr.fn.impl.StreamingMedianHelpers.StreamingDoubleMedianHelper;
+
+@SuppressWarnings("unused")
+
+public class ${aggrtype.className}Functions {
+  static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(${aggrtype.className}Functions.class);
+
+<#list aggrtype.types as type>
+
+<#if aggrtype.aliasName == "">
+  @FunctionTemplate(name = "${aggrtype.funcName}", scope = FunctionTemplate.FunctionScope.POINT_AGGREGATE)
+<#else>
+  @FunctionTemplate(names = {"${aggrtype.funcName}", "${aggrtype.aliasName}"}, scope = FunctionTemplate.FunctionScope.POINT_AGGREGATE)
+</#if>
+
+  public static class ${type.inputType}${aggrtype.className} implements DrillAggFunc {
+
+    @Param ${type.inputType}Holder input;
+    @Output ${type.outputType}Holder median;
+    @Workspace ObjectHolder utils;
+
+    public void setup() {
+      // Initialize the ObjectHolder
+      utils = new ObjectHolder();
+      utils.obj = new org.apache.drill.exec.expr.fn.impl.StreamingMedianHelpers.${type.medianHelper}();
+    }
+
+    @Override
+    public void add() {
+      org.apache.drill.exec.expr.fn.impl.StreamingMedianHelpers.${type.medianHelper} medianHelper = (org.apache.drill.exec.expr.fn.impl.StreamingMedianHelpers.${type.medianHelper}) utils.obj;
+      <#if type.inputType?starts_with("Nullable")>
+          sout: {
+      if (input.isSet == 0) {
+        // processing nullable input and the value is null, so don't do anything...
+        break sout;
+      }
+	    </#if>
+      medianHelper.addNextNumber(input.value);
+
+      <#if type.inputType?starts_with("Nullable")>
+      } // end of sout block
+    </#if>
+    }
+
+    @Override
+    public void output() {
+      org.apache.drill.exec.expr.fn.impl.StreamingMedianHelpers.${type.medianHelper} medianHelper = (org.apache.drill.exec.expr.fn.impl.StreamingMedianHelpers.${type.medianHelper}) utils.obj;
+      median.value = medianHelper.getMedian();
+      median.isSet = 1;
+    }
+
+    @Override
+    public void reset() {
+      org.apache.drill.exec.expr.fn.impl.StreamingMedianHelpers.StreamingMedianHelper medianHelper = (org.apache.drill.exec.expr.fn.impl.StreamingMedianHelpers.${type.medianHelper}) utils.obj;
+      medianHelper.reset();
+    }
+  }
+
+
+</#list>
+}
+</#list>

exec/java-exec/src/main/java/org/apache/drill/exec/expr/fn/impl/StreamingMedianHelpers.java

@@ -0,0 +1,133 @@
+/*
+ * 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.drill.exec.expr.fn.impl;
+
+import java.util.Comparator;
+import java.util.PriorityQueue;
+
+/**
+ * This class implements a heap-based streaming median.
+ *
+ *<p>
+ * Reference: <a href="https://www.baeldung.com/java-stream-integers-median-using-heap">Stream Integers Median using Heap</a>
+ * </p>
+ */
+
+public class StreamingMedianHelpers {
+
+  public interface StreamingMedianHelper {
+    void reset();
+  }
+
+  public static class StreamingIntMedianHelper implements StreamingMedianHelper {
+    private final PriorityQueue<Long> minHeap;
+    private final PriorityQueue<Long> maxHeap;
+
+    public StreamingIntMedianHelper() {
+      super();
+      this.minHeap = new PriorityQueue<>();
+      this.maxHeap = new PriorityQueue<>(Comparator.reverseOrder());
+    }
+
+    public void addNextNumber(Long n) {
+      if (!minHeap.isEmpty() && n < minHeap.peek()) {
+        maxHeap.offer(n);
+        if (maxHeap.size() > minHeap.size() + 1) {
+          minHeap.offer(maxHeap.poll());
+        }
+      } else {
+        minHeap.offer(n);
+        if (minHeap.size() > maxHeap.size() + 1) {
+          maxHeap.offer(minHeap.poll());
+        }
+      }
+    }
+
+    public void addNextNumber(int n) {
+      addNextNumber((long) n);
+    }
+
+    public Long getMedian() {
+      Long median;
+      if (minHeap.size() < maxHeap.size()) {
+        median = maxHeap.peek();
+      } else if (minHeap.size() > maxHeap.size()) {
+        median = minHeap.peek();
+      } else if (minHeap.isEmpty() && maxHeap.isEmpty()) {
+        median = 0L;
+      } else {
+        median = (minHeap.peek() + maxHeap.peek()) / 2;
+      }
+      return median;
+    }
+
+    public void reset() {
+      minHeap.clear();
+      maxHeap.clear();
+    }
+  }
+
+  public static class StreamingDoubleMedianHelper implements StreamingMedianHelper {
+    private final PriorityQueue<Double> minHeap;
+    private final PriorityQueue<Double> maxHeap;
+
+    public StreamingDoubleMedianHelper() {
+      super();
+      this.minHeap = new PriorityQueue<>();
+      this.maxHeap = new PriorityQueue<>(Comparator.reverseOrder());
+    }
+
+    public void addNextNumber(Double n) {
+      if (!minHeap.isEmpty() && n < minHeap.peek()) {
+        maxHeap.offer(n);
+        if (maxHeap.size() > minHeap.size() + 1) {
+          minHeap.offer(maxHeap.poll());
+        }
+      } else {
+        minHeap.offer(n);
+        if (minHeap.size() > maxHeap.size() + 1) {
+          maxHeap.offer(minHeap.poll());
+        }
+      }
+    }
+
+    public void addNextNumber(float n) {
+      addNextNumber((double) n);
+    }
+
+    public Double getMedian() {
+      Double median;
+      if (minHeap.size() < maxHeap.size()) {
+        median = maxHeap.peek();
+      } else if (minHeap.size() > maxHeap.size()) {
+        median = minHeap.peek();
+      } else if (minHeap.isEmpty() && maxHeap.isEmpty()) {
+        median = 0.0;
+      } else {
+        median = (minHeap.peek() + maxHeap.peek()) / 2;
+      }
+      return median;
+    }
+
+    public void reset() {
+      minHeap.clear();
+      maxHeap.clear();
+    }
+  }
+}

exec/java-exec/src/test/java/org/apache/drill/exec/fn/impl/TestAggregateFunctions.java

@@ -1289,4 +1289,46 @@ public void testAggregateWithPivot() throws Exception {
       .baselineValues("Partial High School", 3555.8064516129034, 101.14516129032258, 3469.7014925373132, 103.3731343283582)
       .go();
   }
+
+  @Test
+  public void testMedian() throws Exception {
+
+    // Test as String from CSV and aggregate
+    String query = "SELECT category, median(`value`) AS med FROM cp.`orgs.csvh` GROUP BY category";
+    testBuilder()
+        .sqlQuery(query)
+        .unOrdered()
+        .baselineColumns("category", "med")
+        .baselineValues("cat1", 5L)
+        .baselineValues("cat2", 5L)
+        .go();
+
+    query = "SELECT category, median(CAST(`value` AS DOUBLE)) AS med FROM cp.`orgs.csvh` GROUP BY category";
+    testBuilder()
+        .sqlQuery(query)
+        .unOrdered()
+        .baselineColumns("category", "med")
+        .baselineValues("cat1", 5.0)
+        .baselineValues("cat2", 5.0)
+        .go();
+
+  }
+
+  @Test
+  public void testMedianWithPivot() throws Exception {
+    String query = "SELECT * FROM (\n" +
+        "SELECT education_level, salary, marital_status, extract(year from age('2023-02-23', birth_date)) age\n" +
+        "FROM cp.`employee.json`)\n" +
+        "PIVOT (median(salary) med_salary, median(age) med_age FOR marital_status IN ('M' married, 'S' single))";
+    testBuilder()
+        .sqlQuery(query)
+        .unOrdered()
+        .baselineColumns("education_level", "married_med_salary", "married_med_age", "single_med_salary", "single_med_age")
+        .baselineValues("Graduate Degree", 4400.0, 109L, 20.0, 109L)
+        .baselineValues("Bachelors Degree", 4400.0,109L, 4400.0, 109L)
+        .baselineValues("Partial College", 4500.0, 109L, 4400.0, 109L)
+        .baselineValues("High School Degree", 3700.0, 109L, 4400.0, 109L)
+        .baselineValues("Partial High School", 4400.0, 109L, 20.0, 109L)
+        .go();
+  }
 }

exec/java-exec/src/test/resources/orgs.csvh

@@ -0,0 +1,11 @@
+category,value
+cat1,5
+cat1,3
+cat1,2
+cat1,7
+cat1,10
+cat2,1
+cat2,5
+cat2,9
+cat2,2
+cat2,11

exec/java-exec/src/test/resources/test_data.csvh

@@ -0,0 +1,5 @@
+col1,col2,col3,col4
+2,25,2.0,0
+3,32,3.0,0
+4,49,4.5,
+5,32,9,9