Add tests for JoinableFactoryWrapper

apache · May 14, 2022 · 439c959 · 439c959
1 parent 0e7a6e2
commit 439c959
Showing 1 changed file with 99 additions and 0 deletions.
diff --git a/processing/src/test/java/org/apache/druid/segment/join/JoinableFactoryWrapperTest.java b/processing/src/test/java/org/apache/druid/segment/join/JoinableFactoryWrapperTest.java
@@ -26,13 +26,15 @@
 import com.google.common.collect.Sets;
 import org.apache.druid.common.config.NullHandling;
 import org.apache.druid.common.config.NullHandlingTest;
+import org.apache.druid.java.util.common.DateTimes;
 import org.apache.druid.java.util.common.IAE;
 import org.apache.druid.java.util.common.Intervals;
 import org.apache.druid.java.util.common.Pair;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.math.expr.ExprMacroTable;
 import org.apache.druid.query.DataSource;
 import org.apache.druid.query.GlobalTableDataSource;
+import org.apache.druid.query.InlineDataSource;
 import org.apache.druid.query.LookupDataSource;
 import org.apache.druid.query.TableDataSource;
 import org.apache.druid.query.TestQuery;
@@ -45,7 +47,12 @@
 import org.apache.druid.query.planning.PreJoinableClause;
 import org.apache.druid.query.spec.MultipleIntervalSegmentSpec;
 import org.apache.druid.segment.SegmentReference;
+import org.apache.druid.segment.column.ColumnType;
+import org.apache.druid.segment.column.RowSignature;
 import org.apache.druid.segment.join.lookup.LookupJoinable;
+import org.apache.druid.segment.join.table.IndexedTable;
+import org.apache.druid.segment.join.table.IndexedTableJoinable;
+import org.apache.druid.segment.join.table.RowBasedIndexedTable;
 import org.easymock.EasyMock;
 import org.junit.Assert;
 import org.junit.Rule;
@@ -61,6 +68,7 @@
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Function;
+import java.util.stream.Collectors;
 
 public class JoinableFactoryWrapperTest extends NullHandlingTest
 {
@@ -81,6 +89,25 @@ public class JoinableFactoryWrapperTest extends NullHandlingTest
       ? TEST_LOOKUP.keySet()
       : Sets.difference(TEST_LOOKUP.keySet(), Collections.singleton(""));
 
+  private static final InlineDataSource INDEXED_TABLE_DS = InlineDataSource.fromIterable(
+      ImmutableList.of(
+          new Object[]{"Mexico"},
+          new Object[]{"Norway"},
+          new Object[]{"El Salvador"},
+          new Object[]{"United States"},
+          new Object[]{"United States"}
+      ),
+      RowSignature.builder().add("country", ColumnType.STRING).build()
+  );
+
+  private static final IndexedTable TEST_INDEXED_TABLE = new RowBasedIndexedTable<>(
+      INDEXED_TABLE_DS.getRowsAsList(),
+      INDEXED_TABLE_DS.rowAdapter(),
+      INDEXED_TABLE_DS.getRowSignature(),
+      ImmutableSet.of("country"),
+      DateTimes.nowUtc().toString()
+  );
+
   @Rule
   public ExpectedException expectedException = ExpectedException.none();
 
@@ -467,6 +494,33 @@ public void test_convertJoinsToFilters_convertInnerJoin()
     );
   }
 
+  @Test
+  public void test_convertJoinsToPartialFilters_convertInnerJoin()
+  {
+    JoinableClause joinableClause = new JoinableClause(
+        "j.",
+        new IndexedTableJoinable(TEST_INDEXED_TABLE),
+        JoinType.INNER,
+        JoinConditionAnalysis.forExpression("x == \"j.country\"", "j.", ExprMacroTable.nil())
+    );
+    final Pair<List<Filter>, List<JoinableClause>> conversion = JoinableFactoryWrapper.convertJoinsToFilters(
+        ImmutableList.of(joinableClause),
+        ImmutableSet.of("x"),
+        Integer.MAX_VALUE
+    );
+
+    Assert.assertEquals(
+        Pair.of(
+            ImmutableList.of(new InDimFilter(
+                "x",
+                INDEXED_TABLE_DS.getRowsAsList().stream().map(row -> row[0].toString()).collect(Collectors.toSet()))
+            ),
+            ImmutableList.of(joinableClause) // the joinable clause remains intact since we've duplicates in country column
+        ),
+        conversion
+    );
+  }
+
   @Test
   public void test_convertJoinsToFilters_convertTwoInnerJoins()
   {
@@ -506,6 +560,51 @@ public void test_convertJoinsToFilters_convertTwoInnerJoins()
     );
   }
 
+  @Test
+  public void test_convertJoinsToPartialAndFullFilters_convertMultipleInnerJoins()
+  {
+    final ImmutableList<JoinableClause> clauses = ImmutableList.of(
+        new JoinableClause(
+            "j.",
+            LookupJoinable.wrap(new MapLookupExtractor(TEST_LOOKUP, false)),
+            JoinType.INNER,
+            JoinConditionAnalysis.forExpression("x == \"j.k\"", "j.", ExprMacroTable.nil())
+        ), // this joinable will be fully converted to a filter
+        new JoinableClause(
+            "_j.",
+            new IndexedTableJoinable(TEST_INDEXED_TABLE),
+            JoinType.INNER,
+            JoinConditionAnalysis.forExpression("x == \"_j.country\"", "_j.", ExprMacroTable.nil())
+        ), // this joinable will be partially converted to a filter since we've duplicates on country column
+        new JoinableClause(
+            "__j.",
+            new IndexedTableJoinable(TEST_INDEXED_TABLE),
+            JoinType.LEFT,
+            JoinConditionAnalysis.forExpression("x == \"__j.country\"", "__j.", ExprMacroTable.nil())
+        ) // this joinable will not be converted to filter since its a LEFT join
+    );
+
+    final Pair<List<Filter>, List<JoinableClause>> conversion = JoinableFactoryWrapper.convertJoinsToFilters(
+        clauses,
+        ImmutableSet.of("x"),
+        Integer.MAX_VALUE
+    );
+
+    Assert.assertEquals(
+        Pair.of(
+            ImmutableList.of(
+                new InDimFilter("x", TEST_LOOKUP_KEYS),
+                new InDimFilter(
+                    "x",
+                    INDEXED_TABLE_DS.getRowsAsList().stream().map(row -> row[0].toString()).collect(Collectors.toSet())
+                )
+            ),
+            ImmutableList.of(clauses.get(1), clauses.get(2))
+        ),
+        conversion
+    );
+  }
+
   @Test
   public void test_convertJoinsToFilters_dontConvertTooManyValues()
   {