/*
 * 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.nexmark;

import static org.apache.beam.sdk.nexmark.NexmarkUtils.PubSubMode.COMBINED;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;

import com.google.api.services.bigquery.model.TableFieldSchema;
import com.google.api.services.bigquery.model.TableRow;
import com.google.api.services.bigquery.model.TableSchema;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ThreadLocalRandom;
import javax.annotation.Nullable;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.io.AvroIO;
import org.apache.beam.sdk.io.TextIO;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubIO;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubMessage;
import org.apache.beam.sdk.io.kafka.KafkaIO;
import org.apache.beam.sdk.nexmark.NexmarkUtils.PubSubMode;
import org.apache.beam.sdk.nexmark.NexmarkUtils.SourceType;
import org.apache.beam.sdk.nexmark.model.Auction;
import org.apache.beam.sdk.nexmark.model.Bid;
import org.apache.beam.sdk.nexmark.model.Event;
import org.apache.beam.sdk.nexmark.model.KnownSize;
import org.apache.beam.sdk.nexmark.model.Person;
import org.apache.beam.sdk.nexmark.queries.BoundedSideInputJoin;
import org.apache.beam.sdk.nexmark.queries.BoundedSideInputJoinModel;
import org.apache.beam.sdk.nexmark.queries.NexmarkQuery;
import org.apache.beam.sdk.nexmark.queries.NexmarkQueryModel;
import org.apache.beam.sdk.nexmark.queries.NexmarkQueryUtil;
import org.apache.beam.sdk.nexmark.queries.Query0;
import org.apache.beam.sdk.nexmark.queries.Query0Model;
import org.apache.beam.sdk.nexmark.queries.Query1;
import org.apache.beam.sdk.nexmark.queries.Query10;
import org.apache.beam.sdk.nexmark.queries.Query11;
import org.apache.beam.sdk.nexmark.queries.Query12;
import org.apache.beam.sdk.nexmark.queries.Query1Model;
import org.apache.beam.sdk.nexmark.queries.Query2;
import org.apache.beam.sdk.nexmark.queries.Query2Model;
import org.apache.beam.sdk.nexmark.queries.Query3;
import org.apache.beam.sdk.nexmark.queries.Query3Model;
import org.apache.beam.sdk.nexmark.queries.Query4;
import org.apache.beam.sdk.nexmark.queries.Query4Model;
import org.apache.beam.sdk.nexmark.queries.Query5;
import org.apache.beam.sdk.nexmark.queries.Query5Model;
import org.apache.beam.sdk.nexmark.queries.Query6;
import org.apache.beam.sdk.nexmark.queries.Query6Model;
import org.apache.beam.sdk.nexmark.queries.Query7;
import org.apache.beam.sdk.nexmark.queries.Query7Model;
import org.apache.beam.sdk.nexmark.queries.Query8;
import org.apache.beam.sdk.nexmark.queries.Query8Model;
import org.apache.beam.sdk.nexmark.queries.Query9;
import org.apache.beam.sdk.nexmark.queries.Query9Model;
import org.apache.beam.sdk.nexmark.queries.SessionSideInputJoin;
import org.apache.beam.sdk.nexmark.queries.SessionSideInputJoinModel;
import org.apache.beam.sdk.nexmark.queries.sql.SqlBoundedSideInputJoin;
import org.apache.beam.sdk.nexmark.queries.sql.SqlQuery0;
import org.apache.beam.sdk.nexmark.queries.sql.SqlQuery1;
import org.apache.beam.sdk.nexmark.queries.sql.SqlQuery2;
import org.apache.beam.sdk.nexmark.queries.sql.SqlQuery3;
import org.apache.beam.sdk.nexmark.queries.sql.SqlQuery7;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testutils.metrics.MetricsReader;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.util.CoderUtils;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.TimestampedValue;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Strings;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
import org.apache.kafka.common.serialization.ByteArrayDeserializer;
import org.apache.kafka.common.serialization.ByteArraySerializer;
import org.apache.kafka.common.serialization.LongDeserializer;
import org.apache.kafka.common.serialization.StringSerializer;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.slf4j.LoggerFactory;

/** Run a single Nexmark query using a given configuration. */
public class NexmarkLauncher<OptionT extends NexmarkOptions> {

  private static final org.slf4j.Logger LOG = LoggerFactory.getLogger(NexmarkLauncher.class);

  /** Command line parameter value for query language. */
  private static final String SQL = "sql";

  /** Minimum number of samples needed for 'stead-state' rate calculation. */
  private static final int MIN_SAMPLES = 9;
  /** Minimum length of time over which to consider samples for 'steady-state' rate calculation. */
  private static final Duration MIN_WINDOW = Duration.standardMinutes(2);
  /** Delay between perf samples. */
  private static final Duration PERF_DELAY = Duration.standardSeconds(15);
  /**
   * How long to let streaming pipeline run after all events have been generated and we've seen no
   * activity.
   */
  private static final Duration DONE_DELAY = Duration.standardMinutes(1);
  /** How long to allow no activity at sources and sinks without warning. */
  private static final Duration STUCK_WARNING_DELAY = Duration.standardMinutes(10);
  /**
   * How long to let streaming pipeline run after we've seen no activity at sources or sinks, even
   * if all events have not been generated.
   */
  private static final Duration STUCK_TERMINATE_DELAY = Duration.standardHours(1);

  /** NexmarkOptions for this run. */
  private final OptionT options;

  /** Which configuration we are running. */
  private NexmarkConfiguration configuration;

  /** If in --pubsubMode=COMBINED, the event monitor for the publisher pipeline. Otherwise null. */
  @Nullable private Monitor<Event> publisherMonitor;

  /**
   * If in --pubsubMode=COMBINED, the pipeline result for the publisher pipeline. Otherwise null.
   */
  @Nullable private PipelineResult publisherResult;

  /** Result for the main pipeline. */
  @Nullable private PipelineResult mainResult;

  /** Query name we are running. */
  @Nullable private String queryName;

  /** Full path of the PubSub topic (when PubSub is enabled). */
  @Nullable private String pubsubTopic;

  /** Full path of the PubSub subscription (when PubSub is enabled). */
  @Nullable private String pubsubSubscription;

  @Nullable private PubsubHelper pubsubHelper;

  public NexmarkLauncher(OptionT options, NexmarkConfiguration configuration) {
    this.options = options;
    this.configuration = configuration;
  }

  /** Is this query running in streaming mode? */
  private boolean isStreaming() {
    return options.isStreaming();
  }

  /** Return maximum number of workers. */
  private int maxNumWorkers() {
    return 5;
  }

  /**
   * Find a 'steady state' events/sec from {@code snapshots} and store it in {@code perf} if found.
   */
  private void captureSteadyState(NexmarkPerf perf, List<NexmarkPerf.ProgressSnapshot> snapshots) {
    if (!options.isStreaming()) {
      return;
    }

    // Find the first sample with actual event and result counts.
    int dataStart = 0;
    for (; dataStart < snapshots.size(); dataStart++) {
      if (snapshots.get(dataStart).numEvents >= 0 && snapshots.get(dataStart).numResults >= 0) {
        break;
      }
    }

    // Find the last sample which demonstrated progress.
    int dataEnd = snapshots.size() - 1;
    for (; dataEnd > dataStart; dataEnd--) {
      if (snapshots.get(dataEnd).anyActivity(snapshots.get(dataEnd - 1))) {
        break;
      }
    }

    int numSamples = dataEnd - dataStart + 1;
    if (numSamples < MIN_SAMPLES) {
      // Not enough samples.
      NexmarkUtils.console(
          "%d samples not enough to calculate steady-state event rate", numSamples);
      return;
    }

    // We'll look at only the middle third samples.
    int sampleStart = dataStart + numSamples / 3;
    int sampleEnd = dataEnd - numSamples / 3;

    double sampleSec =
        snapshots.get(sampleEnd).secSinceStart - snapshots.get(sampleStart).secSinceStart;
    if (sampleSec < MIN_WINDOW.getStandardSeconds()) {
      // Not sampled over enough time.
      NexmarkUtils.console(
          "sample of %.1f sec not long enough to calculate steady-state event rate", sampleSec);
      return;
    }

    // Find rate with least squares error.
    double sumxx = 0.0;
    double sumxy = 0.0;
    long prevNumEvents = -1;
    for (int i = sampleStart; i <= sampleEnd; i++) {
      if (prevNumEvents == snapshots.get(i).numEvents) {
        // Skip samples with no change in number of events since they contribute no data.
        continue;
      }
      // Use the effective runtime instead of wallclock time so we can
      // insulate ourselves from delays and stutters in the query manager.
      double x = snapshots.get(i).runtimeSec;
      prevNumEvents = snapshots.get(i).numEvents;
      double y = prevNumEvents;
      sumxx += x * x;
      sumxy += x * y;
    }
    double eventsPerSec = sumxy / sumxx;
    NexmarkUtils.console("revising events/sec from %.1f to %.1f", perf.eventsPerSec, eventsPerSec);
    perf.eventsPerSec = eventsPerSec;
  }

  /** Return the current performance given {@code eventMonitor} and {@code resultMonitor}. */
  private NexmarkPerf currentPerf(
      long startMsSinceEpoch,
      long now,
      PipelineResult result,
      List<NexmarkPerf.ProgressSnapshot> snapshots,
      Monitor<?> eventMonitor,
      Monitor<?> resultMonitor) {
    NexmarkPerf perf = new NexmarkPerf();

    MetricsReader eventMetrics = new MetricsReader(result, eventMonitor.name);

    long numEvents = eventMetrics.getCounterMetric(eventMonitor.prefix + ".elements");
    long numEventBytes = eventMetrics.getCounterMetric(eventMonitor.prefix + ".bytes");
    long eventStart = eventMetrics.getStartTimeMetric(eventMonitor.prefix + ".startTime");
    long eventEnd = eventMetrics.getEndTimeMetric(eventMonitor.prefix + ".endTime");

    MetricsReader resultMetrics = new MetricsReader(result, resultMonitor.name);

    long numResults = resultMetrics.getCounterMetric(resultMonitor.prefix + ".elements");
    long numResultBytes = resultMetrics.getCounterMetric(resultMonitor.prefix + ".bytes");
    long resultStart = resultMetrics.getStartTimeMetric(resultMonitor.prefix + ".startTime");
    long resultEnd = resultMetrics.getEndTimeMetric(resultMonitor.prefix + ".endTime");
    long timestampStart =
        resultMetrics.getStartTimeMetric(resultMonitor.prefix + ".startTimestamp");
    long timestampEnd = resultMetrics.getEndTimeMetric(resultMonitor.prefix + ".endTimestamp");

    long effectiveEnd = -1;
    if (eventEnd >= 0 && resultEnd >= 0) {
      // It is possible for events to be generated after the last result was emitted.
      // (Eg Query 2, which only yields results for a small prefix of the event stream.)
      // So use the max of last event and last result times.
      effectiveEnd = Math.max(eventEnd, resultEnd);
    } else if (resultEnd >= 0) {
      effectiveEnd = resultEnd;
    } else if (eventEnd >= 0) {
      // During startup we may have no result yet, but we would still like to track how
      // long the pipeline has been running.
      effectiveEnd = eventEnd;
    }

    if (effectiveEnd >= 0 && eventStart >= 0 && effectiveEnd >= eventStart) {
      perf.runtimeSec = (effectiveEnd - eventStart) / 1000.0;
    }

    if (numEvents >= 0) {
      perf.numEvents = numEvents;
    }

    if (numEvents >= 0 && perf.runtimeSec > 0.0) {
      // For streaming we may later replace this with a 'steady-state' value calculated
      // from the progress snapshots.
      perf.eventsPerSec = numEvents / perf.runtimeSec;
    }

    if (numEventBytes >= 0 && perf.runtimeSec > 0.0) {
      perf.eventBytesPerSec = numEventBytes / perf.runtimeSec;
    }

    if (numResults >= 0) {
      perf.numResults = numResults;
    }

    if (numResults >= 0 && perf.runtimeSec > 0.0) {
      perf.resultsPerSec = numResults / perf.runtimeSec;
    }

    if (numResultBytes >= 0 && perf.runtimeSec > 0.0) {
      perf.resultBytesPerSec = numResultBytes / perf.runtimeSec;
    }

    if (eventStart >= 0) {
      perf.startupDelaySec = (eventStart - startMsSinceEpoch) / 1000.0;
    }

    if (resultStart >= 0 && eventStart >= 0 && resultStart >= eventStart) {
      perf.processingDelaySec = (resultStart - eventStart) / 1000.0;
    }

    if (timestampStart >= 0 && timestampEnd >= 0 && perf.runtimeSec > 0.0) {
      double eventRuntimeSec = (timestampEnd - timestampStart) / 1000.0;
      perf.timeDilation = eventRuntimeSec / perf.runtimeSec;
    }

    if (resultEnd >= 0) {
      // Fill in the shutdown delay assuming the job has now finished.
      perf.shutdownDelaySec = (now - resultEnd) / 1000.0;
    }

    // As soon as available, try to capture cumulative cost at this point too.

    NexmarkPerf.ProgressSnapshot snapshot = new NexmarkPerf.ProgressSnapshot();
    snapshot.secSinceStart = (now - startMsSinceEpoch) / 1000.0;
    snapshot.runtimeSec = perf.runtimeSec;
    snapshot.numEvents = numEvents;
    snapshot.numResults = numResults;
    snapshots.add(snapshot);

    captureSteadyState(perf, snapshots);

    return perf;
  }

  /** Build and run a pipeline using specified options. */
  interface PipelineBuilder<OptionT extends NexmarkOptions> {
    void build(OptionT publishOnlyOptions);
  }

  /** Invoke the builder with options suitable for running a publish-only child pipeline. */
  private void invokeBuilderForPublishOnlyPipeline(PipelineBuilder<NexmarkOptions> builder) {
    String jobName = options.getJobName();
    String appName = options.getAppName();
    int numWorkers = options.getNumWorkers();
    int maxNumWorkers = options.getMaxNumWorkers();

    options.setJobName("p-" + jobName);
    options.setAppName("p-" + appName);
    int eventGeneratorWorkers = configuration.numEventGenerators;
    // TODO: assign one generator per core rather than one per worker.
    if (numWorkers > 0 && eventGeneratorWorkers > 0) {
      options.setNumWorkers(Math.min(numWorkers, eventGeneratorWorkers));
    }
    if (maxNumWorkers > 0 && eventGeneratorWorkers > 0) {
      options.setMaxNumWorkers(Math.min(maxNumWorkers, eventGeneratorWorkers));
    }
    try {
      builder.build(options);
    } finally {
      options.setJobName(jobName);
      options.setAppName(appName);
      options.setNumWorkers(numWorkers);
      options.setMaxNumWorkers(maxNumWorkers);
    }
  }

  /**
   * Monitor the performance and progress of a running job. Return final performance if it was
   * measured.
   */
  @Nullable
  private NexmarkPerf monitor(NexmarkQuery query) {
    if (!options.getMonitorJobs()) {
      return null;
    }

    if (configuration.debug) {
      NexmarkUtils.console("Waiting for main pipeline to 'finish'");
    } else {
      NexmarkUtils.console("--debug=false, so job will not self-cancel");
    }

    PipelineResult job = mainResult;
    PipelineResult publisherJob = publisherResult;
    List<NexmarkPerf.ProgressSnapshot> snapshots = new ArrayList<>();
    long startMsSinceEpoch = System.currentTimeMillis();
    long endMsSinceEpoch = -1;
    if (options.getRunningTimeMinutes() != null) {
      endMsSinceEpoch =
          startMsSinceEpoch
              + Duration.standardMinutes(options.getRunningTimeMinutes()).getMillis()
              - Duration.standardSeconds(configuration.preloadSeconds).getMillis();
    }
    long lastActivityMsSinceEpoch = -1;
    NexmarkPerf perf = null;
    boolean waitingForShutdown = false;
    boolean cancelJob = false;
    boolean publisherCancelled = false;
    List<String> errors = new ArrayList<>();

    while (true) {
      long now = System.currentTimeMillis();
      if (endMsSinceEpoch >= 0 && now > endMsSinceEpoch && !waitingForShutdown) {
        NexmarkUtils.console("Reached end of test, cancelling job");
        try {
          cancelJob = true;
          job.cancel();
        } catch (IOException e) {
          throw new RuntimeException("Unable to cancel main job: ", e);
        }
        if (publisherResult != null) {
          try {
            publisherJob.cancel();
          } catch (IOException e) {
            throw new RuntimeException("Unable to cancel publisher job: ", e);
          }
          publisherCancelled = true;
        }
        waitingForShutdown = true;
      }

      PipelineResult.State state = job.getState();
      NexmarkUtils.console(
          "%s %s%s", state, queryName, waitingForShutdown ? " (waiting for shutdown)" : "");

      NexmarkPerf currPerf;
      if (configuration.debug) {
        currPerf =
            currentPerf(
                startMsSinceEpoch, now, job, snapshots, query.eventMonitor, query.resultMonitor);
      } else {
        currPerf = null;
      }

      if (perf == null || perf.anyActivity(currPerf)) {
        lastActivityMsSinceEpoch = now;
      }

      if (options.isStreaming() && !waitingForShutdown) {
        Duration quietFor = new Duration(lastActivityMsSinceEpoch, now);
        long fatalCount = new MetricsReader(job, query.getName()).getCounterMetric("fatal");

        if (fatalCount == -1) {
          fatalCount = 0;
        }

        if (fatalCount > 0) {
          NexmarkUtils.console("ERROR: job has fatal errors, cancelling.");
          errors.add(String.format("Pipeline reported %s fatal errors", fatalCount));
          waitingForShutdown = true;
          cancelJob = true;
        } else if (configuration.debug
            && configuration.numEvents > 0
            && currPerf.numEvents == configuration.numEvents
            && currPerf.numResults >= 0
            && quietFor.isLongerThan(DONE_DELAY)) {
          NexmarkUtils.console("streaming query appears to have finished waiting for completion.");
          waitingForShutdown = true;
        } else if (quietFor.isLongerThan(STUCK_TERMINATE_DELAY)) {
          NexmarkUtils.console(
              "ERROR: streaming query appears to have been stuck for %d minutes, cancelling job.",
              quietFor.getStandardMinutes());
          errors.add(
              String.format(
                  "Cancelling streaming job since it appeared stuck for %d min.",
                  quietFor.getStandardMinutes()));
          waitingForShutdown = true;
          cancelJob = true;
        } else if (quietFor.isLongerThan(STUCK_WARNING_DELAY)) {
          NexmarkUtils.console(
              "WARNING: streaming query appears to have been stuck for %d min.",
              quietFor.getStandardMinutes());
        }

        if (cancelJob) {
          try {
            job.cancel();
          } catch (IOException e) {
            throw new RuntimeException("Unable to cancel main job: ", e);
          }
        }
      }

      perf = currPerf;

      boolean running = true;
      switch (state) {
        case UNKNOWN:
        case STOPPED:
        case RUNNING:
          // Keep going.
          break;
        case DONE:
          // All done.
          running = false;
          break;
        case CANCELLED:
          running = false;
          if (!cancelJob) {
            errors.add("Job was unexpectedly cancelled");
          }
          break;
        case FAILED:
        case UPDATED:
          // Abnormal termination.
          running = false;
          errors.add("Job was unexpectedly updated");
          break;
      }

      if (!running) {
        break;
      }

      if (lastActivityMsSinceEpoch == now) {
        NexmarkUtils.console("new perf %s", perf);
      } else {
        NexmarkUtils.console("no activity");
      }

      try {
        Thread.sleep(PERF_DELAY.getMillis());
      } catch (InterruptedException e) {
        Thread.interrupted();
        NexmarkUtils.console("Interrupted: pipeline is still running");
      }
    }

    perf.errors = errors;
    perf.snapshots = snapshots;

    if (publisherResult != null) {
      NexmarkUtils.console("Shutting down publisher pipeline.");
      try {
        if (!publisherCancelled) {
          publisherJob.cancel();
        }
        publisherJob.waitUntilFinish(Duration.standardMinutes(5));
      } catch (IOException e) {
        throw new RuntimeException("Unable to cancel publisher job: ", e);
      }
    }

    return perf;
  }

  // ================================================================================
  // Basic sources and sinks
  // ================================================================================

  /** Return a topic name. */
  private String shortTopic(long now) {
    String baseTopic = options.getPubsubTopic();
    if (Strings.isNullOrEmpty(baseTopic)) {
      throw new RuntimeException("Missing --pubsubTopic");
    }
    switch (options.getResourceNameMode()) {
      case VERBATIM:
        return baseTopic;
      case QUERY:
        return String.format("%s_%s_source", baseTopic, queryName);
      case QUERY_AND_SALT:
        return String.format("%s_%s_%d_source", baseTopic, queryName, now);
      case QUERY_RUNNER_AND_MODE:
        return String.format(
            "%s_%s_%s_%s_source",
            baseTopic, queryName, options.getRunner().getSimpleName(), options.isStreaming());
    }
    throw new RuntimeException("Unrecognized enum " + options.getResourceNameMode());
  }

  /** Return a subscription name. */
  private String shortSubscription(long now) {
    String baseSubscription = options.getPubsubSubscription();
    if (Strings.isNullOrEmpty(baseSubscription)) {
      throw new RuntimeException("Missing --pubsubSubscription");
    }
    switch (options.getResourceNameMode()) {
      case VERBATIM:
        return baseSubscription;
      case QUERY:
        return String.format("%s_%s_source", baseSubscription, queryName);
      case QUERY_AND_SALT:
        return String.format("%s_%s_%d_source", baseSubscription, queryName, now);
      case QUERY_RUNNER_AND_MODE:
        return String.format(
            "%s_%s_%s_%s_source",
            baseSubscription,
            queryName,
            options.getRunner().getSimpleName(),
            options.isStreaming());
    }
    throw new RuntimeException("Unrecognized enum " + options.getResourceNameMode());
  }

  /** Return a file name for plain text. */
  private String textFilename(long now) {
    String baseFilename = options.getOutputPath();
    if (Strings.isNullOrEmpty(baseFilename)) {
      throw new RuntimeException("Missing --outputPath");
    }
    switch (options.getResourceNameMode()) {
      case VERBATIM:
        return baseFilename;
      case QUERY:
        return String.format("%s/nexmark_%s.txt", baseFilename, queryName);
      case QUERY_AND_SALT:
        return String.format("%s/nexmark_%s_%d.txt", baseFilename, queryName, now);
      case QUERY_RUNNER_AND_MODE:
        return String.format(
            "%s/nexmark_%s_%s_%s",
            baseFilename, queryName, options.getRunner().getSimpleName(), options.isStreaming());
    }
    throw new RuntimeException("Unrecognized enum " + options.getResourceNameMode());
  }

  /** Return a directory for logs. */
  private String logsDir(long now) {
    String baseFilename = options.getOutputPath();
    if (Strings.isNullOrEmpty(baseFilename)) {
      throw new RuntimeException("Missing --outputPath");
    }
    switch (options.getResourceNameMode()) {
      case VERBATIM:
        return baseFilename;
      case QUERY:
        return String.format("%s/logs_%s", baseFilename, queryName);
      case QUERY_AND_SALT:
        return String.format("%s/logs_%s_%d", baseFilename, queryName, now);
      case QUERY_RUNNER_AND_MODE:
        return String.format(
            "%s/logs_%s_%s_%s",
            baseFilename, queryName, options.getRunner().getSimpleName(), options.isStreaming());
    }
    throw new RuntimeException("Unrecognized enum " + options.getResourceNameMode());
  }

  /** Return a source of synthetic events. */
  private PCollection<Event> sourceEventsFromSynthetic(Pipeline p) {
    if (isStreaming()) {
      NexmarkUtils.console("Generating %d events in streaming mode", configuration.numEvents);
      return p.apply(queryName + ".ReadUnbounded", NexmarkUtils.streamEventsSource(configuration));
    } else {
      NexmarkUtils.console("Generating %d events in batch mode", configuration.numEvents);
      return p.apply(queryName + ".ReadBounded", NexmarkUtils.batchEventsSource(configuration));
    }
  }

  /** Return source of events from Pubsub. */
  private PCollection<Event> sourceEventsFromPubsub(Pipeline p) {
    NexmarkUtils.console("Reading events from Pubsub %s", pubsubSubscription);

    PubsubIO.Read<PubsubMessage> io =
        PubsubIO.readMessagesWithAttributes()
            .fromSubscription(pubsubSubscription)
            .withIdAttribute(NexmarkUtils.PUBSUB_ID);
    if (!configuration.usePubsubPublishTime) {
      io = io.withTimestampAttribute(NexmarkUtils.PUBSUB_TIMESTAMP);
    }

    return p.apply(queryName + ".ReadPubsubEvents", io)
        .apply(queryName + ".PubsubMessageToEvent", ParDo.of(new PubsubMessageEventDoFn()));
  }

  static final DoFn<Event, byte[]> EVENT_TO_BYTEARRAY =
      new DoFn<Event, byte[]>() {
        @ProcessElement
        public void processElement(ProcessContext c) throws IOException {
          byte[] encodedEvent = CoderUtils.encodeToByteArray(Event.CODER, c.element());
          c.output(encodedEvent);
        }
      };

  /** Send {@code events} to Kafka. */
  private void sinkEventsToKafka(PCollection<Event> events) {
    checkArgument((options.getBootstrapServers() != null), "Missing --bootstrapServers");
    NexmarkUtils.console("Writing events to Kafka Topic %s", options.getKafkaTopic());

    PCollection<byte[]> eventToBytes = events.apply("Event to bytes", ParDo.of(EVENT_TO_BYTEARRAY));
    eventToBytes.apply(
        KafkaIO.<Void, byte[]>write()
            .withBootstrapServers(options.getBootstrapServers())
            .withTopic(options.getKafkaTopic())
            .withValueSerializer(ByteArraySerializer.class)
            .values());
  }

  static final DoFn<KV<Long, byte[]>, Event> BYTEARRAY_TO_EVENT =
      new DoFn<KV<Long, byte[]>, Event>() {
        @ProcessElement
        public void processElement(ProcessContext c) throws IOException {
          byte[] encodedEvent = c.element().getValue();
          Event event = CoderUtils.decodeFromByteArray(Event.CODER, encodedEvent);
          c.output(event);
        }
      };

  /** Return source of events from Kafka. */
  private PCollection<Event> sourceEventsFromKafka(Pipeline p, final Instant now) {
    checkArgument((options.getBootstrapServers() != null), "Missing --bootstrapServers");
    NexmarkUtils.console("Reading events from Kafka Topic %s", options.getKafkaTopic());

    KafkaIO.Read<Long, byte[]> read =
        KafkaIO.<Long, byte[]>read()
            .withBootstrapServers(options.getBootstrapServers())
            .withTopic(options.getKafkaTopic())
            .withKeyDeserializer(LongDeserializer.class)
            .withValueDeserializer(ByteArrayDeserializer.class)
            .withStartReadTime(now)
            .withMaxNumRecords(
                options.getNumEvents() != null ? options.getNumEvents() : Long.MAX_VALUE);

    return p.apply(queryName + ".ReadKafkaEvents", read.withoutMetadata())
        .apply(queryName + ".KafkaToEvents", ParDo.of(BYTEARRAY_TO_EVENT));
  }

  /** Return Avro source of events from {@code options.getInputFilePrefix}. */
  private PCollection<Event> sourceEventsFromAvro(Pipeline p) {
    String filename = options.getInputPath();
    if (Strings.isNullOrEmpty(filename)) {
      throw new RuntimeException("Missing --inputPath");
    }
    NexmarkUtils.console("Reading events from Avro files at %s", filename);
    return p.apply(
            queryName + ".ReadAvroEvents", AvroIO.read(Event.class).from(filename + "*.avro"))
        .apply("OutputWithTimestamp", NexmarkQueryUtil.EVENT_TIMESTAMP_FROM_DATA);
  }

  /** Send {@code events} to Pubsub. */
  private void sinkEventsToPubsub(PCollection<Event> events) {
    checkState(pubsubTopic != null, "Pubsub topic needs to be set up before initializing sink");
    NexmarkUtils.console("Writing events to Pubsub %s", pubsubTopic);

    PubsubIO.Write<PubsubMessage> io =
        PubsubIO.writeMessages().to(pubsubTopic).withIdAttribute(NexmarkUtils.PUBSUB_ID);
    if (!configuration.usePubsubPublishTime) {
      io = io.withTimestampAttribute(NexmarkUtils.PUBSUB_TIMESTAMP);
    }

    events
        .apply(queryName + ".EventToPubsubMessage", ParDo.of(new EventPubsubMessageDoFn()))
        .apply(queryName + ".WritePubsubEvents", io);
  }

  /** Send {@code formattedResults} to Kafka. */
  private void sinkResultsToKafka(PCollection<String> formattedResults) {
    checkArgument((options.getBootstrapServers() != null), "Missing --bootstrapServers");
    NexmarkUtils.console("Writing results to Kafka Topic %s", options.getKafkaResultsTopic());

    formattedResults.apply(
        queryName + ".WriteKafkaResults",
        KafkaIO.<Void, String>write()
            .withBootstrapServers(options.getBootstrapServers())
            .withTopic(options.getKafkaResultsTopic())
            .withValueSerializer(StringSerializer.class)
            .values());
  }

  /** Send {@code formattedResults} to Pubsub. */
  private void sinkResultsToPubsub(PCollection<String> formattedResults, long now) {
    String shortTopic = shortTopic(now);
    NexmarkUtils.console("Writing results to Pubsub %s", shortTopic);
    PubsubIO.Write<String> io =
        PubsubIO.writeStrings().to(shortTopic).withIdAttribute(NexmarkUtils.PUBSUB_ID);
    if (!configuration.usePubsubPublishTime) {
      io = io.withTimestampAttribute(NexmarkUtils.PUBSUB_TIMESTAMP);
    }
    formattedResults.apply(queryName + ".WritePubsubResults", io);
  }

  /**
   * Sink all raw Events in {@code source} to {@code options.getOutputPath}. This will configure the
   * job to write the following files:
   *
   * <ul>
   *   <li>{@code $outputPath/event*.avro} All Event entities.
   *   <li>{@code $outputPath/auction*.avro} Auction entities.
   *   <li>{@code $outputPath/bid*.avro} Bid entities.
   *   <li>{@code $outputPath/person*.avro} Person entities.
   * </ul>
   *
   * @param source A PCollection of events.
   */
  private void sinkEventsToAvro(PCollection<Event> source) {
    String filename = options.getOutputPath();
    if (Strings.isNullOrEmpty(filename)) {
      throw new RuntimeException("Missing --outputPath");
    }
    NexmarkUtils.console("Writing events to Avro files at %s", filename);
    source.apply(
        queryName + ".WriteAvroEvents",
        AvroIO.write(Event.class).to(filename + "/event").withSuffix(".avro"));
    source
        .apply(NexmarkQueryUtil.JUST_BIDS)
        .apply(
            queryName + ".WriteAvroBids",
            AvroIO.write(Bid.class).to(filename + "/bid").withSuffix(".avro"));
    source
        .apply(NexmarkQueryUtil.JUST_NEW_AUCTIONS)
        .apply(
            queryName + ".WriteAvroAuctions",
            AvroIO.write(Auction.class).to(filename + "/auction").withSuffix(".avro"));
    source
        .apply(NexmarkQueryUtil.JUST_NEW_PERSONS)
        .apply(
            queryName + ".WriteAvroPeople",
            AvroIO.write(Person.class).to(filename + "/person").withSuffix(".avro"));
  }

  /** Send {@code formattedResults} to text files. */
  private void sinkResultsToText(PCollection<String> formattedResults, long now) {
    String filename = textFilename(now);
    NexmarkUtils.console("Writing results to text files at %s", filename);
    formattedResults.apply(queryName + ".WriteTextResults", TextIO.write().to(filename));
  }

  private static class StringToTableRow extends DoFn<String, TableRow> {
    @ProcessElement
    public void processElement(ProcessContext c) {
      int n = ThreadLocalRandom.current().nextInt(10);
      List<TableRow> records = new ArrayList<>(n);
      for (int i = 0; i < n; i++) {
        records.add(new TableRow().set("index", i).set("value", Integer.toString(i)));
      }
      c.output(new TableRow().set("result", c.element()).set("records", records));
    }
  }

  /** Send {@code formattedResults} to BigQuery. */
  private void sinkResultsToBigQuery(
      PCollection<String> formattedResults, long now, String version) {
    String tableSpec = NexmarkUtils.tableSpec(options, queryName, now, version);
    TableSchema tableSchema =
        new TableSchema()
            .setFields(
                ImmutableList.of(
                    new TableFieldSchema().setName("result").setType("STRING"),
                    new TableFieldSchema()
                        .setName("records")
                        .setMode("REPEATED")
                        .setType("RECORD")
                        .setFields(
                            ImmutableList.of(
                                new TableFieldSchema().setName("index").setType("INTEGER"),
                                new TableFieldSchema().setName("value").setType("STRING")))));
    NexmarkUtils.console("Writing results to BigQuery table %s", tableSpec);
    BigQueryIO.Write io =
        BigQueryIO.write()
            .to(tableSpec)
            .withSchema(tableSchema)
            .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED)
            .withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_APPEND);
    formattedResults
        .apply(queryName + ".StringToTableRow", ParDo.of(new StringToTableRow()))
        .apply(queryName + ".WriteBigQueryResults", io);
  }

  /** Creates or reuses PubSub topics and subscriptions as configured. */
  private void setupPubSubResources(long now) throws IOException {
    String shortTopic = shortTopic(now);
    String shortSubscription = shortSubscription(now);

    if (!options.getManageResources() || configuration.pubSubMode == PubSubMode.SUBSCRIBE_ONLY) {
      // The topic should already have been created by the user or
      // a companion 'PUBLISH_ONLY' process.
      pubsubTopic = pubsubHelper.reuseTopic(shortTopic).getPath();
    } else {
      // Create a fresh topic. It will be removed when the job is done.
      pubsubTopic = pubsubHelper.createTopic(shortTopic).getPath();
    }

    // Create/confirm the subscription.
    if (configuration.pubSubMode == PubSubMode.PUBLISH_ONLY) {
      // Nothing to consume.
    } else if (options.getManageResources()) {
      pubsubSubscription = pubsubHelper.createSubscription(shortTopic, shortSubscription).getPath();
    } else {
      // The subscription should already have been created by the user.
      pubsubSubscription = pubsubHelper.reuseSubscription(shortTopic, shortSubscription).getPath();
    }
  }

  // ================================================================================
  // Construct overall pipeline
  // ================================================================================

  /** Return source of events for this run, or null if we are simply publishing events to Pubsub. */
  private PCollection<Event> createSource(Pipeline p, final Instant now) throws IOException {
    PCollection<Event> source = null;

    switch (configuration.sourceType) {
      case DIRECT:
        source = sourceEventsFromSynthetic(p);
        break;
      case AVRO:
        source = sourceEventsFromAvro(p);
        break;
      case KAFKA:
      case PUBSUB:
        if (configuration.sourceType == SourceType.PUBSUB) {
          setupPubSubResources(now.getMillis());
        }
        // Setup the sink for the publisher.
        switch (configuration.pubSubMode) {
          case SUBSCRIBE_ONLY:
            // Nothing to publish.
            break;
          case PUBLISH_ONLY:
            {
              // Send synthesized events to Kafka or Pubsub in this job.
              PCollection<Event> events =
                  sourceEventsFromSynthetic(p)
                      .apply(queryName + ".Snoop", NexmarkUtils.snoop(queryName));
              if (configuration.sourceType == NexmarkUtils.SourceType.KAFKA) {
                sinkEventsToKafka(events);
              } else { // pubsub
                sinkEventsToPubsub(events);
              }
            }
            break;
          case COMBINED:
            // Send synthesized events to Kafka or Pubsub in separate publisher job.
            // We won't start the main pipeline until the publisher has sent the pre-load events.
            // We'll shutdown the publisher job when we notice the main job has finished.
            invokeBuilderForPublishOnlyPipeline(
                publishOnlyOptions -> {
                  Pipeline sp = Pipeline.create(publishOnlyOptions);
                  NexmarkUtils.setupPipeline(configuration.coderStrategy, sp);
                  publisherMonitor = new Monitor<>(queryName, "publisher");
                  PCollection<Event> events =
                      sourceEventsFromSynthetic(sp)
                          .apply(queryName + ".Monitor", publisherMonitor.getTransform());
                  if (configuration.sourceType == NexmarkUtils.SourceType.KAFKA) {
                    sinkEventsToKafka(events);
                  } else { // pubsub
                    sinkEventsToPubsub(events);
                  }
                  publisherResult = sp.run();
                  NexmarkUtils.console("Publisher job is started.");
                });
            break;
        }

        // Setup the source for the consumer.
        switch (configuration.pubSubMode) {
          case PUBLISH_ONLY:
            // Nothing to consume. Leave source null.
            break;
          case SUBSCRIBE_ONLY:
          case COMBINED:
            {
              // Read events from Kafka or Pubsub.
              if (configuration.sourceType == NexmarkUtils.SourceType.KAFKA) {
                // We need to have the same indexes for Publisher (sink) and Subscriber (source)
                // pipelines in COMBINED mode (when we run them in sequence). It means that
                // Subscriber should start reading from the same index as Publisher started to write
                // pre-load events even if we run Subscriber right after Publisher has been
                // finished. In other case. when pubSubMode=SUBSCRIBE_ONLY, now should be null and
                // it will be ignored.
                source =
                    sourceEventsFromKafka(p, configuration.pubSubMode == COMBINED ? now : null);
              } else {
                source = sourceEventsFromPubsub(p);
              }
            }
            break;
        }
        break;
    }
    return source;
  }

  private static final TupleTag<String> MAIN = new TupleTag<String>() {};
  private static final TupleTag<String> SIDE = new TupleTag<String>() {};

  private static class PartitionDoFn extends DoFn<String, String> {
    @ProcessElement
    public void processElement(ProcessContext c) {
      if (c.element().hashCode() % 2 == 0) {
        c.output(c.element());
      } else {
        c.output(SIDE, c.element());
      }
    }
  }

  /** Consume {@code results}. */
  private void sink(PCollection<TimestampedValue<KnownSize>> results, long now) {
    if (configuration.sinkType == NexmarkUtils.SinkType.COUNT_ONLY) {
      // Avoid the cost of formatting the results.
      results.apply(queryName + ".DevNull", NexmarkUtils.devNull(queryName));
      return;
    }

    PCollection<String> formattedResults =
        results.apply(queryName + ".Format", NexmarkUtils.format(queryName));
    if (options.getLogResults()) {
      formattedResults =
          formattedResults.apply(
              queryName + ".Results.Log", NexmarkUtils.log(queryName + ".Results"));
    }

    switch (configuration.sinkType) {
      case DEVNULL:
        // Discard all results
        formattedResults.apply(queryName + ".DevNull", NexmarkUtils.devNull(queryName));
        break;
      case PUBSUB:
        sinkResultsToPubsub(formattedResults, now);
        break;
      case KAFKA:
        sinkResultsToKafka(formattedResults);
        break;
      case TEXT:
        sinkResultsToText(formattedResults, now);
        break;
      case AVRO:
        NexmarkUtils.console(
            "WARNING: with --sinkType=AVRO, actual query results will be discarded.");
        break;
      case BIGQUERY:
        // Multiple BigQuery backends to mimic what most customers do.
        PCollectionTuple res =
            formattedResults.apply(
                queryName + ".Partition",
                ParDo.of(new PartitionDoFn()).withOutputTags(MAIN, TupleTagList.of(SIDE)));
        sinkResultsToBigQuery(res.get(MAIN), now, "main");
        sinkResultsToBigQuery(res.get(SIDE), now, "side");
        sinkResultsToBigQuery(formattedResults, now, "copy");
        break;
      case COUNT_ONLY:
        // Short-circuited above.
        throw new RuntimeException();
    }
  }

  // ================================================================================
  // Entry point
  // ================================================================================

  /**
   * Calculate the distribution of the expected rate of results per minute (in event time, not
   * wallclock time).
   */
  private void modelResultRates(NexmarkQueryModel model) {
    List<Long> counts = Lists.newArrayList(model.simulator().resultsPerWindow());
    Collections.sort(counts);
    int n = counts.size();
    if (n < 5) {
      NexmarkUtils.console("Query%s: only %d samples", model.configuration.query, n);
    } else {
      NexmarkUtils.console(
          "Query%d: N:%d; min:%d; 1st%%:%d; mean:%d; 3rd%%:%d; max:%d",
          model.configuration.query,
          n,
          counts.get(0),
          counts.get(n / 4),
          counts.get(n / 2),
          counts.get(n - 1 - n / 4),
          counts.get(n - 1));
    }
  }

  /** Run {@code configuration} and return its performance if possible. */
  @Nullable
  public NexmarkPerf run() throws IOException {
    if (options.getManageResources() && !options.getMonitorJobs()) {
      throw new RuntimeException("If using --manageResources then must also use --monitorJobs.");
    }

    //
    // Setup per-run state.
    //
    checkState(queryName == null);
    if (configuration.sourceType.equals(SourceType.PUBSUB)) {
      pubsubHelper = PubsubHelper.create(options);
    }

    try {
      NexmarkUtils.console("Running %s", configuration.toShortString());

      if (configuration.numEvents < 0) {
        NexmarkUtils.console("skipping since configuration is disabled");
        return null;
      }

      NexmarkQuery<? extends KnownSize> query = getNexmarkQuery();
      if (query == null) {
        NexmarkUtils.console("skipping since configuration is not implemented");
        return null;
      }

      queryName = query.getName();

      // Append queryName to temp location
      if (!"".equals(options.getTempLocation())) {
        options.setTempLocation(options.getTempLocation() + "/" + queryName);
      }

      NexmarkQueryModel model = getNexmarkQueryModel();

      if (options.getJustModelResultRate()) {
        if (model == null) {
          throw new RuntimeException(String.format("No model for %s", queryName));
        }
        modelResultRates(model);
        return null;
      }

      final Instant now = Instant.now();
      Pipeline p = Pipeline.create(options);
      NexmarkUtils.setupPipeline(configuration.coderStrategy, p);

      // Generate events.
      PCollection<Event> source = createSource(p, now);

      if (query.getTransform().needsSideInput()) {
        query.getTransform().setSideInput(NexmarkUtils.prepareSideInput(p, configuration));
      }

      if (options.getLogEvents()) {
        source = source.apply(queryName + ".Events.Log", NexmarkUtils.log(queryName + ".Events"));
      }

      // Source will be null if source type is PUBSUB and mode is PUBLISH_ONLY.
      // In that case there's nothing more to add to pipeline.
      if (source != null) {
        // Optionally sink events in Avro format.
        // (Query results are ignored).
        if (configuration.sinkType == NexmarkUtils.SinkType.AVRO) {
          sinkEventsToAvro(source);
        }

        // Query 10 logs all events to Google Cloud storage files. It could generate a lot of logs,
        // so, set parallelism. Also set the output path where to write log files.
        if (configuration.query == NexmarkQueryName.LOG_TO_SHARDED_FILES) {
          String path = null;
          if (options.getOutputPath() != null && !options.getOutputPath().isEmpty()) {
            path = logsDir(now.getMillis());
          }
          ((Query10) query.getTransform()).setOutputPath(path);
          ((Query10) query.getTransform()).setMaxNumWorkers(maxNumWorkers());
        }

        // Apply query.
        PCollection<TimestampedValue<KnownSize>> results =
            (PCollection<TimestampedValue<KnownSize>>) source.apply(query);

        if (options.getAssertCorrectness()) {
          if (model == null) {
            throw new RuntimeException(String.format("No model for %s", queryName));
          }
          // We know all our streams have a finite number of elements.
          results.setIsBoundedInternal(PCollection.IsBounded.BOUNDED);
          // If we have a finite number of events then assert our pipeline's
          // results match those of a model using the same sequence of events.
          PAssert.that(results).satisfies(model.assertionFor());
        }

        // Output results.
        sink(results, now.getMillis());
      }

      mainResult = p.run();
      mainResult.waitUntilFinish(Duration.standardSeconds(configuration.streamTimeout));
      return monitor(query);
    } finally {
      if (pubsubHelper != null) {
        pubsubHelper.cleanup();
        pubsubHelper = null;
      }
      configuration = null;
      queryName = null;
    }
  }

  private boolean isSql() {
    return SQL.equalsIgnoreCase(options.getQueryLanguage());
  }

  private NexmarkQueryModel getNexmarkQueryModel() {
    Map<NexmarkQueryName, NexmarkQueryModel> models = createQueryModels();
    return models.get(configuration.query);
  }

  private NexmarkQuery<?> getNexmarkQuery() {
    Map<NexmarkQueryName, NexmarkQuery> queries = createQueries();
    return queries.get(configuration.query);
  }

  private Map<NexmarkQueryName, NexmarkQueryModel> createQueryModels() {
    return isSql() ? createSqlQueryModels() : createJavaQueryModels();
  }

  private Map<NexmarkQueryName, NexmarkQueryModel> createSqlQueryModels() {
    return ImmutableMap.of();
  }

  private Map<NexmarkQueryName, NexmarkQueryModel> createJavaQueryModels() {
    return ImmutableMap.<NexmarkQueryName, NexmarkQueryModel>builder()
        .put(NexmarkQueryName.PASSTHROUGH, new Query0Model(configuration))
        .put(NexmarkQueryName.CURRENCY_CONVERSION, new Query1Model(configuration))
        .put(NexmarkQueryName.SELECTION, new Query2Model(configuration))
        .put(NexmarkQueryName.LOCAL_ITEM_SUGGESTION, new Query3Model(configuration))
        .put(NexmarkQueryName.AVERAGE_PRICE_FOR_CATEGORY, new Query4Model(configuration))
        .put(NexmarkQueryName.HOT_ITEMS, new Query5Model(configuration))
        .put(NexmarkQueryName.AVERAGE_SELLING_PRICE_BY_SELLER, new Query6Model(configuration))
        .put(NexmarkQueryName.HIGHEST_BID, new Query7Model(configuration))
        .put(NexmarkQueryName.MONITOR_NEW_USERS, new Query8Model(configuration))
        .put(NexmarkQueryName.WINNING_BIDS, new Query9Model(configuration))
        .put(NexmarkQueryName.BOUNDED_SIDE_INPUT_JOIN, new BoundedSideInputJoinModel(configuration))
        .put(NexmarkQueryName.SESSION_SIDE_INPUT_JOIN, new SessionSideInputJoinModel(configuration))
        .build();
  }

  private Map<NexmarkQueryName, NexmarkQuery> createQueries() {
    return isSql() ? createSqlQueries() : createJavaQueries();
  }

  private Map<NexmarkQueryName, NexmarkQuery> createSqlQueries() {
    return ImmutableMap.<NexmarkQueryName, NexmarkQuery>builder()
        .put(NexmarkQueryName.PASSTHROUGH, new NexmarkQuery(configuration, new SqlQuery0()))
        .put(NexmarkQueryName.CURRENCY_CONVERSION, new NexmarkQuery(configuration, new SqlQuery1()))
        .put(
            NexmarkQueryName.SELECTION,
            new NexmarkQuery(configuration, new SqlQuery2(configuration.auctionSkip)))
        .put(
            NexmarkQueryName.LOCAL_ITEM_SUGGESTION,
            new NexmarkQuery(configuration, new SqlQuery3(configuration)))

        // SqlQuery5 is disabled for now, uses non-equi-joins,
        // never worked right, was giving incorrect results.
        // Gets rejected after PR/8301, causing failures.
        //
        // See:
        //   https://issues.apache.org/jira/browse/BEAM-7072
        //   https://github.com/apache/beam/pull/8301
        //   https://github.com/apache/beam/pull/8422#issuecomment-487676350
        //
        //        .put(
        //            NexmarkQueryName.HOT_ITEMS,
        //            new NexmarkQuery(configuration, new SqlQuery5(configuration)))
        .put(
            NexmarkQueryName.HIGHEST_BID,
            new NexmarkQuery(configuration, new SqlQuery7(configuration)))
        .put(
            NexmarkQueryName.BOUNDED_SIDE_INPUT_JOIN,
            new NexmarkQuery(configuration, new SqlBoundedSideInputJoin(configuration)))
        .build();
  }

  private Map<NexmarkQueryName, NexmarkQuery> createJavaQueries() {
    return ImmutableMap.<NexmarkQueryName, NexmarkQuery>builder()
        .put(NexmarkQueryName.PASSTHROUGH, new NexmarkQuery(configuration, new Query0()))
        .put(
            NexmarkQueryName.CURRENCY_CONVERSION,
            new NexmarkQuery(configuration, new Query1(configuration)))
        .put(NexmarkQueryName.SELECTION, new NexmarkQuery(configuration, new Query2(configuration)))
        .put(
            NexmarkQueryName.LOCAL_ITEM_SUGGESTION,
            new NexmarkQuery(configuration, new Query3(configuration)))
        .put(
            NexmarkQueryName.AVERAGE_PRICE_FOR_CATEGORY,
            new NexmarkQuery(configuration, new Query4(configuration)))
        .put(NexmarkQueryName.HOT_ITEMS, new NexmarkQuery(configuration, new Query5(configuration)))
        .put(
            NexmarkQueryName.AVERAGE_SELLING_PRICE_BY_SELLER,
            new NexmarkQuery(configuration, new Query6(configuration)))
        .put(
            NexmarkQueryName.HIGHEST_BID,
            new NexmarkQuery(configuration, new Query7(configuration)))
        .put(
            NexmarkQueryName.MONITOR_NEW_USERS,
            new NexmarkQuery(configuration, new Query8(configuration)))
        .put(
            NexmarkQueryName.WINNING_BIDS,
            new NexmarkQuery(configuration, new Query9(configuration)))
        .put(
            NexmarkQueryName.LOG_TO_SHARDED_FILES,
            new NexmarkQuery(configuration, new Query10(configuration)))
        .put(
            NexmarkQueryName.USER_SESSIONS,
            new NexmarkQuery(configuration, new Query11(configuration)))
        .put(
            NexmarkQueryName.PROCESSING_TIME_WINDOWS,
            new NexmarkQuery(configuration, new Query12(configuration)))
        .put(
            NexmarkQueryName.BOUNDED_SIDE_INPUT_JOIN,
            new NexmarkQuery(configuration, new BoundedSideInputJoin(configuration)))
        .put(
            NexmarkQueryName.SESSION_SIDE_INPUT_JOIN,
            new NexmarkQuery(configuration, new SessionSideInputJoin(configuration)))
        .build();
  }

  private static class PubsubMessageEventDoFn extends DoFn<PubsubMessage, Event> {
    @ProcessElement
    public void processElement(ProcessContext c) throws IOException {
      byte[] payload = c.element().getPayload();
      Event event = CoderUtils.decodeFromByteArray(Event.CODER, payload);
      c.output(event);
    }
  }

  private static class EventPubsubMessageDoFn extends DoFn<Event, PubsubMessage> {
    @ProcessElement
    public void processElement(ProcessContext c) throws IOException {
      byte[] payload = CoderUtils.encodeToByteArray(Event.CODER, c.element());
      c.output(new PubsubMessage(payload, Collections.emptyMap()));
    }
  }
}
