yashmayya commented on code in PR #19210: URL: https://github.com/apache/pinot/pull/19210#discussion_r3761929623
########## 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: Good suggestion, done. ########## 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: Good point, but I think wrapping won't help and `BYTES` dimension primary keys are probably broken today, everywhere. The dimension table's map key holds a raw `byte[]`: `PinotSegmentRecordReader.getRecordValues` → `PinotSegmentColumnReader.getValue` → either `BytesDictionary.get()` (returns `byte[]`, while the `ByteArray`-returning variant is `getInternal`, which is not called here) or `ForwardIndexReader.getBytes()`. And `DimensionTableDataManager.HASH_STRATEGY` is `Arrays.hashCode`/`Arrays.equals` over `Object[]`, which delegate to each element's own `hashCode`/`equals` - identity for `byte[]`. So nothing can ever probe a `BYTES` primary key: not a `ByteArray` (different class, content-based hash), not even a fresh `byte[]` with the same bytes. That includes `LookupTransformFunction` itself - so I don't think the code you've linked is a working reference. Let's track this separately since it doesn't really relate to the fix here. ########## 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: I think the same hazard exists elsewhere too TBH, so this isn't lookup join specific. I've updated this. -- 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]
