package com.uber.athena.udf.common.udaf;

import com.uber.athena.athenax.vm.api.AthenaXAggregateFunction;
import java.util.ArrayList;
import java.util.Iterator;
import org.apache.commons.math3.stat.descriptive.rank.Percentile;

public class PercentileUDAF extends AthenaXAggregateFunction<Double, PercentileAccumulator> {
  public PercentileAccumulator createAccumulator() {
    return new PercentileAccumulator();
  }

  public Double getValue(PercentileAccumulator accumulator) {
    return getValue(accumulator, accumulator.percentile);
  }

  public Double getValue(PercentileAccumulator accumulator, int percentile) {
    double[] valuesArray = new double[accumulator.values.size()];
    for (int i = 0; i < accumulator.values.size(); i++) {
      valuesArray[i] = accumulator.values.get(i);
    }

    Percentile p = new Percentile(percentile);
    return p.evaluate(valuesArray);
  }

  // Processes the input values and update the provided accumulator instance
  public void accumulate(PercentileAccumulator accumulator, Double value, int percentile) {
    if (value == null) {
      return;
    }
    accumulator.values.add(value);
    accumulator.percentile = percentile;
  }

  // Retracts the input values from the accumulator instance
  public void retract(PercentileAccumulator accumulator, Double value) {
    accumulator.values.remove(value);
  }

  // Merges a group of accumulator instances into one accumulator instance
  public void merge(
      PercentileAccumulator accumulator, java.lang.Iterable<PercentileAccumulator> its) {
    Iterator<PercentileAccumulator> iterator = its.iterator();
    while (iterator.hasNext()) {
      PercentileAccumulator nextAccumulator = iterator.next();
      accumulator.values.addAll(nextAccumulator.values);
    }
  }

  public void resetAccumulator(PercentileAccumulator accumulator) {
    accumulator.values = new ArrayList<Double>();
  }
}

ListView

ArrayList

ArrayList