Fix potential NPEs in joins

processing/pom.xml

@@ -193,6 +193,12 @@
             <artifactId>caliper</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-core</artifactId>
+            <version>3.2.4</version>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>nl.jqno.equalsverifier</groupId>
             <artifactId>equalsverifier</artifactId>

processing/src/main/java/org/apache/druid/segment/filter/cnf/CalciteCnfHelper.java

@@ -59,9 +59,7 @@ public static Filter pull(Filter rex)
       final List<Filter> list = new ArrayList<>();
       for (Filter operand : operands) {
         Filter removed = removeFactor(factors, operand);
-        if (removed != null) {
-          list.add(removed);
-        }
+        list.add(removed);
       }
       if (list.isEmpty()) {
         return and(factors.values());

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

@@ -236,6 +236,7 @@ public Sequence<Cursor> makeCursors(
     return Sequences.map(
         baseCursorSequence,
         cursor -> {
+          assert cursor != null;
           Cursor retVal = cursor;
 
           for (JoinableClause clause : clauses) {

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

@@ -79,7 +79,7 @@ private 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());
+    rightKeyColumns = getEquiConditions().stream().map(Equality::getRightColumn).collect(Collectors.toSet());
   }
 
   /**
@@ -108,14 +108,18 @@ public static JoinConditionAnalysis forExpression(
         nonEquiConditions.add(childExpr);
       } else {
         final Pair<Expr, Expr> decomposed = maybeDecomposed.get();
-        final Expr lhs = decomposed.lhs;
-        final Expr rhs = decomposed.rhs;
+        final Expr lhs = Objects.requireNonNull(decomposed.lhs);
+        final Expr rhs = Objects.requireNonNull(decomposed.rhs);
 
         if (isLeftExprAndRightColumn(lhs, rhs, rightPrefix)) {
           // rhs is a right-hand column; lhs is an expression solely of the left-hand side.
-          equiConditions.add(new Equality(lhs, rhs.getBindingIfIdentifier().substring(rightPrefix.length())));
+          equiConditions.add(
+              new Equality(lhs, Objects.requireNonNull(rhs.getBindingIfIdentifier()).substring(rightPrefix.length()))
+          );
         } else if (isLeftExprAndRightColumn(rhs, lhs, rightPrefix)) {
-          equiConditions.add(new Equality(rhs, lhs.getBindingIfIdentifier().substring(rightPrefix.length())));
+          equiConditions.add(
+              new Equality(rhs, Objects.requireNonNull(lhs.getBindingIfIdentifier()).substring(rightPrefix.length()))
+          );
         } else {
           nonEquiConditions.add(childExpr);
         }

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

@@ -81,9 +81,9 @@ JoinMatcher makeJoinMatcher(
    * Searches a column from this Joinable for a particular value, finds rows that match,
    * and returns values of a second column for those rows.
    *
-   * @param searchColumnName Name of the search column
-   * @param searchColumnValue Target value of the search column
-   * @param retrievalColumnName The column to retrieve values from
+   * @param searchColumnName Name of the search column. This is the column that is being used in the filter
+   * @param searchColumnValue Target value of the search column. This is the value that is being filtered on.
+   * @param retrievalColumnName The column to retrieve values from. This is the column that is being joined against.
    * @param maxCorrelationSetSize Maximum number of values to retrieve. If we detect that more values would be
    *                              returned than this limit, return an empty set.
    * @param allowNonKeyColumnSearch If true, allow searchs on non-key columns. If this is false,

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

@@ -138,7 +138,10 @@ public int lookupId(@Nullable String name)
       // id 0 is always null for this selector impl.
       return 0;
     } else {
-      return baseSelector.idLookup().lookupId(name) + nullAdjustment;
+      IdLookup idLookup = baseSelector.idLookup();
+      // idLookup is null here because callers are expected to check this condition before calling lookupId
+      assert idLookup != null;
+      return idLookup.lookupId(name) + nullAdjustment;
     }
   }
 

processing/src/main/java/org/apache/druid/segment/join/filter/JoinFilterAnalyzer.java

@@ -38,6 +38,7 @@
 import org.apache.druid.segment.join.JoinableClause;
 import org.apache.druid.segment.virtual.ExpressionVirtualColumn;
 
+import javax.annotation.Nullable;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -175,9 +176,7 @@ public static JoinFilterPreAnalysis computeJoinFilterPreAnalysis(
       for (Equality equality : clause.getCondition().getEquiConditions()) {
         Set<Expr> exprsForRhs = equiconditions.computeIfAbsent(
             clause.getPrefix() + equality.getRightColumn(),
-            (rhs) -> {
-              return new HashSet<>();
-            }
+            (rhs) -> new HashSet<>()
         );
         exprsForRhs.add(equality.getLeftExpr());
       }
@@ -263,9 +262,7 @@ public static JoinFilterPreAnalysis computeJoinFilterPreAnalysis(
           Optional<List<JoinFilterColumnCorrelationAnalysis>> perColumnCorrelations =
               correlationsByFilteringColumn.computeIfAbsent(
                   rhsRewriteCandidate.getRhsColumn(),
-                  (rhsCol) -> {
-                    return Optional.of(new ArrayList<>());
-                  }
+                  (rhsCol) -> Optional.of(new ArrayList<>())
               );
           perColumnCorrelations.get().add(correlationForPrefix.getValue());
           correlationForPrefix.getValue().getCorrelatedValuesMap().computeIfAbsent(
@@ -350,6 +347,7 @@ public static JoinFilterSplit splitFilter(
           joinFilterPreAnalysis
       );
       if (joinFilterAnalysis.isCanPushDown()) {
+        //noinspection OptionalGetWithoutIsPresent isCanPushDown checks isPresent
         leftFilters.add(joinFilterAnalysis.getPushDownFilter().get());
         if (!joinFilterAnalysis.getPushDownVirtualColumns().isEmpty()) {
           pushDownVirtualColumns.addAll(joinFilterAnalysis.getPushDownVirtualColumns());
@@ -438,6 +436,7 @@ private static JoinFilterAnalysis rewriteOrFilter(
         if (!rewritten.isCanPushDown()) {
           return JoinFilterAnalysis.createNoPushdownFilterAnalysis(orFilter);
         } else {
+          //noinspection OptionalGetWithoutIsPresent isCanPushDown checks isPresent
           newFilters.add(rewritten.getPushDownFilter().get());
         }
       } else {
@@ -762,6 +761,7 @@ private static boolean filterMatchesNull(Filter filter)
     return valueMatcher.matches();
   }
 
+  @Nullable
   private static JoinableClause isColumnFromJoin(
       List<JoinableClause> joinableClauses,
       String column

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

@@ -275,6 +275,7 @@ public void matchRemainder()
     } else if (condition.isAlwaysTrue()) {
       currentIterator = Collections.emptyIterator();
     } else {
+      //noinspection ConstantConditions - entry can not be null because extractor.iterable() prevents this
       currentIterator = Iterators.filter(
           extractor.iterable().iterator(),
           entry -> !matchedKeys.contains(entry.getKey())

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

@@ -31,6 +31,7 @@
 import org.apache.druid.segment.join.Joinable;
 
 import javax.annotation.Nullable;
+import java.util.Collections;
 import java.util.List;
 import java.util.Set;
 
@@ -95,18 +96,23 @@ public Set<String> getCorrelatedColumnValues(
       boolean allowNonKeyColumnSearch
   )
   {
+    if (!ALL_COLUMNS.contains(searchColumnName) || !ALL_COLUMNS.contains(retrievalColumnName)) {
+      return ImmutableSet.of();
+    }
     Set<String> correlatedValues;
     if (LookupColumnSelectorFactory.KEY_COLUMN.equals(searchColumnName)) {
       if (LookupColumnSelectorFactory.KEY_COLUMN.equals(retrievalColumnName)) {
         correlatedValues = ImmutableSet.of(searchColumnValue);
       } else {
-        correlatedValues = ImmutableSet.of(extractor.apply(searchColumnName));
+        // This should not happen in practice because the column to be joined on must be a key.
+        correlatedValues = Collections.singleton(extractor.apply(searchColumnValue));
       }
     } else {
       if (!allowNonKeyColumnSearch) {
         return ImmutableSet.of();
       }
       if (LookupColumnSelectorFactory.VALUE_COLUMN.equals(retrievalColumnName)) {
+        // This should not happen in practice because the column to be joined on must be a key.
         correlatedValues = ImmutableSet.of(searchColumnValue);
       } else {
         // Lookup extractor unapply only provides a list of strings, so we can't respect

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

@@ -52,6 +52,7 @@ public double getDouble()
 
     // Otherwise this shouldn't have been called (due to isNull returning true).
     assert NullHandling.replaceWithDefault();
+    //noinspection ConstantConditions assert statement above guarantees this is non null.
     return NullHandling.defaultDoubleValue();
   }
 
@@ -70,6 +71,7 @@ public float getFloat()
 
     // Otherwise this shouldn't have been called (due to isNull returning true).
     assert NullHandling.replaceWithDefault();
+    //noinspection ConstantConditions assert statement above guarantees this is non null.
     return NullHandling.defaultFloatValue();
   }
 
@@ -88,6 +90,7 @@ public long getLong()
 
     // Otherwise this shouldn't have been called (due to isNull returning true).
     assert NullHandling.replaceWithDefault();
+    //noinspection ConstantConditions assert statement above guarantees this is non null.
     return NullHandling.defaultLongValue();
   }
 

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

@@ -30,6 +30,7 @@
 import javax.annotation.Nullable;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Objects;
 import java.util.Set;
 
 public class IndexedTableJoinable implements Joinable
@@ -103,7 +104,8 @@ public Set<String> getCorrelatedColumnValues(
       IntList rowIndex = index.find(searchColumnValue);
       for (int i = 0; i < rowIndex.size(); i++) {
         int rowNum = rowIndex.getInt(i);
-        correlatedValues.add(reader.read(rowNum).toString());
+        String correlatedDimVal = Objects.toString(reader.read(rowNum), null);
+        correlatedValues.add(correlatedDimVal);
 
         if (correlatedValues.size() > maxCorrelationSetSize) {
           return ImmutableSet.of();
@@ -118,11 +120,13 @@ public Set<String> getCorrelatedColumnValues(
       IndexedTable.Reader dimNameReader = table.columnReader(filterColumnPosition);
       IndexedTable.Reader correlatedColumnReader = table.columnReader(correlatedColumnPosition);
       for (int i = 0; i < table.numRows(); i++) {
-        if (searchColumnValue.equals(dimNameReader.read(i).toString())) {
-          correlatedValues.add(correlatedColumnReader.read(i).toString());
-        }
-        if (correlatedValues.size() > maxCorrelationSetSize) {
-          return ImmutableSet.of();
+        String dimVal = Objects.toString(dimNameReader.read(i), null);
+        if (searchColumnValue.equals(dimVal)) {
+          String correlatedDimVal = Objects.toString(correlatedColumnReader.read(i), null);
+          correlatedValues.add(correlatedDimVal);
+          if (correlatedValues.size() > maxCorrelationSetSize) {
+            return ImmutableSet.of();
+          }
         }
       }
 

processing/src/main/java/org/apache/druid/segment/transform/Transformer.java

@@ -156,6 +156,7 @@ public long getTimestampFromEpoch()
     {
       final RowFunction transform = transforms.get(ColumnHolder.TIME_COLUMN_NAME);
       if (transform != null) {
+        //noinspection ConstantConditions time column is never null
         return Rows.objectToNumber(ColumnHolder.TIME_COLUMN_NAME, transform.eval(row), true).longValue();
       } else {
         return row.getTimestampFromEpoch();

processing/src/main/java/org/apache/druid/segment/vector/QueryableIndexVectorColumnSelectorFactory.java

@@ -96,6 +96,8 @@ public MultiValueDimensionVectorSelector makeMultiValueDimensionSelector(final D
           final DictionaryEncodedColumn<String> dictionaryEncodedColumn = (DictionaryEncodedColumn<String>)
               getCachedColumn(spec.getDimension());
 
+          // dictionaryEncodedColumn is not null because of holder null check above
+          assert dictionaryEncodedColumn != null;
           final MultiValueDimensionVectorSelector selector = dictionaryEncodedColumn.makeMultiValueDimensionVectorSelector(
               offset
           );
@@ -132,6 +134,8 @@ public SingleValueDimensionVectorSelector makeSingleValueDimensionSelector(final
           final DictionaryEncodedColumn<String> dictionaryEncodedColumn = (DictionaryEncodedColumn<String>)
               getCachedColumn(spec.getDimension());
 
+          // dictionaryEncodedColumn is not null because of holder null check above
+          assert dictionaryEncodedColumn != null;
           final SingleValueDimensionVectorSelector selector =
               dictionaryEncodedColumn.makeSingleValueDimensionVectorSelector(offset);
 

processing/src/main/java/org/apache/druid/segment/virtual/MultiValueExpressionDimensionSelector.java

@@ -55,6 +55,7 @@ ExprEval getEvaluated()
     return baseSelector.getObject();
   }
 
+  @Nullable
   String getValue(ExprEval evaluated)
   {
     assert !evaluated.isArray();
@@ -64,15 +65,18 @@ String getValue(ExprEval evaluated)
   List<String> getArray(ExprEval evaluated)
   {
     assert evaluated.isArray();
+    //noinspection ConstantConditions
     return Arrays.stream(evaluated.asStringArray())
                  .map(NullHandling::emptyToNullIfNeeded)
                  .collect(Collectors.toList());
   }
 
+  @Nullable
   String getArrayValue(ExprEval evaluated, int i)
   {
     assert evaluated.isArray();
     String[] stringArray = evaluated.asStringArray();
+    //noinspection ConstantConditions because of assert statement above
     assert i < stringArray.length;
     return NullHandling.emptyToNullIfNeeded(stringArray[i]);
   }
@@ -83,7 +87,8 @@ public IndexedInts getRow()
     ExprEval evaluated = getEvaluated();
     if (evaluated.isArray()) {
       RangeIndexedInts ints = new RangeIndexedInts();
-      ints.setSize(evaluated.asArray() != null ? evaluated.asArray().length : 0);
+      Object[] evaluatedArray = evaluated.asArray();
+      ints.setSize(evaluatedArray != null ? evaluatedArray.length : 0);
       return ints;
     }
     return ZeroIndexedInts.instance();

processing/src/main/java/org/apache/druid/segment/virtual/SingleInputBindings.java

@@ -25,6 +25,7 @@
 
 public class SingleInputBindings implements Expr.ObjectBinding
 {
+  @Nullable
   private Object value;
 
   @Override