如何在Apache Calcite自定义聚合函数中使用第二个参数？

要让你的UDAF支持AGG(column, 5)这种带常量参数的调用，需要调整UDAF类的结构，让它能接收并保存第二个参数，同时修改注册逻辑匹配函数签名。以下是具体步骤：

1. 修改UDAF类，支持接收常量参数 #

原来的UDAF类没有处理第二个参数的逻辑，需要添加构造函数接收参数，并自定义Accumulator来保存参数和聚合数据：

public class CalciteUDAF {
    // 保存第二个数字参数，由构造函数传入
    private final int threshold;

    // 构造函数，接收第二个常量参数
    public CalciteUDAF(int threshold) {
        this.threshold = threshold;
    }

    // 自定义Accumulator，同时存储聚合列数据和常量参数
    public static class Accumulator {
        List<Object> values;
        int threshold;

        public Accumulator(int threshold) {
            this.values = new ArrayList<>();
            this.threshold = threshold;
        }
    }

    // 初始化Accumulator时传入常量参数
    public Accumulator init() {
        return new Accumulator(this.threshold);
    }

    // 聚合阶段：添加当前行的列值到Accumulator
    public Accumulator add(Accumulator accumulator, Object columnValue) {
        accumulator.values.add(columnValue);
        return accumulator;
    }

    // 合并阶段（可选但建议实现，支持并行聚合）
    public Accumulator merge(Accumulator a, Accumulator b) {
        a.values.addAll(b.values);
        return a;
    }

    // 计算最终结果：示例为统计大于阈值的元素数量
    public Object result(Accumulator accumulator) {
        return accumulator.values.stream()
                .mapToInt(v -> ((Number) v).intValue())
                .filter(v -> v > accumulator.threshold)
                .count();
    }
}

2. 调整UDAF注册逻辑，匹配函数签名 #

原来的注册方式没有声明函数的参数数量和类型，Calcite无法正确传递第二个参数。需要用SqlAggFunction定义完整的函数签名后再注册：

// 定义函数的元信息：名称、参数类型、返回值等
SqlAggFunction aggFunc = new SqlAggFunction(
        "AGG",
        null,
        SqlKind.OTHER_FUNCTION,
        // 自定义返回值推断逻辑，示例用默认实现
        SqlReturnTypeInference.DEFAULT,
        SqlOperandTypeInference.DEFAULT,
        // 指定函数接收2个参数（任意类型列 + 整数）
        SqlOperandTypeChecker.FACTORY.any(2),
        // 绑定自定义UDAF实现类
        AggregateFunctionImpl.create(CalciteUDAF.class),
        false, false, false
);

// 注册到Schema
schemaPlus.add("AGG", aggFunc);

关键说明 #

• Calcite会在解析SQL时，将第二个常量参数（如示例中的5）传递给UDAF的构造函数，因此UDAF类必须提供对应参数类型的构造函数。
• 自定义Accumulator是为了在聚合的各个阶段（add、merge、result）都能访问到常量参数，避免依赖UDAF实例的字段（Accumulator可能在分布式场景下序列化传输）。
• 如果第二个参数是其他数字类型（如long、double），只需修改构造函数和Accumulator中的参数类型即可。

内容的提问来源于stack exchange，提问作者Mr.zs