content: /*
   * 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.io.Serializable;
  import java.math.BigDecimal;
  import java.math.MathContext;
  import java.math.RoundingMode;
  import java.util.Map;
  import java.util.function.Function;
  import org.apache.beam.sdk.coders.BigDecimalCoder;
  import org.apache.beam.sdk.coders.BigEndianIntegerCoder;
  import org.apache.beam.sdk.coders.CannotProvideCoderException;
  import org.apache.beam.sdk.coders.Coder;
  import org.apache.beam.sdk.coders.CoderRegistry;
  import org.apache.beam.sdk.coders.KvCoder;
  import org.apache.beam.sdk.coders.NullableCoder;
  import org.apache.beam.sdk.extensions.sql.impl.transform.agg.CountIf;
  import org.apache.beam.sdk.extensions.sql.impl.transform.agg.CovarianceFn;
  import org.apache.beam.sdk.extensions.sql.impl.transform.agg.VarianceFn;
  import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils;
  import org.apache.beam.sdk.schemas.Schema;
  import org.apache.beam.sdk.schemas.Schema.FieldType;
  import org.apache.beam.sdk.schemas.Schema.TypeName;
  import org.apache.beam.sdk.transforms.Combine;
  import org.apache.beam.sdk.transforms.Combine.CombineFn;
  import org.apache.beam.sdk.transforms.Count;
  import org.apache.beam.sdk.transforms.Max;
  import org.apache.beam.sdk.transforms.Min;
  import org.apache.beam.sdk.transforms.Sample;
  import org.apache.beam.sdk.transforms.Sum;
  import org.apache.beam.sdk.values.KV;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Predicates;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
  import org.checkerframework.checker.nullness.qual.Nullable;
  
  /** Built-in aggregations functions for COUNT/MAX/MIN/SUM/AVG/VAR_POP/VAR_SAMP. */
  @SuppressWarnings({
    "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  public class BeamBuiltinAggregations {
  
    public static final Map<String, Function<Schema.FieldType, CombineFn<?, ?, ?>>>
        BUILTIN_AGGREGATOR_FACTORIES =
            ImmutableMap.<String, Function<Schema.FieldType, CombineFn<?, ?, ?>>>builder()
                .put("ANY_VALUE", typeName -> Sample.anyValueCombineFn())
                // Drop null elements for these aggregations.
                .put("COUNT", typeName -> new DropNullFnWithDefault(Count.combineFn()))
                .put("MAX", typeName -> new DropNullFn(BeamBuiltinAggregations.createMax(typeName)))
                .put("MIN", typeName -> new DropNullFn(BeamBuiltinAggregations.createMin(typeName)))
                .put("SUM", typeName -> new DropNullFn(BeamBuiltinAggregations.createSum(typeName)))
                .put(
                    "$SUM0",
                    typeName ->
                        new DropNullFnWithDefault(BeamBuiltinAggregations.createSum0(typeName)))
                .put("AVG", typeName -> new DropNullFn(BeamBuiltinAggregations.createAvg(typeName)))
                .put(
                    "BIT_OR",
                    typeName -> new DropNullFn(BeamBuiltinAggregations.createBitOr(typeName)))
                .put(
                    "BIT_XOR",
                    typeName -> new DropNullFn(BeamBuiltinAggregations.createBitXOr(typeName)))
                // JIRA link:https://issues.apache.org/jira/browse/BEAM-10379
                .put(
                    "BIT_AND",
                    typeName -> new DropNullFn(BeamBuiltinAggregations.createBitAnd(typeName)))
                .put("VAR_POP", t -> VarianceFn.newPopulation(t.getTypeName()))
                .put("VAR_SAMP", t -> VarianceFn.newSample(t.getTypeName()))
                .put("COVAR_POP", t -> CovarianceFn.newPopulation(t.getTypeName()))
                .put("COVAR_SAMP", t -> CovarianceFn.newSample(t.getTypeName()))
                .put("COUNTIF", typeName -> CountIf.combineFn())
                .build();
  
    private static MathContext mc = new MathContext(10, RoundingMode.HALF_UP);
  
    public static CombineFn<?, ?, ?> create(String functionName, Schema.FieldType fieldType) {
  
      Function<Schema.FieldType, CombineFn<?, ?, ?>> aggregatorFactory =
          BUILTIN_AGGREGATOR_FACTORIES.get(functionName);
  
      if (aggregatorFactory != null) {
        return aggregatorFactory.apply(fieldType);
      }
  
      throw new UnsupportedOperationException(
          String.format("Aggregator [%s] is not supported", functionName));
    }
  
    /** {@link CombineFn} for MAX based on {@link Max} and {@link Combine.BinaryCombineFn}. */
    static CombineFn createMax(FieldType fieldType) {
      if (CalciteUtils.isDateTimeType(fieldType)) {
        return new CustMax<>();
      }
      switch (fieldType.getTypeName()) {
        case BOOLEAN:
        case INT16:
        case BYTE:
        case FLOAT:
        case DATETIME:
        case DECIMAL:
        case STRING:
          return new CustMax<>();
        case INT32:
          return Max.ofIntegers();
        case INT64:
          return Max.ofLongs();
        case DOUBLE:
          return Max.ofDoubles();
        default:
          throw new UnsupportedOperationException(
              String.format("[%s] is not supported in MAX", fieldType));
      }
    }
  
    /** {@link CombineFn} for MIN based on {@link Min} and {@link Combine.BinaryCombineFn}. */
    static CombineFn createMin(Schema.FieldType fieldType) {
      if (CalciteUtils.isDateTimeType(fieldType)) {
        return new CustMin();
      }
      switch (fieldType.getTypeName()) {
        case BOOLEAN:
        case BYTE:
        case INT16:
        case FLOAT:
        case DATETIME:
        case DECIMAL:
        case STRING:
          return new CustMin();
        case INT32:
          return Min.ofIntegers();
        case INT64:
          return Min.ofLongs();
        case DOUBLE:
          return Min.ofDoubles();
        default:
          throw new UnsupportedOperationException(
              String.format("[%s] is not supported in MIN", fieldType));
      }
    }
  
    /** {@link CombineFn} for Sum based on {@link Sum} and {@link Combine.BinaryCombineFn}. */
    static CombineFn createSum(Schema.FieldType fieldType) {
      switch (fieldType.getTypeName()) {
        case INT32:
          return Sum.ofIntegers();
        case INT16:
          return new ShortSum();
        case BYTE:
          return new ByteSum();
        case INT64:
          return new LongSum();
        case FLOAT:
          return new FloatSum();
        case DOUBLE:
          return Sum.ofDoubles();
        case DECIMAL:
          return new BigDecimalSum();
        default:
          throw new UnsupportedOperationException(
              String.format("[%s] is not supported in SUM", fieldType));
      }
    }
  
    /**
     * {@link CombineFn} for Sum0 where sum of null returns 0 based on {@link Sum} and {@link
     * Combine.BinaryCombineFn}.
     */
    static CombineFn createSum0(Schema.FieldType fieldType) {
      switch (fieldType.getTypeName()) {
        case INT32:
          return new IntegerSum0();
        case INT16:
          return new ShortSum0();
        case BYTE:
          return new ByteSum0();
        case INT64:
          return new LongSum0();
        case FLOAT:
          return new FloatSum0();
        case DOUBLE:
          return new DoubleSum0();
        case DECIMAL:
          return new BigDecimalSum0();
        default:
          throw new UnsupportedOperationException(
              String.format("[%s] is not supported in SUM0", fieldType));
      }
    }
  
    /** {@link CombineFn} for AVG. */
    static CombineFn createAvg(Schema.FieldType fieldType) {
      switch (fieldType.getTypeName()) {
        case INT32:
          return new IntegerAvg();
        case INT16:
          return new ShortAvg();
        case BYTE:
          return new ByteAvg();
        case INT64:
          return new LongAvg();
        case FLOAT:
          return new FloatAvg();
        case DOUBLE:
          return new DoubleAvg();
        case DECIMAL:
          return new BigDecimalAvg();
        default:
          throw new UnsupportedOperationException(
              String.format("[%s] is not supported in AVG", fieldType));
      }
    }
  
    static CombineFn createBitOr(Schema.FieldType fieldType) {
      if (fieldType.getTypeName() == TypeName.INT64) {
        return new BitOr();
      }
      throw new UnsupportedOperationException(
          String.format("[%s] is not supported in BIT_OR", fieldType));
    }
  
    static CombineFn createBitAnd(Schema.FieldType fieldType) {
      if (fieldType.getTypeName() == TypeName.INT64) {
        return new BitAnd();
      }
      throw new UnsupportedOperationException(
          String.format("[%s] is not supported in BIT_AND", fieldType));
    }
  
    public static CombineFn createBitXOr(Schema.FieldType fieldType) {
      if (fieldType.getTypeName() == TypeName.INT64) {
        return new BitXOr();
      }
      throw new UnsupportedOperationException(
          String.format("[%s] is not supported in BIT_XOR", fieldType));
    }
  
    static class CustMax<T extends Comparable<T>> extends Combine.BinaryCombineFn<T> {
      @Override
      public T apply(T left, T right) {
        return (right == null || right.compareTo(left) < 0) ? left : right;
      }
    }
  
    static class CustMin<T extends Comparable<T>> extends Combine.BinaryCombineFn<T> {
      @Override
      public T apply(T left, T right) {
        return (left == null || left.compareTo(right) < 0) ? left : right;
      }
    }
  
    static class IntegerSum extends Combine.BinaryCombineFn<Integer> {
      @Override
      public Integer apply(Integer left, Integer right) {
        return (int) (left + right);
      }
    }
  
    static class ShortSum extends Combine.BinaryCombineFn<Short> {
      @Override
      public Short apply(Short left, Short right) {
        return (short) (left + right);
      }
    }
  
    static class ByteSum extends Combine.BinaryCombineFn<Byte> {
      @Override
      public Byte apply(Byte left, Byte right) {
        return (byte) (left + right);
      }
    }
  
    static class FloatSum extends Combine.BinaryCombineFn<Float> {
      @Override
      public Float apply(Float left, Float right) {
        return left + right;
      }
    }
  
    static class DoubleSum extends Combine.BinaryCombineFn<Double> {
      @Override
      public Double apply(Double left, Double right) {
        return (double) left + right;
      }
    }
  
    static class LongSum extends Combine.BinaryCombineFn<Long> {
      @Override
      public Long apply(Long left, Long right) {
        return Math.addExact(left, right);
      }
    }
  
    static class BigDecimalSum extends Combine.BinaryCombineFn<BigDecimal> {
      @Override
      public BigDecimal apply(BigDecimal left, BigDecimal right) {
        return left.add(right);
      }
    }
  
    static class IntegerSum0 extends IntegerSum {
      @Override
      public @Nullable Integer identity() {
        return 0;
      }
    }
  
    static class ShortSum0 extends ShortSum {
      @Override
      public @Nullable Short identity() {
        return 0;
      }
    }
  
    static class ByteSum0 extends ByteSum {
      @Override
      public @Nullable Byte identity() {
        return 0;
      }
    }
  
    static class FloatSum0 extends FloatSum {
      @Override
      public @Nullable Float identity() {
        return 0F;
      }
    }
  
    static class DoubleSum0 extends DoubleSum {
      @Override
      public @Nullable Double identity() {
        return 0D;
      }
    }
  
    static class LongSum0 extends LongSum {
      @Override
      public @Nullable Long identity() {
        return 0L;
      }
    }
  
    static class BigDecimalSum0 extends BigDecimalSum {
      @Override
      public @Nullable BigDecimal identity() {
        return BigDecimal.ZERO;
      }
    }
  
    private static class DropNullFn<InputT, AccumT, OutputT>
        extends CombineFn<InputT, AccumT, OutputT> {
      protected final CombineFn<InputT, AccumT, OutputT> combineFn;
  
      DropNullFn(CombineFn<InputT, AccumT, OutputT> combineFn) {
        this.combineFn = combineFn;
      }
  
      @Override
      public AccumT createAccumulator() {
        return null;
      }
  
      @Override
      public AccumT addInput(AccumT accumulator, InputT input) {
        if (input == null) {
          return accumulator;
        }
  
        if (accumulator == null) {
          accumulator = combineFn.createAccumulator();
        }
        return combineFn.addInput(accumulator, input);
      }
  
      @Override
      public AccumT mergeAccumulators(Iterable<AccumT> accumulators) {
        // filter out nulls
        accumulators = Iterables.filter(accumulators, Predicates.notNull());
  
        // handle only nulls
        if (!accumulators.iterator().hasNext()) {
          return null;
        }
  
        return combineFn.mergeAccumulators(accumulators);
      }
  
      @Override
      public OutputT extractOutput(AccumT accumulator) {
        if (accumulator == null) {
          return null;
        }
        return combineFn.extractOutput(accumulator);
      }
  
      @Override
      public Coder<AccumT> getAccumulatorCoder(CoderRegistry registry, Coder<InputT> inputCoder)
          throws CannotProvideCoderException {
        Coder<AccumT> coder = combineFn.getAccumulatorCoder(registry, inputCoder);
        if (coder instanceof NullableCoder) {
          return coder;
        }
        return NullableCoder.of(coder);
      }
    }
  
    private static class DropNullFnWithDefault<InputT, AccumT, OutputT>
        extends DropNullFn<InputT, AccumT, OutputT> {
  
      DropNullFnWithDefault(CombineFn<InputT, AccumT, OutputT> combineFn) {
        super(combineFn);
      }
  
      @Override
      public AccumT createAccumulator() {
        return combineFn.createAccumulator();
      }
    }
  
    /** {@link CombineFn} for <em>AVG</em> on {@link Number} types. */
    abstract static class Avg<T extends Number> extends CombineFn<T, KV<Integer, BigDecimal>, T> {
      @Override
      public KV<Integer, BigDecimal> createAccumulator() {
        return KV.of(0, BigDecimal.ZERO);
      }
  
      @Override
      public KV<Integer, BigDecimal> addInput(KV<Integer, BigDecimal> accumulator, T input) {
        return KV.of(accumulator.getKey() + 1, accumulator.getValue().add(toBigDecimal(input)));
      }
  
      @Override
      public KV<Integer, BigDecimal> mergeAccumulators(
          Iterable<KV<Integer, BigDecimal>> accumulators) {
        int size = 0;
        BigDecimal acc = BigDecimal.ZERO;
        for (KV<Integer, BigDecimal> ele : accumulators) {
          size += ele.getKey();
          acc = acc.add(ele.getValue());
        }
        return KV.of(size, acc);
      }
  
      @Override
      public Coder<KV<Integer, BigDecimal>> getAccumulatorCoder(
          CoderRegistry registry, Coder<T> inputCoder) {
        return KvCoder.of(BigEndianIntegerCoder.of(), BigDecimalCoder.of());
      }
  
      protected BigDecimal prepareOutput(KV<Integer, BigDecimal> accumulator) {
        return accumulator.getValue().divide(new BigDecimal(accumulator.getKey()), mc);
      }
  
      @Override
      public abstract T extractOutput(KV<Integer, BigDecimal> accumulator);
  
      public abstract BigDecimal toBigDecimal(T record);
    }
  
    static class IntegerAvg extends Avg<Integer> {
      @Override
      public @Nullable Integer extractOutput(KV<Integer, BigDecimal> accumulator) {
        return accumulator.getKey() == 0 ? null : prepareOutput(accumulator).intValue();
      }
  
      @Override
      public BigDecimal toBigDecimal(Integer record) {
        return new BigDecimal(record);
      }
    }
  
    static class LongAvg extends Avg<Long> {
      @Override
      public @Nullable Long extractOutput(KV<Integer, BigDecimal> accumulator) {
        return accumulator.getKey() == 0 ? null : prepareOutput(accumulator).longValue();
      }
  
      @Override
      public BigDecimal toBigDecimal(Long record) {
        return new BigDecimal(record);
      }
    }
  
    static class ShortAvg extends Avg<Short> {
      @Override
      public @Nullable Short extractOutput(KV<Integer, BigDecimal> accumulator) {
        return accumulator.getKey() == 0 ? null : prepareOutput(accumulator).shortValue();
      }
  
      @Override
      public BigDecimal toBigDecimal(Short record) {
        return new BigDecimal(record);
      }
    }
  
    static class ByteAvg extends Avg<Byte> {
      @Override
      public @Nullable Byte extractOutput(KV<Integer, BigDecimal> accumulator) {
        return accumulator.getKey() == 0 ? null : prepareOutput(accumulator).byteValue();
      }
  
      @Override
      public BigDecimal toBigDecimal(Byte record) {
        return new BigDecimal(record);
      }
    }
  
    static class FloatAvg extends Avg<Float> {
      @Override
      public @Nullable Float extractOutput(KV<Integer, BigDecimal> accumulator) {
        return accumulator.getKey() == 0 ? null : prepareOutput(accumulator).floatValue();
      }
  
      @Override
      public BigDecimal toBigDecimal(Float record) {
        return new BigDecimal(record);
      }
    }
  
    static class DoubleAvg extends Avg<Double> {
      @Override
      public @Nullable Double extractOutput(KV<Integer, BigDecimal> accumulator) {
        return accumulator.getKey() == 0 ? null : prepareOutput(accumulator).doubleValue();
      }
  
      @Override
      public BigDecimal toBigDecimal(Double record) {
        return new BigDecimal(record);
      }
    }
  
    static class BigDecimalAvg extends Avg<BigDecimal> {
      @Override
      public @Nullable BigDecimal extractOutput(KV<Integer, BigDecimal> accumulator) {
        return accumulator.getKey() == 0 ? null : prepareOutput(accumulator);
      }
  
      @Override
      public BigDecimal toBigDecimal(BigDecimal record) {
        return record;
      }
    }
  
    static class BitOr<T extends Number> extends CombineFn<T, BitOr.Accum, Long> {
      static class Accum implements Serializable {
        /** True if no inputs have been seen yet. */
        boolean isEmpty = true;
        /** The bitwise-or of the inputs seen so far. */
        long bitOr = 0L;
      }
  
      @Override
      public BitOr.Accum createAccumulator() {
        return new BitOr.Accum();
      }
  
      @Override
      public BitOr.Accum addInput(BitOr.Accum accum, T input) {
        accum.isEmpty = false;
        accum.bitOr |= input.longValue();
        return accum;
      }
  
      @Override
      public BitOr.Accum mergeAccumulators(Iterable<BitOr.Accum> accums) {
        BitOr.Accum merged = createAccumulator();
        for (BitOr.Accum accum : accums) {
          if (accum.isEmpty) {
            continue;
          }
          merged.isEmpty = false;
          merged.bitOr |= accum.bitOr;
        }
        return merged;
      }
  
      @Override
      public Long extractOutput(BitOr.Accum accum) {
        if (accum.isEmpty) {
          return null;
        }
        return accum.bitOr;
      }
    }
  
    /**
     * Bitwise AND function implementation.
     *
     * <p>Note: null values are ignored when mixed with non-null values.
     * (https://issues.apache.org/jira/browse/BEAM-10379)
     */
    static class BitAnd<T extends Number> extends CombineFn<T, BitAnd.Accum, Long> {
      static class Accum implements Serializable {
        /** True if no inputs have been seen yet. */
        boolean isEmpty = true;
        /** The bitwise-and of the inputs seen so far. */
        long bitAnd = -1L;
      }
  
      @Override
      public BitAnd.Accum createAccumulator() {
        return new BitAnd.Accum();
      }
  
      @Override
      public BitAnd.Accum addInput(BitAnd.Accum accum, T input) {
        accum.isEmpty = false;
        accum.bitAnd &= input.longValue();
        return accum;
      }
  
      @Override
      public BitAnd.Accum mergeAccumulators(Iterable<BitAnd.Accum> accums) {
        BitAnd.Accum merged = createAccumulator();
        for (BitAnd.Accum accum : accums) {
          if (accum.isEmpty) {
            continue;
          }
          merged.isEmpty = false;
          merged.bitAnd &= accum.bitAnd;
        }
        return merged;
      }
  
      @Override
      public Long extractOutput(BitAnd.Accum accum) {
        if (accum.isEmpty) {
          return null;
        }
        return accum.bitAnd;
      }
    }
  
    public static class BitXOr<T extends Number> extends CombineFn<T, BitXOr.Accum, Long> {
  
      static class Accum implements Serializable {
        /** True if no inputs have been seen yet. */
        boolean isEmpty = true;
        /** The bitwise-and of the inputs seen so far. */
        long bitXOr = 0L;
      }
  
      @Override
      public BitXOr.Accum createAccumulator() {
        return new BitXOr.Accum();
      }
  
      @Override
      public BitXOr.Accum addInput(Accum mutableAccumulator, T input) {
        if (input != null) {
          mutableAccumulator.isEmpty = false;
          mutableAccumulator.bitXOr ^= input.longValue();
        }
        return mutableAccumulator;
      }
  
      @Override
      public BitXOr.Accum mergeAccumulators(Iterable<BitXOr.Accum> accumulators) {
        BitXOr.Accum merged = createAccumulator();
        for (BitXOr.Accum accum : accumulators) {
          if (accum.isEmpty) {
            continue;
          }
          merged.isEmpty = false;
          merged.bitXOr ^= accum.bitXOr;
        }
        return merged;
      }
  
      @Override
      public Long extractOutput(BitXOr.Accum accumulator) {
        if (accumulator.isEmpty) {
          return null;
        }
        return accumulator.bitXOr;
      }
    }
  }
  