timothy-e commented on code in PR #19210:
URL: https://github.com/apache/pinot/pull/19210#discussion_r3758955388


##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LookupJoinOperator.java:
##########
@@ -97,10 +122,194 @@ public LookupJoinOperator(OpChainExecutionContext 
context, MultiStageOperator le
     _resultSchema = node.getDataSchema();
     _resultColumnSize = _resultSchema.size();
     List<RexExpression> nonEquiConditions = node.getNonEquiConditions();
+    // SEMI and ANTI joins project the left columns only, so an evaluator 
built over the join result schema cannot
+    // reference a dimension table column. Reject the combination here, 
otherwise the loop below fails with an index
+    // error that says nothing about the cause.
+    Preconditions.checkState(nonEquiConditions.isEmpty() || 
_joinType.projectsRight(),
+        "Lookup join type: %s does not support non-equi join conditions, got: 
%s", _joinType, nonEquiConditions);
     _nonEquiEvaluators = new ArrayList<>(nonEquiConditions.size());
     for (RexExpression nonEquiCondition : nonEquiConditions) {
       
_nonEquiEvaluators.add(TransformOperandFactory.getTransformOperand(nonEquiCondition,
 _resultSchema));
     }
+
+    KeyPlan keyPlan =
+        compileKeyPlan(node, rightTableName, 
_rightTable.getPrimaryKeyColumns(), _rightInput.getDataSchema(),
+            _leftColumnSize);
+    _keySize = keyPlan._sources.length;
+    _keySources = keyPlan._sources;
+    _keyConstants = keyPlan._constants;
+    _neverMatches = keyPlan._neverMatches;
+  }
+
+  /// Works out where each value of the lookup key comes from.
+  ///
+  /// The key has one position per dimension table primary key column, in the 
order the dimension table schema declares
+  /// them. Each position is bound in two passes:
+  ///
+  /// 1. Equi-join keys. `rightKeys[i]` names a dimension column, and that 
column's position in the primary key decides
+  ///    where `leftKeys[i]` lands. This is what makes the key independent of 
the order of the join condition.
+  /// 2. Constants. A non-equi condition of the form `dim_column = literal` 
binds a position that pass 1 left open.
+  ///    A constant never replaces an equi-join key, because the equi-join key 
is not kept anywhere else and dropping it
+  ///    would silently widen the join. A constant that pass 1 already bound 
stays in [#_nonEquiEvaluators] and runs as
+  ///    a filter after the lookup, which is what the SQL semantics require.
+  ///
+  /// The method rejects a join condition that cannot produce exactly one 
value per primary key column. Every rejected
+  /// case returned no rows or wrong rows before this validation existed, so 
an error is the better outcome. This is
+  /// also the contract that the single-stage `lookup` transform function 
enforces.
+  @VisibleForTesting
+  static KeyPlan compileKeyPlan(JoinNode node, String tableName, @Nullable 
List<String> primaryKeyColumns,
+      DataSchema rightSchema, int leftColumnSize) {
+    Preconditions.checkState(CollectionUtils.isNotEmpty(primaryKeyColumns),
+        "Failed to find primary key columns for dimension table: %s", 
tableName);
+    String[] rightColumns = rightSchema.getColumnNames();
+    int keySize = primaryKeyColumns.size();
+    int[] sources = new int[keySize];
+    Arrays.fill(sources, KEY_SOURCE_UNBOUND);
+    Object[] constants = new Object[keySize];
+
+    // Pass 1: bind key positions from the equi-join keys.
+    List<Integer> leftKeys = node.getLeftKeys();
+    List<Integer> rightKeys = node.getRightKeys();
+    int numEquiKeys = leftKeys.size();
+    for (int i = 0; i < numEquiKeys; i++) {
+      String rightColumn = rightColumns[rightKeys.get(i)];
+      int keyPosition = primaryKeyColumns.indexOf(rightColumn);
+      Preconditions.checkState(keyPosition >= 0,
+          "Lookup join on dimension table: %s has a join key on column: %s, 
which is not a primary key column. "
+              + "Primary key columns: %s", tableName, rightColumn, 
primaryKeyColumns);
+      Preconditions.checkState(sources[keyPosition] == KEY_SOURCE_UNBOUND,
+          "Lookup join on dimension table: %s has multiple join keys on 
primary key column: %s", tableName,
+          rightColumn);
+      sources[keyPosition] = leftKeys.get(i);
+    }
+
+    // Pass 2: bind the remaining key positions from constant equality 
conditions.
+    boolean neverMatches = false;
+    for (RexExpression nonEquiCondition : node.getNonEquiConditions()) {
+      int rightColumnId = getConstantEqualityColumnId(nonEquiCondition, 
leftColumnSize, rightColumns.length);
+      if (rightColumnId < 0) {
+        continue;
+      }
+      int keyPosition = primaryKeyColumns.indexOf(rightColumns[rightColumnId]);
+      if (keyPosition < 0 || sources[keyPosition] != KEY_SOURCE_UNBOUND) {
+        continue;
+      }
+      sources[keyPosition] = KEY_SOURCE_CONSTANT;
+      Object value = getConstantValue(nonEquiCondition);
+      constants[keyPosition] =
+          toStoredValue(value, rightSchema.getColumnDataType(rightColumnId), 
tableName, rightColumns[rightColumnId]);
+      neverMatches |= constants[keyPosition] == null;
+    }
+
+    List<String> unboundColumns = new ArrayList<>();
+    for (int i = 0; i < keySize; i++) {
+      if (sources[i] == KEY_SOURCE_UNBOUND) {
+        unboundColumns.add(primaryKeyColumns.get(i));
+      }
+    }
+    Preconditions.checkState(unboundColumns.isEmpty(),
+        "Lookup join on dimension table: %s cannot determine primary key 
columns: %s from the join condition. "
+            + "A lookup join reads the dimension table by primary key, so the 
join condition must have an equality on "
+            + "every primary key column: %s. Add the missing conditions, or 
remove the lookup join hint to use a hash "
+            + "join instead.", tableName, unboundColumns, primaryKeyColumns);
+    return new KeyPlan(sources, constants, neverMatches);
+  }
+
+  /// Returns the dimension table column id of a `dim_column = literal` 
condition, or -1 when the condition does not
+  /// have that shape. Only an equality against a single literal can serve as 
a key value. A condition such as
+  /// `dim_column IN ('a', 'b')` reaches this method as a disjunction and 
returns -1, because a hash lookup cannot read
+  /// a set of keys.
+  ///
+  /// Non-equi conditions index the joined row, so the dimension table columns 
start at `leftColumnSize`.
+  private static int getConstantEqualityColumnId(RexExpression condition, int 
leftColumnSize, int numRightColumns) {
+    if (!(condition instanceof RexExpression.FunctionCall)) {
+      return -1;
+    }
+    RexExpression.FunctionCall functionCall = (RexExpression.FunctionCall) 
condition;
+    if (!functionCall.getFunctionName().equals(SqlKind.EQUALS.name())) {
+      return -1;
+    }
+    List<RexExpression> operands = functionCall.getFunctionOperands();
+    if (operands.size() != 2) {
+      return -1;
+    }
+    RexExpression inputRef = operands.get(0) instanceof RexExpression.InputRef 
? operands.get(0) : operands.get(1);
+    RexExpression literal = operands.get(0) instanceof RexExpression.InputRef 
? operands.get(1) : operands.get(0);
+    if (!(inputRef instanceof RexExpression.InputRef) || !(literal instanceof 
RexExpression.Literal)) {
+      return -1;
+    }
+    int columnId = ((RexExpression.InputRef) inputRef).getIndex() - 
leftColumnSize;
+    return columnId >= 0 && columnId < numRightColumns ? columnId : -1;
+  }
+
+  /// Returns the literal value of a condition that 
[#getConstantEqualityColumnId] accepted.
+  @Nullable
+  private static Object getConstantValue(RexExpression condition) {
+    List<RexExpression> operands = ((RexExpression.FunctionCall) 
condition).getFunctionOperands();
+    RexExpression literal = operands.get(0) instanceof RexExpression.Literal ? 
operands.get(0) : operands.get(1);
+    return ((RexExpression.Literal) literal).getValue();
+  }
+
+  /// Converts a literal value to the representation that the dimension table 
stores.
+  ///
+  /// A literal already holds Pinot's internal value, but its numeric width 
follows the type that the planner gave the
+  /// literal, which can be narrower or wider than the dimension column. 
[PrimaryKey] compares values with `equals`,
+  /// where an `Integer` never equals a `Long`, so a literal of the wrong 
width silently misses every row.
+  ///
+  /// The switch rejects every type it cannot convert, rather than passing the 
value through. A value that does not
+  /// match the stored representation misses every row, and this operator 
reports no rows the same way whether the key
+  /// is genuinely absent or malformed. The single-stage `lookup` transform 
function rejects the same way.
+  ///
+  /// BIG_DECIMAL is rejected because `BigDecimal#equals` compares the scale, 
so a literal of `1.5` never matches a
+  /// stored `1.50`. BYTES is rejected because the literal is a 
[org.apache.pinot.spi.utils.ByteArray] while the
+  /// dimension table stores `byte[]`, whose `equals` is identity.

Review Comment:
   the SSE lookup join wraps `byte[]` in a `ByteArray` to compare, can we do 
the same thing here? 
https://github.com/apache/pinot/blob/master/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/LookupTransformFunction.java#L218



##########
pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LookupJoinOperatorTest.java:
##########
@@ -0,0 +1,237 @@
+/**
+ * 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.pinot.query.runtime.operator;
+
+import java.math.BigDecimal;
+import java.util.List;
+import javax.annotation.Nullable;
+import org.apache.calcite.rel.core.JoinRelType;
+import org.apache.calcite.sql.SqlKind;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.query.planner.logical.RexExpression;
+import org.apache.pinot.query.planner.plannode.JoinNode;
+import org.apache.pinot.query.planner.plannode.PlanNode;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.expectThrows;
+
+
+/// Tests [LookupJoinOperator#compileKeyPlan], which decides where each value 
of the dimension table lookup key comes
+/// from.
+///
+/// The dimension table is a hash map keyed by the primary key values, so a 
key is only usable when it holds one value
+/// per primary key column, in the order the dimension table schema declares 
them, and with the stored type of each
+/// column. These tests cover the cases that a query alone cannot reach, such 
as a null literal and a literal of the
+/// wrong numeric width. End-to-end coverage is in `LookupJoin.json`.
+public class LookupJoinOperatorTest {
+  private static final String TABLE_NAME = "dim_tbl_OFFLINE";
+
+  /// Dimension table columns, in the order that the leaf stage reports them. 
The order is not the primary key order,
+  /// which is what makes the key positions worth testing.
+  private static final String[] RIGHT_COLUMNS = {"currency", "rate", 
"rate_start_date"};
+  private static final DataSchema RIGHT_SCHEMA = new DataSchema(RIGHT_COLUMNS,
+      new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.INT, 
ColumnDataType.LONG});
+  private static final List<String> PRIMARY_KEY_COLUMNS = List.of("currency", 
"rate_start_date");
+
+  /// The fact table has two columns, so the dimension table columns start at 
index 2 of the joined row.
+  private static final int LEFT_COLUMN_SIZE = 2;
+
+  @Test
+  public void testKeyPositionsFollowPrimaryKeyOrderNotConditionOrder() {
+    // ON dim.rate_start_date = fact.col1 AND dim.currency = fact.col0
+    // The conditions are in reverse primary key order, so the key values must 
still land in primary key order.
+    LookupJoinOperator.KeyPlan keyPlan =
+        compileKeyPlan(List.of(1, 0), List.of(2, 0), List.of());

Review Comment:
   If we added something like
   ```java
   private static final int FACT_COL0 = 0;
   private static final int FACT_COL1 = 1;
   ...
   private static ... eq(int leftCol, int RightCol)
   ```
   
   then we could make these tests look like 
   ```java
   compileKeyPlan(List.of(eq(FACT_COL1, DIM_RATE_START_DATE), eq(FACT_COL0, 
DIM_CURRENCY)), List.of())
   ```
   instead of 
   ```java
   compileKeyPlan(List.of(1, 0), List.of(2, 0), List.of());
   ```
   
   which makes the intention much easier to read



##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LookupJoinOperator.java:
##########
@@ -97,10 +122,194 @@ public LookupJoinOperator(OpChainExecutionContext 
context, MultiStageOperator le
     _resultSchema = node.getDataSchema();
     _resultColumnSize = _resultSchema.size();
     List<RexExpression> nonEquiConditions = node.getNonEquiConditions();
+    // SEMI and ANTI joins project the left columns only, so an evaluator 
built over the join result schema cannot
+    // reference a dimension table column. Reject the combination here, 
otherwise the loop below fails with an index
+    // error that says nothing about the cause.
+    Preconditions.checkState(nonEquiConditions.isEmpty() || 
_joinType.projectsRight(),
+        "Lookup join type: %s does not support non-equi join conditions, got: 
%s", _joinType, nonEquiConditions);
     _nonEquiEvaluators = new ArrayList<>(nonEquiConditions.size());
     for (RexExpression nonEquiCondition : nonEquiConditions) {
       
_nonEquiEvaluators.add(TransformOperandFactory.getTransformOperand(nonEquiCondition,
 _resultSchema));
     }
+
+    KeyPlan keyPlan =
+        compileKeyPlan(node, rightTableName, 
_rightTable.getPrimaryKeyColumns(), _rightInput.getDataSchema(),
+            _leftColumnSize);
+    _keySize = keyPlan._sources.length;
+    _keySources = keyPlan._sources;
+    _keyConstants = keyPlan._constants;
+    _neverMatches = keyPlan._neverMatches;
+  }
+
+  /// Works out where each value of the lookup key comes from.
+  ///
+  /// The key has one position per dimension table primary key column, in the 
order the dimension table schema declares
+  /// them. Each position is bound in two passes:
+  ///
+  /// 1. Equi-join keys. `rightKeys[i]` names a dimension column, and that 
column's position in the primary key decides
+  ///    where `leftKeys[i]` lands. This is what makes the key independent of 
the order of the join condition.
+  /// 2. Constants. A non-equi condition of the form `dim_column = literal` 
binds a position that pass 1 left open.
+  ///    A constant never replaces an equi-join key, because the equi-join key 
is not kept anywhere else and dropping it
+  ///    would silently widen the join. A constant that pass 1 already bound 
stays in [#_nonEquiEvaluators] and runs as
+  ///    a filter after the lookup, which is what the SQL semantics require.
+  ///
+  /// The method rejects a join condition that cannot produce exactly one 
value per primary key column. Every rejected
+  /// case returned no rows or wrong rows before this validation existed, so 
an error is the better outcome. This is
+  /// also the contract that the single-stage `lookup` transform function 
enforces.
+  @VisibleForTesting
+  static KeyPlan compileKeyPlan(JoinNode node, String tableName, @Nullable 
List<String> primaryKeyColumns,
+      DataSchema rightSchema, int leftColumnSize) {
+    Preconditions.checkState(CollectionUtils.isNotEmpty(primaryKeyColumns),
+        "Failed to find primary key columns for dimension table: %s", 
tableName);
+    String[] rightColumns = rightSchema.getColumnNames();
+    int keySize = primaryKeyColumns.size();
+    int[] sources = new int[keySize];
+    Arrays.fill(sources, KEY_SOURCE_UNBOUND);
+    Object[] constants = new Object[keySize];
+
+    // Pass 1: bind key positions from the equi-join keys.
+    List<Integer> leftKeys = node.getLeftKeys();
+    List<Integer> rightKeys = node.getRightKeys();
+    int numEquiKeys = leftKeys.size();
+    for (int i = 0; i < numEquiKeys; i++) {
+      String rightColumn = rightColumns[rightKeys.get(i)];
+      int keyPosition = primaryKeyColumns.indexOf(rightColumn);
+      Preconditions.checkState(keyPosition >= 0,
+          "Lookup join on dimension table: %s has a join key on column: %s, 
which is not a primary key column. "
+              + "Primary key columns: %s", tableName, rightColumn, 
primaryKeyColumns);
+      Preconditions.checkState(sources[keyPosition] == KEY_SOURCE_UNBOUND,
+          "Lookup join on dimension table: %s has multiple join keys on 
primary key column: %s", tableName,
+          rightColumn);
+      sources[keyPosition] = leftKeys.get(i);
+    }
+
+    // Pass 2: bind the remaining key positions from constant equality 
conditions.
+    boolean neverMatches = false;
+    for (RexExpression nonEquiCondition : node.getNonEquiConditions()) {
+      int rightColumnId = getConstantEqualityColumnId(nonEquiCondition, 
leftColumnSize, rightColumns.length);
+      if (rightColumnId < 0) {
+        continue;
+      }
+      int keyPosition = primaryKeyColumns.indexOf(rightColumns[rightColumnId]);
+      if (keyPosition < 0 || sources[keyPosition] != KEY_SOURCE_UNBOUND) {
+        continue;
+      }
+      sources[keyPosition] = KEY_SOURCE_CONSTANT;
+      Object value = getConstantValue(nonEquiCondition);
+      constants[keyPosition] =
+          toStoredValue(value, rightSchema.getColumnDataType(rightColumnId), 
tableName, rightColumns[rightColumnId]);
+      neverMatches |= constants[keyPosition] == null;
+    }
+
+    List<String> unboundColumns = new ArrayList<>();
+    for (int i = 0; i < keySize; i++) {
+      if (sources[i] == KEY_SOURCE_UNBOUND) {
+        unboundColumns.add(primaryKeyColumns.get(i));
+      }
+    }
+    Preconditions.checkState(unboundColumns.isEmpty(),
+        "Lookup join on dimension table: %s cannot determine primary key 
columns: %s from the join condition. "
+            + "A lookup join reads the dimension table by primary key, so the 
join condition must have an equality on "
+            + "every primary key column: %s. Add the missing conditions, or 
remove the lookup join hint to use a hash "
+            + "join instead.", tableName, unboundColumns, primaryKeyColumns);
+    return new KeyPlan(sources, constants, neverMatches);
+  }
+
+  /// Returns the dimension table column id of a `dim_column = literal` 
condition, or -1 when the condition does not
+  /// have that shape. Only an equality against a single literal can serve as 
a key value. A condition such as
+  /// `dim_column IN ('a', 'b')` reaches this method as a disjunction and 
returns -1, because a hash lookup cannot read
+  /// a set of keys.
+  ///
+  /// Non-equi conditions index the joined row, so the dimension table columns 
start at `leftColumnSize`.
+  private static int getConstantEqualityColumnId(RexExpression condition, int 
leftColumnSize, int numRightColumns) {
+    if (!(condition instanceof RexExpression.FunctionCall)) {
+      return -1;
+    }
+    RexExpression.FunctionCall functionCall = (RexExpression.FunctionCall) 
condition;
+    if (!functionCall.getFunctionName().equals(SqlKind.EQUALS.name())) {
+      return -1;
+    }
+    List<RexExpression> operands = functionCall.getFunctionOperands();
+    if (operands.size() != 2) {
+      return -1;
+    }
+    RexExpression inputRef = operands.get(0) instanceof RexExpression.InputRef 
? operands.get(0) : operands.get(1);
+    RexExpression literal = operands.get(0) instanceof RexExpression.InputRef 
? operands.get(1) : operands.get(0);
+    if (!(inputRef instanceof RexExpression.InputRef) || !(literal instanceof 
RexExpression.Literal)) {
+      return -1;
+    }
+    int columnId = ((RexExpression.InputRef) inputRef).getIndex() - 
leftColumnSize;
+    return columnId >= 0 && columnId < numRightColumns ? columnId : -1;
+  }
+
+  /// Returns the literal value of a condition that 
[#getConstantEqualityColumnId] accepted.
+  @Nullable
+  private static Object getConstantValue(RexExpression condition) {
+    List<RexExpression> operands = ((RexExpression.FunctionCall) 
condition).getFunctionOperands();
+    RexExpression literal = operands.get(0) instanceof RexExpression.Literal ? 
operands.get(0) : operands.get(1);
+    return ((RexExpression.Literal) literal).getValue();
+  }
+
+  /// Converts a literal value to the representation that the dimension table 
stores.
+  ///
+  /// A literal already holds Pinot's internal value, but its numeric width 
follows the type that the planner gave the
+  /// literal, which can be narrower or wider than the dimension column. 
[PrimaryKey] compares values with `equals`,
+  /// where an `Integer` never equals a `Long`, so a literal of the wrong 
width silently misses every row.
+  ///
+  /// The switch rejects every type it cannot convert, rather than passing the 
value through. A value that does not
+  /// match the stored representation misses every row, and this operator 
reports no rows the same way whether the key
+  /// is genuinely absent or malformed. The single-stage `lookup` transform 
function rejects the same way.
+  ///
+  /// BIG_DECIMAL is rejected because `BigDecimal#equals` compares the scale, 
so a literal of `1.5` never matches a

Review Comment:
   How does BIG_DECIMAL work on other join types?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to