sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamBuiltinAggregations.java - beam - Git at Google

 /*
  * 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;
     }
   }
 }
	/*
	* 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;
	}
	}
	}