sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamBuiltinAnalyticFunctions.java

/*
 * 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.beam.sdk.extensions.sql.impl.transform;

import java.math.BigDecimal;
import java.util.Map;
import java.util.Optional;
import java.util.function.Function;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.transforms.Combine;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;

/** Built-in Analytic Functions for the aggregation analytics functionality. */
@SuppressWarnings({
  "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
public class BeamBuiltinAnalyticFunctions {
  public static final Map<String, Function<Schema.FieldType, Combine.CombineFn<?, ?, ?>>>
      BUILTIN_ANALYTIC_FACTORIES =
          ImmutableMap.<String, Function<Schema.FieldType, Combine.CombineFn<?, ?, ?>>>builder()
              // Aggregate Analytic Functions
              .putAll(BeamBuiltinAggregations.BUILTIN_AGGREGATOR_FACTORIES)
              // Navigation Functions
              .put("FIRST_VALUE", typeName -> navigationFirstValue())
              .put("LAST_VALUE", typeName -> navigationLastValue())
              // Numbering Functions
              .put("ROW_NUMBER", typeName -> numberingRowNumber())
              .put("DENSE_RANK", typeName -> numberingDenseRank())
              .put("RANK", typeName -> numberingRank())
              .put("PERCENT_RANK", typeName -> numberingPercentRank())
              .build();

  public static Combine.CombineFn<?, ?, ?> create(String functionName, Schema.FieldType fieldType) {
    Function<Schema.FieldType, Combine.CombineFn<?, ?, ?>> aggregatorFactory =
        BUILTIN_ANALYTIC_FACTORIES.get(functionName);
    if (aggregatorFactory != null) {
      return aggregatorFactory.apply(fieldType);
    }
    throw new UnsupportedOperationException(
        String.format("Analytics Function [%s] is not supported", functionName));
  }

  // Navigation functions
  public static <T> Combine.CombineFn<T, ?, T> navigationFirstValue() {
    return new FirstValueCombineFn();
  }

  public static <T> Combine.CombineFn<T, ?, T> navigationLastValue() {
    return new LastValueCombineFn();
  }

  private static class FirstValueCombineFn<T> extends Combine.CombineFn<T, Optional<T>, T> {
    private FirstValueCombineFn() {}

    @Override
    public Optional<T> createAccumulator() {
      return Optional.empty();
    }

    @Override
    public Optional<T> addInput(Optional<T> accumulator, T input) {
      Optional<T> r = accumulator;
      if (!accumulator.isPresent()) {
        r = Optional.of(input);
      }
      return r;
    }

    @Override
    public Optional<T> mergeAccumulators(Iterable<Optional<T>> accumulators) {
      throw new UnsupportedOperationException();
    }

    @Override
    public T extractOutput(Optional<T> accumulator) {
      return accumulator.isPresent() ? accumulator.get() : null;
    }
  }

  private static class LastValueCombineFn<T> extends Combine.CombineFn<T, Optional<T>, T> {
    private LastValueCombineFn() {}

    @Override
    public Optional<T> createAccumulator() {
      return Optional.empty();
    }

    @Override
    public Optional<T> addInput(Optional<T> accumulator, T input) {
      Optional<T> r = Optional.of(input);
      return r;
    }

    @Override
    public Optional<T> mergeAccumulators(Iterable<Optional<T>> accumulators) {
      throw new UnsupportedOperationException();
    }

    @Override
    public T extractOutput(Optional<T> accumulator) {
      return accumulator.isPresent() ? accumulator.get() : null;
    }
  }

  // Numbering functions
  public static <T> Combine.CombineFn<T, ?, T> numberingRowNumber() {
    return new RowNumberCombineFn();
  }

  public static <T> Combine.CombineFn<T, ?, T> numberingDenseRank() {
    return new DenseRankCombineFn();
  }

  public static <T> Combine.CombineFn<T, ?, T> numberingRank() {
    return new RankCombineFn();
  }

  public static <T> Combine.CombineFn<T, ?, T> numberingPercentRank() {
    return new PercentRankCombineFn();
  }

  public abstract static class PositionAwareCombineFn<InputT, AccumT, OutputT>
      extends Combine.CombineFn<InputT, AccumT, OutputT> {
    public abstract AccumT addInput(
        AccumT accumulator,
        InputT input,
        Long cursorWindow,
        Long cursorPartition,
        Long countPartition);

    @Override
    public AccumT addInput(AccumT mutableAccumulator, InputT input) {
      throw new UnsupportedOperationException();
    }

    @Override
    public AccumT mergeAccumulators(Iterable<AccumT> accumulators) {
      throw new UnsupportedOperationException();
    }
  }

  private static class RowNumberCombineFn<T>
      extends PositionAwareCombineFn<BigDecimal, Optional<Long>, Long> {

    @Override
    public Optional<Long> addInput(
        Optional<Long> accumulator,
        BigDecimal input,
        Long cursorPosition,
        Long cursorPartition,
        Long countPartition) {
      return Optional.of(cursorPartition);
    }

    @Override
    public Optional<Long> createAccumulator() {
      return Optional.empty();
    }

    @Override
    public Long extractOutput(Optional<Long> accumulator) {
      // 1-based result
      return accumulator.isPresent() ? accumulator.get() + 1 : null;
    }
  }

  private static class DenseRankCombineFn<T>
      extends PositionAwareCombineFn<BigDecimal, KV<BigDecimal, Long>, Long> {

    @Override
    public KV<BigDecimal, Long> addInput(
        KV<BigDecimal, Long> accumulator,
        BigDecimal input,
        Long cursorPosition,
        Long cursorPartition,
        Long countPartition) {
      KV<BigDecimal, Long> r = null;
      if (accumulator == null) {
        r = KV.of(input, 0L);
      } else {
        if (accumulator.getKey().compareTo(input) == 0) {
          r = KV.of(input, accumulator.getValue());
        } else {
          r = KV.of(input, accumulator.getValue() + 1);
        }
      }
      return r;
    }

    @Override
    public KV<BigDecimal, Long> createAccumulator() {
      return null;
    }

    @Override
    public Long extractOutput(KV<BigDecimal, Long> accumulator) {
      // 1-based result
      return accumulator != null ? accumulator.getValue() + 1 : null;
    }
  }

  private static class RankCombineFn<T>
      extends PositionAwareCombineFn<BigDecimal, KV<BigDecimal, Long>, Long> {

    @Override
    public KV<BigDecimal, Long> addInput(
        KV<BigDecimal, Long> accumulator,
        BigDecimal input,
        Long cursorPosition,
        Long cursorPartition,
        Long countPartition) {
      KV<BigDecimal, Long> r = null;
      if (accumulator == null) {
        r = KV.of(input, 0L);
      } else {
        if (accumulator.getKey().compareTo(input) == 0) {
          r = KV.of(input, accumulator.getValue());
        } else {
          r = KV.of(input, cursorPosition);
        }
      }
      return r;
    }

    @Override
    public KV<BigDecimal, Long> createAccumulator() {
      return null;
    }

    @Override
    public Long extractOutput(KV<BigDecimal, Long> accumulator) {
      // 1-based result
      return accumulator != null ? accumulator.getValue() + 1 : null;
    }
  }

  private static class PercentRankCombineFn<T>
      extends PositionAwareCombineFn<BigDecimal, KV<Optional<Long>, KV<BigDecimal, Long>>, Double> {

    RankCombineFn internalRank;

    PercentRankCombineFn() {
      internalRank = new RankCombineFn();
    }

    @Override
    public KV<Optional<Long>, KV<BigDecimal, Long>> addInput(
        KV<Optional<Long>, KV<BigDecimal, Long>> accumulator,
        BigDecimal input,
        Long cursorPosition,
        Long cursorPartition,
        Long countPartition) {
      KV<BigDecimal, Long> ac1 =
          internalRank.addInput(
              accumulator.getValue(), input, cursorPosition, cursorPartition, countPartition);
      Optional<Long> ac2 = Optional.of(countPartition);
      return KV.of(ac2, ac1);
    }

    @Override
    public KV<Optional<Long>, KV<BigDecimal, Long>> createAccumulator() {
      return KV.of(Optional.empty(), internalRank.createAccumulator());
    }

    @Override
    public Double extractOutput(KV<Optional<Long>, KV<BigDecimal, Long>> accumulator) {
      Long nr = accumulator.getKey().orElse(null);
      Long rk = internalRank.extractOutput(accumulator.getValue());
      Double r = 0.0;
      if (nr != null && rk != null && nr > 1L) {
        r = (rk.doubleValue() - 1) / (nr.doubleValue() - 1);
      }
      return r;
    }
  }
}