linq4j/src/test/java/org/apache/calcite/linq4j/test/CorrelateJoinTest.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.calcite.linq4j.test;

import org.apache.calcite.linq4j.Enumerable;
import org.apache.calcite.linq4j.Enumerator;
import org.apache.calcite.linq4j.ExtendedEnumerable;
import org.apache.calcite.linq4j.JoinType;
import org.apache.calcite.linq4j.Linq4j;
import org.apache.calcite.linq4j.function.Function1;
import org.apache.calcite.linq4j.function.Function2;

import com.google.common.collect.ImmutableList;

import org.junit.jupiter.api.Test;

import java.util.ArrayList;
import java.util.List;

import static org.junit.jupiter.api.Assertions.assertArrayEquals;

/**
 * Tests {@link ExtendedEnumerable#correlateJoin(JoinType, Function1, Function2)}.
 */
class CorrelateJoinTest {
  static final Function2<Integer, Integer, Integer[]> SELECT_BOTH =
      (v0, v1) -> new Integer[]{v0, v1};

  @Test void testInner() {
    testJoin(JoinType.INNER, new Integer[][]{
        {2, 20},
        {3, -30},
        {3, -60},
        {20, 200},
        {30, -300},
        {30, -600}});
  }

  @Test void testLeft() {
    testJoin(JoinType.LEFT, new Integer[][]{
        {1, null},
        {2, 20},
        {3, -30},
        {3, -60},
        {10, null},
        {20, 200},
        {30, -300},
        {30, -600}});
  }

  @Test void testSemi() {
    testJoin(JoinType.SEMI, new Integer[][]{
        {2, null},
        {3, null},
        {20, null},
        {30, null}});
  }

  @Test void testAnti() {
    testJoin(JoinType.ANTI, new Integer[][]{
        {1, null},
        {10, null}});
  }

  public void testJoin(JoinType joinType, Integer[][] expected) {
    Enumerable<Integer[]> join =
        Linq4j.asEnumerable(ImmutableList.of(1, 2, 3, 10, 20, 30))
            .correlateJoin(joinType, a0 -> {
              if (a0 == 1 || a0 == 10) {
                return Linq4j.emptyEnumerable();
              }
              if (a0 == 2 || a0 == 20) {
                return Linq4j.singletonEnumerable(a0 * 10);
              }
              if (a0 == 3 || a0 == 30) {
                return Linq4j.asEnumerable(
                    ImmutableList.of(-a0 * 10, -a0 * 20));
              }
              throw new IllegalArgumentException(
                  "Unexpected input " + a0);
            }, SELECT_BOTH);
    for (int i = 0; i < 2; i++) {
      Enumerator<Integer[]> e = join.enumerator();
      checkResults(e, expected);
      e.close();
    }
  }

  private void checkResults(Enumerator<Integer[]> e, Integer[][] expected) {
    List<Integer[]> res = new ArrayList<>();
    while (e.moveNext()) {
      res.add(e.current());
    }
    Integer[][] actual = res.toArray(new Integer[res.size()][]);
    assertArrayEquals(expected, actual);
  }
}