content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.beam.sdk.extensions.sql.impl.rel;
  
  import static org.apache.beam.sdk.extensions.sql.impl.cep.CEPUtils.makeOrderKeysFromCollation;
  import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
  
  import java.util.ArrayList;
  import java.util.List;
  import java.util.Map;
  import java.util.SortedSet;
  import org.apache.beam.sdk.coders.KvCoder;
  import org.apache.beam.sdk.coders.RowCoder;
  import org.apache.beam.sdk.extensions.sql.impl.cep.CEPCall;
  import org.apache.beam.sdk.extensions.sql.impl.cep.CEPFieldRef;
  import org.apache.beam.sdk.extensions.sql.impl.cep.CEPKind;
  import org.apache.beam.sdk.extensions.sql.impl.cep.CEPLiteral;
  import org.apache.beam.sdk.extensions.sql.impl.cep.CEPMeasure;
  import org.apache.beam.sdk.extensions.sql.impl.cep.CEPOperation;
  import org.apache.beam.sdk.extensions.sql.impl.cep.CEPPattern;
  import org.apache.beam.sdk.extensions.sql.impl.cep.CEPUtils;
  import org.apache.beam.sdk.extensions.sql.impl.cep.OrderKey;
  import org.apache.beam.sdk.extensions.sql.impl.nfa.NFA;
  import org.apache.beam.sdk.extensions.sql.impl.planner.BeamCostModel;
  import org.apache.beam.sdk.extensions.sql.impl.planner.BeamRelMetadataQuery;
  import org.apache.beam.sdk.extensions.sql.impl.planner.NodeStats;
  import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils;
  import org.apache.beam.sdk.schemas.Schema;
  import org.apache.beam.sdk.transforms.DoFn;
  import org.apache.beam.sdk.transforms.GroupByKey;
  import org.apache.beam.sdk.transforms.PTransform;
  import org.apache.beam.sdk.transforms.ParDo;
  import org.apache.beam.sdk.values.KV;
  import org.apache.beam.sdk.values.PCollection;
  import org.apache.beam.sdk.values.PCollectionList;
  import org.apache.beam.sdk.values.Row;
  import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptCluster;
  import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptPlanner;
  import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelTraitSet;
  import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.RelCollation;
  import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.RelNode;
  import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.core.Match;
  import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.type.RelDataType;
  import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rex.RexCall;
  import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rex.RexNode;
  import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlKind;
  import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.util.ImmutableBitSet;
  import org.checkerframework.checker.nullness.qual.Nullable;
  
  /**
   * {@code BeamRelNode} to replace a {@code Match} node.
   *
   * <p>The {@code BeamMatchRel} is the Beam implementation of {@code MATCH_RECOGNIZE} in SQL.
   *
   * <p>For now, the underline implementation is based on java.util.regex.
   */
  @SuppressWarnings({
    "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  public class BeamMatchRel extends Match implements BeamRelNode {
  
    public BeamMatchRel(
        RelOptCluster cluster,
        RelTraitSet traitSet,
        RelNode input,
        RelDataType rowType,
        RexNode pattern,
        boolean strictStart,
        boolean strictEnd,
        Map<String, RexNode> patternDefinitions,
        Map<String, RexNode> measures,
        RexNode after,
        Map<String, ? extends SortedSet<String>> subsets,
        boolean allRows,
        ImmutableBitSet partitionKeys,
        RelCollation orderKeys,
        @Nullable RexNode interval) {
  
      super(
          cluster,
          traitSet,
          input,
          rowType,
          pattern,
          strictStart,
          strictEnd,
          patternDefinitions,
          measures,
          after,
          subsets,
          allRows,
          partitionKeys,
          orderKeys,
          interval);
    }
  
    @Override
    public BeamCostModel beamComputeSelfCost(RelOptPlanner planner, BeamRelMetadataQuery mq) {
      return BeamCostModel.FACTORY.makeTinyCost(); // return constant costModel for now
    }
  
    @Override
    public NodeStats estimateNodeStats(BeamRelMetadataQuery mq) {
      // a simple way of getting some estimate data
      // to be examined further
      NodeStats inputEstimate = BeamSqlRelUtils.getNodeStats(input, mq);
      double numRows = inputEstimate.getRowCount();
      double winSize = inputEstimate.getWindow();
      double rate = inputEstimate.getRate();
  
      return NodeStats.create(numRows, rate, winSize).multiply(0.5);
    }
  
    @Override
    public PTransform<PCollectionList<Row>, PCollection<Row>> buildPTransform() {
      return new MatchTransform(
          partitionKeys, orderKeys, measures, allRows, pattern, patternDefinitions);
    }
  
    private class MatchTransform extends PTransform<PCollectionList<Row>, PCollection<Row>> {
  
      private final ImmutableBitSet partitionKeys;
      private final RelCollation orderKeys;
      private final Map<String, RexNode> measures;
      private final boolean allRows;
      private final RexNode pattern;
      private final Map<String, RexNode> patternDefs;
  
      public MatchTransform(
          ImmutableBitSet partitionKeys,
          RelCollation orderKeys,
          Map<String, RexNode> measures,
          boolean allRows,
          RexNode pattern,
          Map<String, RexNode> patternDefs) {
        this.partitionKeys = partitionKeys;
        this.orderKeys = orderKeys;
        this.measures = measures;
        this.allRows = allRows;
        this.pattern = pattern;
        this.patternDefs = patternDefs;
      }
  
      @Override
      public PCollection<Row> expand(PCollectionList<Row> pinput) {
        checkArgument(
            pinput.size() == 1,
            "Wrong number of inputs for %s: %s",
            BeamMatchRel.class.getSimpleName(),
            pinput);
        PCollection<Row> upstream = pinput.get(0);
        Schema upstreamSchema = upstream.getSchema();
        Schema outSchema = CalciteUtils.toSchema(getRowType());
  
        Schema.Builder schemaBuilder = new Schema.Builder();
        for (int index : partitionKeys.asList()) {
          schemaBuilder.addField(upstreamSchema.getField(index));
        }
        Schema partitionKeySchema = schemaBuilder.build();
  
        // partition according to the partition keys
        PCollection<KV<Row, Row>> keyedUpstream =
            upstream.apply(ParDo.of(new MapKeys(partitionKeySchema)));
  
        // group by keys
        PCollection<KV<Row, Iterable<Row>>> groupedUpstream =
            keyedUpstream
                .setCoder(KvCoder.of(RowCoder.of(partitionKeySchema), RowCoder.of(upstreamSchema)))
                .apply(GroupByKey.create());
  
        // sort within each keyed partition
        ArrayList<OrderKey> orderKeyList = makeOrderKeysFromCollation(orderKeys);
        // This will rely on an assumption that Fusion will fuse
        // operators here so the sorted result will be preserved
        // for the next match transform.
        // In most of the runners (if not all) this should be true.
        PCollection<KV<Row, Iterable<Row>>> orderedUpstream =
            groupedUpstream.apply(ParDo.of(new SortPerKey(orderKeyList)));
  
        // apply the pattern match in each partition
        ArrayList<CEPPattern> cepPattern =
            CEPUtils.getCEPPatternFromPattern(upstreamSchema, pattern, patternDefs);
        List<CEPMeasure> cepMeasures = new ArrayList<>();
        for (Map.Entry<String, RexNode> i : measures.entrySet()) {
          String outTableName = i.getKey();
          CEPOperation measureOperation;
          // TODO: support FINAL clause, for now, get rid of the FINAL operation
          if (i.getValue().getClass() == RexCall.class) {
            RexCall rexCall = (RexCall) i.getValue();
            if (rexCall.getOperator().getKind() == SqlKind.FINAL) {
              measureOperation = CEPOperation.of(rexCall.getOperands().get(0));
              cepMeasures.add(new CEPMeasure(upstreamSchema, outTableName, measureOperation));
              continue;
            }
          }
          measureOperation = CEPOperation.of(i.getValue());
          cepMeasures.add(new CEPMeasure(upstreamSchema, outTableName, measureOperation));
        }
  
        // apply the ParDo for the match process and measures clause
        // for now, support FINAL only
        // TODO: add support for FINAL/RUNNING
        List<CEPFieldRef> cepParKeys = CEPUtils.getCEPFieldRefFromParKeys(partitionKeys);
        PCollection<Row> outStream =
            orderedUpstream
                .apply(
                    ParDo.of(
                        new MatchPattern(
                            upstreamSchema, cepParKeys, cepPattern, cepMeasures, allRows, outSchema)))
                .setRowSchema(outSchema);
  
        return outStream;
      }
    }
  
    // TODO: support both ALL ROWS PER MATCH and ONE ROW PER MATCH.
    // support only one row per match for now.
    private static class MatchPattern extends DoFn<KV<Row, Iterable<Row>>, Row> {
  
      private final Schema upstreamSchema;
      private final Schema outSchema;
      private final List<CEPFieldRef> partitionKeys;
      private final ArrayList<CEPPattern> pattern;
      private final List<CEPMeasure> measures;
      private final boolean allRows;
  
      MatchPattern(
          Schema upstreamSchema,
          List<CEPFieldRef> partitionKeys,
          ArrayList<CEPPattern> pattern,
          List<CEPMeasure> measures,
          boolean allRows,
          Schema outSchema) {
        this.upstreamSchema = upstreamSchema;
        this.partitionKeys = partitionKeys;
        this.pattern = pattern;
        this.measures = measures;
        this.allRows = allRows;
        this.outSchema = outSchema;
      }
  
      @ProcessElement
      public void processElement(@Element KV<Row, Iterable<Row>> keyRows, OutputReceiver<Row> out) {
        NFA partNFA = NFA.compile(pattern, upstreamSchema);
        Iterable<Row> partRows = keyRows.getValue();
        Map<String, ArrayList<Row>> result;
  
        for (Row singleRow : partRows) {
          // output each matched sequence as specified by the Measure clause
          result = partNFA.processNewRow(singleRow);
          if (result == null) {
            // if there isn't match
            continue;
          }
  
          if (allRows) {
            for (ArrayList<Row> i : result.values()) {
              for (Row j : i) {
                out.output(j);
              }
            }
          } else {
  
            // output corresponding columns according to the measures schema
            Row.Builder newRowBuilder = Row.withSchema(outSchema);
            Row.FieldValueBuilder newFieldBuilder = null;
  
            // add partition key columns
            for (CEPFieldRef i : partitionKeys) {
              int colIndex = i.getIndex();
              Schema.Field parSchema = upstreamSchema.getField(colIndex);
              String fieldName = parSchema.getName();
              if (!result.isEmpty()) {
                Row partitionKeyRow = keyRows.getKey();
                if (newFieldBuilder == null) {
                  newFieldBuilder =
                      newRowBuilder.withFieldValue(fieldName, partitionKeyRow.getValue(fieldName));
                } else {
                  newFieldBuilder =
                      newFieldBuilder.withFieldValue(fieldName, partitionKeyRow.getValue(fieldName));
                }
              } else {
                break;
              }
            }
  
            // add measure columns
            for (CEPMeasure i : measures) {
              String outName = i.getName();
              CEPFieldRef patternRef = i.getField();
              String patternVar = patternRef.getAlpha();
              List<Row> patternRows = result.get(patternVar);
  
              // implement CEPOperation as functions
              CEPOperation opr = i.getOperation();
              if (opr.getClass() == CEPCall.class) {
                CEPCall call = (CEPCall) opr;
                CEPKind funcName = call.getOperator().getCepKind();
                switch (funcName) {
                  case FIRST:
                    CEPFieldRef colFirstField = (CEPFieldRef) call.getOperands().get(0);
                    CEPLiteral colFirstIndex = (CEPLiteral) call.getOperands().get(1);
                    Row rowFirstToProc = patternRows.get(colFirstIndex.getDecimal().intValue());
                    if (newFieldBuilder == null) {
                      newFieldBuilder =
                          newRowBuilder.withFieldValue(
                              outName, rowFirstToProc.getValue(colFirstField.getIndex()));
                    } else {
                      newFieldBuilder =
                          newFieldBuilder.withFieldValue(
                              outName, rowFirstToProc.getValue(colFirstField.getIndex()));
                    }
                    break;
                  case LAST:
                    CEPFieldRef colLastField = (CEPFieldRef) call.getOperands().get(0);
                    CEPLiteral colLastIndex = (CEPLiteral) call.getOperands().get(1);
                    Row rowLastToProc =
                        patternRows.get(
                            patternRows.size() - 1 - colLastIndex.getDecimal().intValue());
                    if (newFieldBuilder == null) {
                      newFieldBuilder =
                          newRowBuilder.withFieldValue(
                              outName, rowLastToProc.getValue(colLastField.getIndex()));
                    } else {
                      newFieldBuilder =
                          newFieldBuilder.withFieldValue(
                              outName, rowLastToProc.getValue(colLastField.getIndex()));
                    }
                    break;
                  default:
                    throw new UnsupportedOperationException(
                        "The measure function is not recognized: " + funcName.name());
                }
              } else if (opr.getClass() == CEPFieldRef.class) {
                Row rowToProc = patternRows.get(0);
                CEPFieldRef fieldRef = (CEPFieldRef) opr;
                if (newFieldBuilder == null) {
                  newFieldBuilder =
                      newRowBuilder.withFieldValue(outName, rowToProc.getValue(fieldRef.getIndex()));
                } else {
                  newFieldBuilder =
                      newFieldBuilder.withFieldValue(
                          outName, rowToProc.getValue(fieldRef.getIndex()));
                }
              } else {
                throw new UnsupportedOperationException(
                    "CEP operation is not recognized: " + opr.getClass().getName());
              }
            }
            Row newRow;
            if (newFieldBuilder == null) {
              newRow = newRowBuilder.build();
            } else {
              newRow = newFieldBuilder.build();
            }
            out.output(newRow);
          }
        }
      }
    }
  
    private static class SortPerKey extends DoFn<KV<Row, Iterable<Row>>, KV<Row, Iterable<Row>>> {
  
      private final ArrayList<OrderKey> orderKeys;
  
      public SortPerKey(ArrayList<OrderKey> orderKeys) {
        this.orderKeys = orderKeys;
      }
  
      @ProcessElement
      public void processElement(
          @Element KV<Row, Iterable<Row>> keyRows, OutputReceiver<KV<Row, Iterable<Row>>> out) {
        ArrayList<Row> rows = new ArrayList<>();
        for (Row i : keyRows.getValue()) {
          rows.add(i);
        }
  
        ArrayList<Integer> fIndexList = new ArrayList<>();
        ArrayList<Boolean> dirList = new ArrayList<>();
        ArrayList<Boolean> nullDirList = new ArrayList<>();
  
        // reversely traverse the order key list
        for (int i = (orderKeys.size() - 1); i >= 0; --i) {
          OrderKey thisKey = orderKeys.get(i);
          fIndexList.add(thisKey.getIndex());
          dirList.add(thisKey.getDir());
          nullDirList.add(thisKey.getNullFirst());
        }
  
        rows.sort(new BeamSortRel.BeamSqlRowComparator(fIndexList, dirList, nullDirList));
  
        out.output(KV.of(keyRows.getKey(), rows));
      }
    }
  
    private static class MapKeys extends DoFn<Row, KV<Row, Row>> {
  
      private final Schema partitionKeySchema;
  
      public MapKeys(Schema partitionKeySchema) {
        this.partitionKeySchema = partitionKeySchema;
      }
  
      @ProcessElement
      public void processElement(@Element Row eleRow, OutputReceiver<KV<Row, Row>> out) {
        Row.Builder newRowBuilder = Row.withSchema(partitionKeySchema);
  
        // no partition specified would result in empty row as keys for rows
        for (Schema.Field i : partitionKeySchema.getFields()) {
          String fieldName = i.getName();
          newRowBuilder.addValue(eleRow.getValue(fieldName));
        }
        KV kvPair = KV.of(newRowBuilder.build(), eleRow);
        out.output(kvPair);
      }
    }
  
    public Match copy(
        RelNode input,
        RelDataType rowType,
        RexNode pattern,
        boolean strictStart,
        boolean strictEnd,
        Map<String, RexNode> patternDefinitions,
        Map<String, RexNode> measures,
        RexNode after,
        Map<String, ? extends SortedSet<String>> subsets,
        boolean allRows,
        ImmutableBitSet partitionKeys,
        RelCollation orderKeys,
        RexNode interval) {
  
      return new BeamMatchRel(
          getCluster(),
          getTraitSet(),
          input,
          rowType,
          pattern,
          strictStart,
          strictEnd,
          patternDefinitions,
          measures,
          after,
          subsets,
          allRows,
          partitionKeys,
          orderKeys,
          interval);
    }
  
    @Override
    public RelNode copy(RelTraitSet traitSet, List<RelNode> inputs) {
      return new BeamMatchRel(
          getCluster(),
          traitSet,
          inputs.get(0),
          rowType,
          pattern,
          strictStart,
          strictEnd,
          patternDefinitions,
          measures,
          after,
          subsets,
          allRows,
          partitionKeys,
          orderKeys,
          interval);
    }
  }
  