Beam SQL 扩展：用户定义函数

如果 Beam SQL 没有标量函数或聚合函数来满足您的需求，则可以在 Java 中编写这些函数，并在 SQL 查询中调用它们。这些函数通常被称为 UDF（用于标量函数）和 UDAF（用于聚合函数）。

创建和指定用户定义函数 (UDF)

UDF 可以是以下内容

任何接受零个或多个标量字段并返回一个标量值的 Java 方法。
SerializableFunction。

以下是一个 UDF 示例以及如何在 DSL 中使用它

/**
 * A example UDF for test.
 */
public static class CubicInteger implements BeamSqlUdf {
  public static Integer eval(Integer input){
    return input * input * input;
  }
}

/**
 * Another example UDF with {@link SerializableFunction}.
 */
public static class CubicIntegerFn implements SerializableFunction<Integer, Integer> {
  @Override
  public Integer apply(Integer input) {
    return input * input * input;
  }
}

// Define a SQL query which calls the above UDFs
String sql =
    "SELECT f_int, cubic1(f_int), cubic2(f_int)"
      + "FROM PCOLLECTION "
      + "WHERE f_int = 2";

// Create and apply the PTransform representing the query.
// Register the UDFs used in the query by calling '.registerUdf()' with
// either a class which implements BeamSqlUdf or with
// an instance of the SerializableFunction;
PCollection<Row> result =
    input.apply(
        "udfExample",
        SqlTransform
            .query(sql)
            .registerUdf("cubic1", CubicInteger.class)
            .registerUdf("cubic2", new CubicIntegerFn())

创建和指定用户定义聚合函数 (UDAF)

Beam SQL 可以接受 CombineFn 作为 UDAF。注册类似于上面的 UDF 示例

/**
 * UDAF(CombineFn) for test, which returns the sum of square.
 */
public static class SquareSum extends CombineFn<Integer, Integer, Integer> {
  @Override
  public Integer createAccumulator() {
    return 0;
  }

  @Override
  public Integer addInput(Integer accumulator, Integer input) {
    return accumulator + input * input;
  }

  @Override
  public Integer mergeAccumulators(Iterable<Integer> accumulators) {
    int v = 0;
    Iterator<Integer> ite = accumulators.iterator();
    while (ite.hasNext()) {
      v += ite.next();
    }
    return v;
  }

  @Override
  public Integer extractOutput(Integer accumulator) {
    return accumulator;
  }
}

// Define a SQL query which calls the above UDAF
String sql =
    "SELECT f_int1, squaresum(f_int2) "
      + "FROM PCOLLECTION "
      + "GROUP BY f_int1";

// Create and apply the PTransform representing the query.
// Register the UDAFs used in the query by calling '.registerUdaf()' by
// providing it an instance of the CombineFn
PCollection<Row> result =
    input.apply(
        "udafExample",
        SqlTransform
            .query(sql)
            .registerUdaf("squaresum", new SquareSum()));