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.schemas.transforms;
  
  import com.google.auto.value.AutoValue;
  import java.io.Serializable;
  import java.util.List;
  import java.util.stream.Collectors;
  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.schemas.FieldAccessDescriptor;
  import org.apache.beam.sdk.schemas.FieldTypeDescriptors;
  import org.apache.beam.sdk.schemas.Schema;
  import org.apache.beam.sdk.schemas.Schema.Field;
  import org.apache.beam.sdk.schemas.Schema.FieldType;
  import org.apache.beam.sdk.schemas.SchemaCoder;
  import org.apache.beam.sdk.schemas.utils.RowSelector;
  import org.apache.beam.sdk.schemas.utils.SelectHelpers;
  import org.apache.beam.sdk.schemas.utils.SelectHelpers.RowSelectorContainer;
  import org.apache.beam.sdk.transforms.Combine.CombineFn;
  import org.apache.beam.sdk.transforms.CombineFns;
  import org.apache.beam.sdk.transforms.CombineFns.CoCombineResult;
  import org.apache.beam.sdk.transforms.CombineFns.ComposedCombineFn;
  import org.apache.beam.sdk.transforms.SimpleFunction;
  import org.apache.beam.sdk.values.Row;
  import org.apache.beam.sdk.values.TupleTag;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
  import org.checkerframework.checker.nullness.qual.Nullable;
  
  /** This is the builder used by {@link Group} to build up a composed {@link CombineFn}. */
  @SuppressWarnings({
    "nullness", // TODO(https://github.com/apache/beam/issues/20497)
    "rawtypes"
  })
  class SchemaAggregateFn {
    static Inner create() {
      return new AutoValue_SchemaAggregateFn_Inner.Builder()
          .setFieldAggregations(Lists.newArrayList())
          .build();
    }
  
    /** Implementation of {@link #create}. */
    @AutoValue
    @AutoValue.CopyAnnotations
    @SuppressWarnings("rawtypes")
    abstract static class Inner extends CombineFn<Row, Object[], Row> {
      // Represents an aggregation of one or more fields.
      static class FieldAggregation<FieldT, AccumT, OutputT> implements Serializable {
        FieldAccessDescriptor fieldsToAggregate;
        private final boolean aggregateBaseValues;
        // The specification of the output field.
        private final Field outputField;
        // The combine function.
        private final CombineFn<FieldT, AccumT, OutputT> fn;
        // The TupleTag identifying this aggregation element in the composed combine fn.
        private final TupleTag<Object> combineTag;
        // The schema corresponding to the subset of input fields being aggregated.
        private final @Nullable Schema inputSubSchema;
        private final @Nullable FieldAccessDescriptor flattenedFieldAccessDescriptor;
        // The flattened version of inputSubSchema.
        private final @Nullable Schema flattenedInputSubSchema;
        // The output schema resulting from the aggregation.
        private final Schema aggregationSchema;
        private final boolean needsFlattening;
  
        FieldAggregation(
            FieldAccessDescriptor fieldsToAggregate,
            boolean aggregateBaseValues,
            Field outputField,
            CombineFn<FieldT, AccumT, OutputT> fn,
            TupleTag<Object> combineTag) {
          this(
              fieldsToAggregate,
              aggregateBaseValues,
              outputField,
              fn,
              combineTag,
              Schema.builder().addField(outputField).build(),
              null);
        }
  
        FieldAggregation(
            FieldAccessDescriptor fieldsToAggregate,
            boolean aggregateBaseValues,
            Field outputField,
            CombineFn<FieldT, AccumT, OutputT> fn,
            TupleTag<Object> combineTag,
            Schema aggregationSchema,
            @Nullable Schema inputSchema) {
          this.aggregateBaseValues = aggregateBaseValues;
          if (inputSchema != null) {
            this.fieldsToAggregate = fieldsToAggregate.resolve(inputSchema);
            if (aggregateBaseValues) {
              Preconditions.checkArgument(fieldsToAggregate.referencesSingleField());
            }
            this.inputSubSchema = SelectHelpers.getOutputSchema(inputSchema, this.fieldsToAggregate);
            this.flattenedFieldAccessDescriptor =
                SelectHelpers.allLeavesDescriptor(inputSubSchema, SelectHelpers.CONCAT_FIELD_NAMES);
            this.flattenedInputSubSchema =
                SelectHelpers.getOutputSchema(inputSubSchema, flattenedFieldAccessDescriptor);
            this.needsFlattening = !inputSchema.equals(flattenedInputSubSchema);
          } else {
            this.fieldsToAggregate = fieldsToAggregate;
            this.inputSubSchema = null;
            this.flattenedFieldAccessDescriptor = null;
            this.flattenedInputSubSchema = null;
            this.needsFlattening = false;
          }
          this.outputField = outputField;
          this.fn = fn;
          this.combineTag = combineTag;
          this.aggregationSchema = aggregationSchema;
        }
  
        // The Schema is not necessarily known when the SchemaAggregateFn is created. Once the schema
        // is known, resolve will be called with the proper schema.
        FieldAggregation<FieldT, AccumT, OutputT> resolve(Schema schema) {
          return new FieldAggregation<>(
              fieldsToAggregate,
              aggregateBaseValues,
              outputField,
              fn,
              combineTag,
              aggregationSchema,
              schema);
        }
      }
  
      abstract Builder toBuilder();
  
      @AutoValue.Builder
      @AutoValue.CopyAnnotations
      @SuppressWarnings("rawtypes")
      abstract static class Builder {
        abstract Builder setInputSchema(@Nullable Schema inputSchema);
  
        abstract Builder setOutputSchema(@Nullable Schema outputSchema);
  
        abstract Builder setComposedCombineFn(@Nullable ComposedCombineFn composedCombineFn);
  
        abstract Builder setFieldAggregations(List<FieldAggregation> fieldAggregations);
  
        abstract Inner build();
      }
  
      abstract @Nullable Schema getInputSchema();
  
      abstract @Nullable Schema getOutputSchema();
  
      abstract @Nullable ComposedCombineFn getComposedCombineFn();
  
      abstract List<FieldAggregation> getFieldAggregations();
  
      /** Once the schema is known, this function is called by the {@link Group} transform. */
      Inner withSchema(Schema inputSchema) {
        List<FieldAggregation> fieldAggregations =
            getFieldAggregations().stream()
                .map(f -> f.resolve(inputSchema))
                .collect(Collectors.toList());
  
        ComposedCombineFn composedCombineFn = null;
        for (int i = 0; i < fieldAggregations.size(); ++i) {
          FieldAggregation fieldAggregation = fieldAggregations.get(i);
          SimpleFunction<Row, ?> extractFunction;
          Coder extractOutputCoder;
          if (fieldAggregation.fieldsToAggregate.referencesSingleField()) {
            extractFunction =
                new ExtractSingleFieldFunction(
                    inputSchema, fieldAggregation.aggregateBaseValues, fieldAggregation);
  
            FieldType fieldType = fieldAggregation.flattenedInputSubSchema.getField(0).getType();
            if (fieldAggregation.aggregateBaseValues) {
              while (fieldType.getTypeName().isLogicalType()) {
                fieldType = fieldType.getLogicalType().getBaseType();
              }
            }
            extractOutputCoder = SchemaCoder.coderForFieldType(fieldType);
          } else {
            extractFunction = new ExtractFieldsFunction(inputSchema, fieldAggregation);
            extractOutputCoder = SchemaCoder.of(fieldAggregation.inputSubSchema);
          }
          if (i == 0) {
            composedCombineFn =
                CombineFns.compose()
                    .with(
                        extractFunction,
                        extractOutputCoder,
                        fieldAggregation.fn,
                        fieldAggregation.combineTag);
          } else {
            composedCombineFn =
                ((ComposedCombineFn) composedCombineFn)
                    .with(
                        extractFunction,
                        extractOutputCoder,
                        fieldAggregation.fn,
                        fieldAggregation.combineTag);
          }
        }
  
        return toBuilder()
            .setInputSchema(inputSchema)
            .setComposedCombineFn(composedCombineFn)
            .setFieldAggregations(fieldAggregations)
            .build();
      }
  
      /** Aggregate all values of a set of fields into an output field. */
      <CombineInputT, AccumT, CombineOutputT> Inner aggregateFields(
          FieldAccessDescriptor fieldsToAggregate,
          boolean aggregateBaseValues,
          CombineFn<CombineInputT, AccumT, CombineOutputT> fn,
          String outputFieldName) {
        return aggregateFields(
            fieldsToAggregate,
            aggregateBaseValues,
            fn,
            Field.of(outputFieldName, FieldTypeDescriptors.fieldTypeForJavaType(fn.getOutputType())));
      }
  
      /** Aggregate all values of a set of fields into an output field. */
      <CombineInputT, AccumT, CombineOutputT> Inner aggregateFields(
          FieldAccessDescriptor fieldsToAggregate,
          boolean aggregateBaseValues,
          CombineFn<CombineInputT, AccumT, CombineOutputT> fn,
          Field outputField) {
        List<FieldAggregation> fieldAggregations = getFieldAggregations();
        TupleTag<Object> combineTag = new TupleTag<>(Integer.toString(fieldAggregations.size()));
        FieldAggregation fieldAggregation =
            new FieldAggregation<>(
                fieldsToAggregate, aggregateBaseValues, outputField, fn, combineTag);
        fieldAggregations.add(fieldAggregation);
  
        return toBuilder()
            .setOutputSchema(getOutputSchema(fieldAggregations))
            .setFieldAggregations(fieldAggregations)
            .build();
      }
  
      private Schema getOutputSchema(List<FieldAggregation> fieldAggregations) {
        Schema.Builder outputSchema = Schema.builder();
        for (FieldAggregation aggregation : fieldAggregations) {
          outputSchema.addField(aggregation.outputField);
        }
        return outputSchema.build();
      }
  
      /** Extract a single field from an input {@link Row}. */
      private static class ExtractSingleFieldFunction<OutputT> extends SimpleFunction<Row, OutputT> {
        private final RowSelector rowSelector;
        private final boolean extractBaseValue;
        private final @Nullable RowSelector flatteningSelector;
        private final FieldAggregation fieldAggregation;
  
        private ExtractSingleFieldFunction(
            Schema inputSchema, boolean extractBaseValue, FieldAggregation fieldAggregation) {
          rowSelector =
              new RowSelectorContainer(inputSchema, fieldAggregation.fieldsToAggregate, true);
          this.extractBaseValue = extractBaseValue;
          flatteningSelector =
              fieldAggregation.needsFlattening
                  ? new RowSelectorContainer(
                      fieldAggregation.inputSubSchema,
                      fieldAggregation.flattenedFieldAccessDescriptor,
                      true)
                  : null;
          this.fieldAggregation = fieldAggregation;
        }
  
        @Override
        public OutputT apply(Row row) {
          Row selected = rowSelector.select(row);
          if (fieldAggregation.needsFlattening) {
            selected = flatteningSelector.select(selected);
          }
          if (extractBaseValue
              && selected.getSchema().getField(0).getType().getTypeName().isLogicalType()) {
            return (OutputT) selected.getBaseValue(0, Object.class);
          }
          return selected.getValue(0);
        }
      }
  
      /** Extract multiple fields from an input {@link Row}. */
      private static class ExtractFieldsFunction extends SimpleFunction<Row, Row> {
        private final RowSelector rowSelector;
  
        private ExtractFieldsFunction(Schema inputSchema, FieldAggregation fieldAggregation) {
          rowSelector =
              new RowSelectorContainer(inputSchema, fieldAggregation.fieldsToAggregate, true);
        }
  
        @Override
        public Row apply(Row row) {
          return rowSelector.select(row);
        }
      }
  
      @Override
      public Object[] createAccumulator() {
        return getComposedCombineFn().createAccumulator();
      }
  
      @Override
      public Object[] addInput(Object[] accumulator, Row input) {
        return getComposedCombineFn().addInput(accumulator, input);
      }
  
      @Override
      public Object[] mergeAccumulators(Iterable<Object[]> accumulator) {
        return getComposedCombineFn().mergeAccumulators(accumulator);
      }
  
      @Override
      public Coder<Object[]> getAccumulatorCoder(CoderRegistry registry, Coder<Row> inputCoder)
          throws CannotProvideCoderException {
        return getComposedCombineFn().getAccumulatorCoder(registry, inputCoder);
      }
  
      @Override
      public Coder<Row> getDefaultOutputCoder(CoderRegistry registry, Coder<Row> inputCoder) {
        return SchemaCoder.of(getOutputSchema());
      }
  
      @Override
      public Row extractOutput(Object[] accumulator) {
        // Build a row containing a field for every aggregate that was registered.
        CoCombineResult coCombineResult = getComposedCombineFn().extractOutput(accumulator);
        Row.Builder output = Row.withSchema(getOutputSchema());
        for (FieldAggregation fieldAggregation : getFieldAggregations()) {
          Object aggregate = coCombineResult.get(fieldAggregation.combineTag);
          output.addValue(aggregate);
        }
        return output.build();
      }
    }
  }
  