User defined aggregation functions

io.lenses:lenses-sql-udf:4.0.0

public interface UserDefinedAggregateFunction extends UserDefinedFunction {

    String name();
    /**
     * Defines a mapping from input to output types.
     *
     * @param argType The data type of the single argument.
     * @return The data type of the result.
     * @throws UdfException if the given data type is not supported by this UDF.
     */
    DataType typer(DataType argType) throws UdfException;

    /**
     * The empty/initial aggregate value of this function.
     *
     * @return the initial value.
     */
    Value empty();

    /**
     * This method can be overridden by users to perform initialization logic.
     * It is guaranteed to be called before empty, add, merge, and in case of table-based aggregations, subtract.
     *
     * @param args The tail of the arguments passed to the UDAF.
     */
    default void init(Value[] args) {}

    /**
     * Handles adding a new value to the current aggregated value.
     *
     * @param aggregateKey The key of the field being aggregated on.
     * @param aggregatedValue The current aggregated value.
     * @param toBeAdded The value to be added to the current aggregated value.
     * @return The updated aggregated value.
     * @throws UdfException if adding the new value to the current aggregated value fails.
     */
    Value add(Value aggregateKey, Value aggregatedValue, Value toBeAdded) throws UdfException;

    /**
     * Merges two aggregated values in case of windowing.
     *
     * @param aggregateKey The key of the field being aggregated on.
     * @param first The first aggregated value.
     * @param second The second aggregated value.
     * @return The combined aggregated value.
     * @throws UdfException if the merge fails.
     */
    Value merge(Value aggregateKey, Value first, Value second) throws UdfException;

    /**
     * An optional final step that calculates and returns the result.
     *
     * @return An optional function that will be called with the aggregated value to calculate the final result.
     */
    default Optional<FinalStep> finalStep() {
        return Optional.empty();
    }

}

SELECT STREAM foo(bar1,bar2,bar3)
FROM source
GROUP BY foo.bar;

public interface UserDefinedTableAggregateFunction extends UserDefinedAggregateFunction {

    /**
     * Handles removal of a value from the aggregation.
     *
     * @param aggregateKey The key of the field being aggregated on.
     * @param aggregatedValue The current aggregated value.
     * @param toBeSubtracted The value to be subtracted from the current aggregated value.
     * @return The updated aggregated value.
     * @throws UdfException if the subtraction failed.
     */
    Value subtract(Value aggregateKey, Value aggregatedValue, Value toBeSubtracted) throws UdfException;

}

public class custom_average implements UserDefinedAggregateFunction {

    private Value calculateFinalResult(Value aggregatedValue) throws UdfException {
        long count = aggregatedValue.asRepeatedValue().get(0).toLongValue().get();
        long sum = aggregatedValue.asRepeatedValue().get(1).toLongValue().get();

        DoubleValue result;
        if (count == 0) {
            result = new DoubleValue(0d);
        } else {
            result = new DoubleValue(sum/count);
        }
        return result;
    }

    @Override
    public String name() {
        return "custom_average";
    }

    @Override
    public DataType typer(DataType argType) {
        return DataType.ltDouble();
    }

    @Override
    public Value empty() {
        LongValue count = new LongValue(0);
        LongValue product = new LongValue(0);
        return RepeatedValue.ofTwo(count, product);
    }

    @Override
    public Value add(Value aggregateKey, Value aggregatedValue, Value toBeAdded) throws UdfException {
        long count = aggregatedValue.asRepeatedValue().get(0).toLongValue().get();
        long product = aggregatedValue.asRepeatedValue().get(1).toLongValue().get();
        long added = toBeAdded.toLongValue().get();

        long updatedCount = count + 1;
        long updatedSum = product + added;

        return RepeatedValue.ofTwo(new LongValue(updatedCount), new LongValue(updatedSum));
    }

    @Override
    public Value merge(Value aggregateKey, Value first, Value second) throws UdfException {
        // sum up the counts and multiply the products
        long count1 = first.asRepeatedValue().get(0).toLongValue().get();
        long sum1 = first.asRepeatedValue().get(1).toLongValue().get();
        long count2 = second.asRepeatedValue().get(0).toLongValue().get();
        long sum2 = second.asRepeatedValue().get(1).toLongValue().get();

        long mergedCount = count1 + count2;
        long mergedProduct = sum1 + sum2;
        return RepeatedValue.ofTwo(new LongValue(mergedCount), new LongValue(mergedProduct));
    }

    @Override
    public Optional<FinalStep> finalStep() {
        return Optional.of(this::calculateFinalResult);
    }
}