sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamSetOperatorsTransforms.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.beam.sdk.extensions.sql.impl.transform;

import org.apache.beam.sdk.extensions.sql.impl.rel.BeamSetOperatorRelBase;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.SimpleFunction;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;

/** Collections of {@code PTransform} and {@code DoFn} used to perform Set operations. */
@SuppressWarnings({
  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
public abstract class BeamSetOperatorsTransforms {
  /** Transform a {@code BeamSqlRow} to a {@code KV<BeamSqlRow, BeamSqlRow>}. */
  public static class BeamSqlRow2KvFn extends SimpleFunction<Row, KV<Row, Row>> {
    @Override
    public KV<Row, Row> apply(Row input) {
      return KV.of(input, input);
    }
  }

  /** Filter function used for Set operators. */
  public static class SetOperatorFilteringDoFn extends DoFn<Row, Row> {
    private final String leftTag;
    private final String rightTag;
    private final BeamSetOperatorRelBase.OpType opType;
    // ALL?
    private final boolean all;

    public SetOperatorFilteringDoFn(
        String leftTag, String rightTag, BeamSetOperatorRelBase.OpType opType, boolean all) {
      this.leftTag = leftTag;
      this.rightTag = rightTag;
      this.opType = opType;
      this.all = all;
    }

    @ProcessElement
    public void processElement(@Element Row element, OutputReceiver<Row> o) {
      Row key = element.getRow("key");
      long numLeftRows = 0;
      long numRightRows = 0;
      if (!Iterables.isEmpty(element.<Row>getIterable(leftTag))) {
        numLeftRows = Iterables.size(element.<Row>getIterable(leftTag));
      }
      if (!Iterables.isEmpty(element.<Row>getIterable(rightTag))) {
        numRightRows = Iterables.size(element.<Row>getIterable(rightTag));
      }

      switch (opType) {
        case UNION:
          if (all) {
            for (int i = 0; i < numLeftRows + numRightRows; i++) {
              o.output(key);
            }
          } else {
            // only output the key
            o.output(key);
          }
          break;
        case INTERSECT:
          if (numLeftRows > 0 && numRightRows > 0) {
            if (all) {
              // Say for Row R, there are m instances on left and n instances on right,
              // INTERSECT ALL outputs MIN(m, n) instances of R.
              for (int i = 0; i < Math.min(numLeftRows, numRightRows); i++) {
                o.output(key);
              }
            } else {
              o.output(key);
            }
          }
          break;
        case MINUS:
          // Say for Row R, there are m instances on left and n instances on right:
          // - EXCEPT ALL outputs MAX(m - n, 0) instances of R.
          // - EXCEPT [DISTINCT] outputs a single instance of R if m > 0 and n == 0, else
          //   they output 0 instances.
          if (numLeftRows > 0 && numRightRows == 0) {
            if (all) {
              // output all
              for (int i = 0; i < numLeftRows; i++) {
                o.output(key);
              }
            } else {
              // only output one
              o.output(key);
            }
          } else if (numLeftRows > 0 && numRightRows > 0) {
            long outputCount = numLeftRows - numRightRows;
            if (outputCount > 0) {
              if (all) {
                while (outputCount > 0) {
                  outputCount--;
                  o.output(key);
                }
              }
              // Dont output any in DISTINCT (if (!all)) case
            }
          }
      }
    }
  }
}