Add getRightEquiConditionKeys to JoinConditionAnalysis (#9287)

* Add getRightColumns to JoinConditionAnalysis

This change other implementations of JoinableFactory to ask the analysis
for the right key columns instead of having to calculate it themselves.

* Address some review comments

* more code review stuff

processing/src/main/java/org/apache/druid/segment/join/JoinConditionAnalysis.java

@@ -32,6 +32,8 @@ import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
 
 /**
  * Represents analysis of a join condition.
@@ -55,6 +57,7 @@ public class JoinConditionAnalysis
   private final boolean isAlwaysFalse;
   private final boolean isAlwaysTrue;
   private final boolean canHashJoin;
+  private final Set<String> rightKeyColumns;
 
   private JoinConditionAnalysis(
       final String originalExpression,
@@ -76,6 +79,7 @@ public class JoinConditionAnalysis
                                                                 .allMatch(expr -> expr.isLiteral() && expr.eval(
                                                                     ExprUtils.nilBindings()).asBoolean());
     canHashJoin = nonEquiConditions.stream().allMatch(Expr::isLiteral);
+    rightKeyColumns = getEquiConditions().stream().map(Equality::getRightColumn).distinct().collect(Collectors.toSet());
   }
 
   /**
@@ -176,6 +180,14 @@ public class JoinConditionAnalysis
     return canHashJoin;
   }
 
+  /**
+   * Returns the distinct column keys from the RHS required to evaluate the equi conditions.
+   */
+  public Set<String> getRightEquiConditionKeys()
+  {
+    return rightKeyColumns;
+  }
+
   @Override
   public boolean equals(Object o)
   {

processing/src/main/java/org/apache/druid/segment/join/lookup/LookupJoinMatcher.java

@@ -191,9 +191,9 @@ public class LookupJoinMatcher implements JoinMatcher
       keyExprs = null;
     } else if (!condition.getNonEquiConditions().isEmpty()) {
       throw new IAE("Cannot join lookup with non-equi condition: %s", condition);
-    } else if (!condition.getEquiConditions()
+    } else if (!condition.getRightEquiConditionKeys()
                          .stream()
-                         .allMatch(eq -> eq.getRightColumn().equals(LookupColumnSelectorFactory.KEY_COLUMN))) {
+                         .allMatch(LookupColumnSelectorFactory.KEY_COLUMN::equals)) {
       throw new IAE("Cannot join lookup with condition referring to non-key column: %s", condition);
     } else {
       keyExprs = condition.getEquiConditions().stream().map(Equality::getLeftExpr).collect(Collectors.toList());

processing/src/main/java/org/apache/druid/segment/join/table/IndexedTable.java

@@ -25,6 +25,7 @@ import org.apache.druid.segment.column.ValueType;
 import javax.annotation.Nullable;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 /**
  * An interface to a table where some columns (the 'key columns') have indexes that enable fast lookups.
@@ -36,7 +37,7 @@ public interface IndexedTable
   /**
    * Returns the columns of this table that have indexes.
    */
-  List<String> keyColumns();
+  Set<String> keyColumns();
 
   /**
    * Returns all columns of this table, including the key and non-key columns.

processing/src/main/java/org/apache/druid/segment/join/table/RowBasedIndexedTable.java

@@ -33,6 +33,7 @@ import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
@@ -48,13 +49,13 @@ public class RowBasedIndexedTable<RowType> implements IndexedTable
   private final List<String> columns;
   private final List<ValueType> columnTypes;
   private final List<Function<RowType, Object>> columnFunctions;
-  private final List<String> keyColumns;
+  private final Set<String> keyColumns;
 
   public RowBasedIndexedTable(
       final List<RowType> table,
       final RowAdapter<RowType> rowAdapter,
       final Map<String, ValueType> rowSignature,
-      final List<String> keyColumns
+      final Set<String> keyColumns
   )
   {
     this.table = table;
@@ -107,7 +108,7 @@ public class RowBasedIndexedTable<RowType> implements IndexedTable
   }
 
   @Override
-  public List<String> keyColumns()
+  public Set<String> keyColumns()
   {
     return keyColumns;
   }

processing/src/test/java/org/apache/druid/segment/join/JoinConditionAnalysisTest.java

@@ -20,6 +20,7 @@
 package org.apache.druid.segment.join;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import nl.jqno.equalsverifier.EqualsVerifier;
 import org.apache.druid.common.config.NullHandling;
 import org.apache.druid.java.util.common.Pair;
@@ -60,6 +61,7 @@ public class JoinConditionAnalysisTest
         ImmutableList.of(),
         exprsToStrings(analysis.getNonEquiConditions())
     );
+    Assert.assertEquals(analysis.getRightEquiConditionKeys(), ImmutableSet.of("y"));
   }
 
   @Test
@@ -80,6 +82,7 @@ public class JoinConditionAnalysisTest
         ImmutableList.of(),
         exprsToStrings(analysis.getNonEquiConditions())
     );
+    Assert.assertEquals(analysis.getRightEquiConditionKeys(), ImmutableSet.of("y"));
   }
 
   @Test
@@ -100,6 +103,7 @@ public class JoinConditionAnalysisTest
         ImmutableList.of(),
         exprsToStrings(analysis.getNonEquiConditions())
     );
+    Assert.assertEquals(analysis.getRightEquiConditionKeys(), ImmutableSet.of("z"));
   }
 
   @Test
@@ -120,6 +124,7 @@ public class JoinConditionAnalysisTest
         ImmutableList.of("(== (+ j.x j.y) z)"),
         exprsToStrings(analysis.getNonEquiConditions())
     );
+    Assert.assertTrue(analysis.getRightEquiConditionKeys().isEmpty());
   }
 
   @Test
@@ -140,6 +145,7 @@ public class JoinConditionAnalysisTest
         ImmutableList.of("(== (+ x j.y) j.z)"),
         exprsToStrings(analysis.getNonEquiConditions())
     );
+    Assert.assertTrue(analysis.getRightEquiConditionKeys().isEmpty());
   }
 
   @Test
@@ -160,6 +166,7 @@ public class JoinConditionAnalysisTest
         ImmutableList.of("2"),
         exprsToStrings(analysis.getNonEquiConditions())
     );
+    Assert.assertTrue(analysis.getRightEquiConditionKeys().isEmpty());
   }
 
   @Test
@@ -180,6 +187,7 @@ public class JoinConditionAnalysisTest
         ImmutableList.of("0"),
         exprsToStrings(analysis.getNonEquiConditions())
     );
+    Assert.assertTrue(analysis.getRightEquiConditionKeys().isEmpty());
   }
 
   @Test
@@ -200,6 +208,7 @@ public class JoinConditionAnalysisTest
         ImmutableList.of("(== x 1)"),
         exprsToStrings(analysis.getNonEquiConditions())
     );
+    Assert.assertTrue(analysis.getRightEquiConditionKeys().isEmpty());
   }
 
   @Test
@@ -220,6 +229,7 @@ public class JoinConditionAnalysisTest
         ImmutableList.of(),
         exprsToStrings(analysis.getNonEquiConditions())
     );
+    Assert.assertEquals(analysis.getRightEquiConditionKeys(), ImmutableSet.of("x"));
   }
 
   @Test
@@ -240,6 +250,7 @@ public class JoinConditionAnalysisTest
         ImmutableList.of(),
         exprsToStrings(analysis.getNonEquiConditions())
     );
+    Assert.assertEquals(analysis.getRightEquiConditionKeys(), ImmutableSet.of("y", "z", "zz"));
   }
 
   @Test
@@ -260,6 +271,7 @@ public class JoinConditionAnalysisTest
         ImmutableList.of("(|| (== (+ x y) j.z) (== z j.zz))"),
         exprsToStrings(analysis.getNonEquiConditions())
     );
+    Assert.assertEquals(analysis.getRightEquiConditionKeys(), ImmutableSet.of("y"));
   }
 
   @Test
@@ -270,8 +282,8 @@ public class JoinConditionAnalysisTest
                   .withIgnoredFields(
                           // These fields are tightly coupled with originalExpression
                           "equiConditions", "nonEquiConditions",
-                          // These fields are calculated from nonEquiConditions
+                          // These fields are calculated from other other fields in the class
-                          "isAlwaysTrue", "isAlwaysFalse", "canHashJoin")
+                          "isAlwaysTrue", "isAlwaysFalse", "canHashJoin", "rightKeyColumns")
                   .verify();
   }
 

processing/src/test/java/org/apache/druid/segment/join/JoinTestHelper.java

@@ -23,6 +23,7 @@ import com.fasterxml.jackson.databind.MappingIterator;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Lists;
 import org.apache.druid.common.config.NullHandling;
 import org.apache.druid.data.input.InputRow;
@@ -252,7 +253,7 @@ public class JoinTestHelper
             rows,
             createMapRowAdapter(COUNTRIES_SIGNATURE),
             COUNTRIES_SIGNATURE,
-            ImmutableList.of("countryNumber", "countryIsoCode")
+            ImmutableSet.of("countryNumber", "countryIsoCode")
         )
     );
   }
@@ -265,7 +266,7 @@ public class JoinTestHelper
             rows,
             createMapRowAdapter(REGIONS_SIGNATURE),
             REGIONS_SIGNATURE,
-            ImmutableList.of("regionIsoCode", "countryIsoCode")
+            ImmutableSet.of("regionIsoCode", "countryIsoCode")
         )
     );
   }

processing/src/test/java/org/apache/druid/segment/join/table/IndexedTableJoinableTest.java

@@ -20,6 +20,7 @@
 package org.apache.druid.segment.join.table;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import org.apache.druid.math.expr.ExprMacroTable;
 import org.apache.druid.query.InlineDataSource;
 import org.apache.druid.query.dimension.DefaultDimensionSpec;
@@ -72,7 +73,7 @@ public class IndexedTableJoinableTest
       inlineDataSource.getRowsAsList(),
       inlineDataSource.rowAdapter(),
       inlineDataSource.getRowSignature(),
-      ImmutableList.of("str")
+      ImmutableSet.of("str")
   );
 
   @Test

processing/src/test/java/org/apache/druid/segment/join/table/RowBasedIndexedTableTest.java

@@ -21,6 +21,7 @@ package org.apache.druid.segment.join.table;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import org.apache.druid.common.config.NullHandling;
 import org.apache.druid.segment.column.ValueType;
 import org.apache.druid.segment.join.JoinTestHelper;
@@ -64,7 +65,7 @@ public class RowBasedIndexedTableTest
   @Test
   public void test_keyColumns_countries()
   {
-    Assert.assertEquals(ImmutableList.of("countryNumber", "countryIsoCode"), countriesTable.keyColumns());
+    Assert.assertEquals(ImmutableSet.of("countryNumber", "countryIsoCode"), countriesTable.keyColumns());
   }
 
   @Test

server/src/main/java/org/apache/druid/segment/join/InlineJoinableFactory.java

@@ -25,9 +25,8 @@ 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 java.util.List;
 import java.util.Optional;
-import java.util.stream.Collectors;
+import java.util.Set;
 
 /**
  * A {@link JoinableFactory} for {@link InlineDataSource}. It works by building an {@link IndexedTable}.
@@ -39,8 +38,7 @@ public class InlineJoinableFactory implements JoinableFactory
   {
     if (condition.canHashJoin() && dataSource instanceof InlineDataSource) {
       final InlineDataSource inlineDataSource = (InlineDataSource) dataSource;
-      final List<String> rightKeyColumns =
+      final Set<String> rightKeyColumns = condition.getRightEquiConditionKeys();
-          condition.getEquiConditions().stream().map(Equality::getRightColumn).distinct().collect(Collectors.toList());
 
       return Optional.of(
           new IndexedTableJoinable(