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.runners.samza.runtime;
  
  import java.util.Collection;
  import java.util.Collections;
  import org.apache.beam.runners.core.DoFnRunner;
  import org.apache.beam.runners.core.DoFnRunners;
  import org.apache.beam.runners.core.GroupAlsoByWindowViaWindowSetNewDoFn;
  import org.apache.beam.runners.core.KeyedWorkItem;
  import org.apache.beam.runners.core.KeyedWorkItemCoder;
  import org.apache.beam.runners.core.KeyedWorkItems;
  import org.apache.beam.runners.core.NullSideInputReader;
  import org.apache.beam.runners.core.StateInternals;
  import org.apache.beam.runners.core.StateInternalsFactory;
  import org.apache.beam.runners.core.StepContext;
  import org.apache.beam.runners.core.SystemReduceFn;
  import org.apache.beam.runners.core.TimerInternals;
  import org.apache.beam.runners.core.TimerInternals.TimerData;
  import org.apache.beam.runners.samza.SamzaExecutionContext;
  import org.apache.beam.runners.samza.SamzaPipelineOptions;
  import org.apache.beam.runners.samza.metrics.DoFnRunnerWithMetrics;
  import org.apache.beam.sdk.coders.Coder;
  import org.apache.beam.sdk.options.PipelineOptionsFactory;
  import org.apache.beam.sdk.transforms.DoFn;
  import org.apache.beam.sdk.transforms.DoFnSchemaInformation;
  import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
  import org.apache.beam.sdk.util.WindowedValueMultiReceiver;
  import org.apache.beam.sdk.values.KV;
  import org.apache.beam.sdk.values.PCollection.IsBounded;
  import org.apache.beam.sdk.values.TupleTag;
  import org.apache.beam.sdk.values.WindowedValue;
  import org.apache.beam.sdk.values.WindowedValues;
  import org.apache.beam.sdk.values.WindowingStrategy;
  import org.apache.samza.config.Config;
  import org.apache.samza.context.Context;
  import org.apache.samza.operators.Scheduler;
  import org.joda.time.Instant;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /** Samza operator for {@link org.apache.beam.sdk.transforms.GroupByKey}. */
  @SuppressWarnings({
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  public class GroupByKeyOp<K, InputT, OutputT>
      implements Op<KeyedWorkItem<K, InputT>, KV<K, OutputT>, K> {
    private static final Logger LOG = LoggerFactory.getLogger(GroupByKeyOp.class);
    private static final String TIMER_STATE_ID = "timer";
  
    private final TupleTag<KV<K, OutputT>> mainOutputTag;
    private final KeyedWorkItemCoder<K, InputT> inputCoder;
    private final WindowingStrategy<?, BoundedWindow> windowingStrategy;
    private final OutputManagerFactory<KV<K, OutputT>> outputManagerFactory;
    private final Coder<K> keyCoder;
    private final SystemReduceFn<K, InputT, ?, OutputT, BoundedWindow> reduceFn;
    private final String transformFullName;
    private final String transformId;
    private final IsBounded isBounded;
  
    private transient StateInternalsFactory<K> stateInternalsFactory;
    private transient SamzaTimerInternalsFactory<K> timerInternalsFactory;
    private transient DoFnRunner<KeyedWorkItem<K, InputT>, KV<K, OutputT>> fnRunner;
    private transient SamzaPipelineOptions pipelineOptions;
  
    public GroupByKeyOp(
        TupleTag<KV<K, OutputT>> mainOutputTag,
        Coder<KeyedWorkItem<K, InputT>> inputCoder,
        SystemReduceFn<K, InputT, ?, OutputT, BoundedWindow> reduceFn,
        WindowingStrategy<?, BoundedWindow> windowingStrategy,
        OutputManagerFactory<KV<K, OutputT>> outputManagerFactory,
        String transformFullName,
        String transformId,
        IsBounded isBounded) {
      this.mainOutputTag = mainOutputTag;
      this.windowingStrategy = windowingStrategy;
      this.outputManagerFactory = outputManagerFactory;
      this.transformFullName = transformFullName;
      this.transformId = transformId;
      this.isBounded = isBounded;
  
      if (!(inputCoder instanceof KeyedWorkItemCoder)) {
        throw new IllegalArgumentException(
            String.format(
                "GroupByKeyOp requires input to use KeyedWorkItemCoder. Got: %s",
                inputCoder.getClass()));
      }
      this.inputCoder = (KeyedWorkItemCoder<K, InputT>) inputCoder;
      this.keyCoder = this.inputCoder.getKeyCoder();
      this.reduceFn = reduceFn;
    }
  
    @Override
    public void open(
        Config config,
        Context context,
        Scheduler<KeyedTimerData<K>> timerRegistry,
        OpEmitter<KV<K, OutputT>> emitter) {
  
      final SamzaExecutionContext samzaExecutionContext =
          (SamzaExecutionContext) context.getApplicationContainerContext();
      this.pipelineOptions = samzaExecutionContext.getPipelineOptions();
  
      final SamzaStoreStateInternals.Factory<?> nonKeyedStateInternalsFactory =
          SamzaStoreStateInternals.createNonKeyedStateInternalsFactory(
              transformId, context.getTaskContext(), pipelineOptions);
  
      final WindowedValueMultiReceiver outputManager = outputManagerFactory.create(emitter);
  
      this.stateInternalsFactory =
          new SamzaStoreStateInternals.Factory<>(
              transformId,
              Collections.singletonMap(
                  SamzaStoreStateInternals.BEAM_STORE,
                  SamzaStoreStateInternals.getBeamStore(context.getTaskContext())),
              keyCoder,
              pipelineOptions.getStoreBatchGetSize());
  
      this.timerInternalsFactory =
          SamzaTimerInternalsFactory.createTimerInternalFactory(
              keyCoder,
              timerRegistry,
              TIMER_STATE_ID,
              nonKeyedStateInternalsFactory,
              windowingStrategy,
              isBounded,
              pipelineOptions);
  
      final DoFn<KeyedWorkItem<K, InputT>, KV<K, OutputT>> doFn =
          GroupAlsoByWindowViaWindowSetNewDoFn.create(
              windowingStrategy,
              stateInternalsFactory,
              timerInternalsFactory,
              NullSideInputReader.of(Collections.emptyList()),
              reduceFn,
              outputManager,
              mainOutputTag);
  
      final KeyedInternals<K> keyedInternals =
          new KeyedInternals<>(stateInternalsFactory, timerInternalsFactory);
  
      final StepContext stepContext =
          new StepContext() {
            @Override
            public StateInternals stateInternals() {
              return keyedInternals.stateInternals();
            }
  
            @Override
            public TimerInternals timerInternals() {
              return keyedInternals.timerInternals();
            }
          };
  
      final DoFnRunner<KeyedWorkItem<K, InputT>, KV<K, OutputT>> doFnRunner =
          DoFnRunners.simpleRunner(
              PipelineOptionsFactory.create(),
              doFn,
              NullSideInputReader.of(Collections.emptyList()),
              outputManager,
              mainOutputTag,
              Collections.emptyList(),
              stepContext,
              null,
              Collections.emptyMap(),
              windowingStrategy,
              DoFnSchemaInformation.create(),
              Collections.emptyMap());
  
      final DoFnRunner<KeyedWorkItem<K, InputT>, KV<K, OutputT>> dropLateDataRunner =
          pipelineOptions.getDropLateData()
              ? DoFnRunners.lateDataDroppingRunner(
                  doFnRunner, keyedInternals.timerInternals(), windowingStrategy)
              : doFnRunner;
  
      final SamzaExecutionContext executionContext =
          (SamzaExecutionContext) context.getApplicationContainerContext();
      final DoFnRunner<KeyedWorkItem<K, InputT>, KV<K, OutputT>> doFnRunnerWithMetrics =
          DoFnRunnerWithMetrics.wrap(
              dropLateDataRunner, executionContext.getMetricsContainer(), transformFullName);
  
      this.fnRunner = new DoFnRunnerWithKeyedInternals<>(doFnRunnerWithMetrics, keyedInternals);
    }
  
    @Override
    public void processElement(
        WindowedValue<KeyedWorkItem<K, InputT>> inputElement, OpEmitter<KV<K, OutputT>> emitter) {
      fnRunner.startBundle();
      fnRunner.processElement(inputElement);
      fnRunner.finishBundle();
    }
  
    @Override
    public void processWatermark(Instant watermark, OpEmitter<KV<K, OutputT>> emitter) {
      timerInternalsFactory.setInputWatermark(watermark);
  
      Collection<KeyedTimerData<K>> readyTimers = timerInternalsFactory.removeReadyTimers();
      if (!readyTimers.isEmpty()) {
        fnRunner.startBundle();
        for (KeyedTimerData<K> keyedTimerData : readyTimers) {
          fireTimer(keyedTimerData.getKey(), keyedTimerData.getTimerData());
        }
        fnRunner.finishBundle();
      }
  
      if (timerInternalsFactory.getOutputWatermark() == null
          || timerInternalsFactory.getOutputWatermark().isBefore(watermark)) {
        timerInternalsFactory.setOutputWatermark(watermark);
        emitter.emitWatermark(timerInternalsFactory.getOutputWatermark());
      }
    }
  
    @Override
    public void processTimer(KeyedTimerData<K> keyedTimerData, OpEmitter<KV<K, OutputT>> emitter) {
      fnRunner.startBundle();
      fireTimer(keyedTimerData.getKey(), keyedTimerData.getTimerData());
      fnRunner.finishBundle();
  
      timerInternalsFactory.removeProcessingTimer(keyedTimerData);
    }
  
    private void fireTimer(K key, TimerData timer) {
      LOG.debug("Firing timer {} for key {}", timer, key);
      fnRunner.processElement(
          WindowedValues.valueInGlobalWindow(
              KeyedWorkItems.timersWorkItem(key, Collections.singletonList(timer))));
    }
  }
  