Reduce array copying for non-equi join condition evaluation for LookupJoinOperator

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseJoinOperator.java

@@ -18,7 +18,6 @@
  */
 package org.apache.pinot.query.runtime.operator;
 
-import java.util.AbstractList;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -35,6 +34,7 @@
 import org.apache.pinot.query.runtime.blocks.MseBlock;
 import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock;
 import org.apache.pinot.query.runtime.blocks.SuccessMseBlock;
+import org.apache.pinot.query.runtime.operator.join.JoinedRowView;
 import org.apache.pinot.query.runtime.operator.operands.TransformOperand;
 import org.apache.pinot.query.runtime.operator.operands.TransformOperandFactory;
 import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
@@ -410,104 +410,4 @@ public StatMap.Type getType() {
       return _type;
     }
   }
-
-  /**
-   * This util class is a view over the left and right row joined together.
-   * Currently, this is used for filtering and input of projection. So if the joined
-   * tuple doesn't pass the predicate, the join result is not materialized into {@code Object[]}.
-   */
-  protected abstract static class JoinedRowView extends AbstractList<Object> implements List<Object> {
-    protected final int _leftSize;
-    protected final int _size;
-
-    protected JoinedRowView(int resultColumnSize, int leftSize) {
-      _leftSize = leftSize;
-      _size = resultColumnSize;
-    }
-
-    private static final class BothNotNullView extends JoinedRowView {
-      private final Object[] _leftRow;
-      private final Object[] _rightRow;
-
-      private BothNotNullView(Object[] leftRow, Object[] rightRow, int resultColumnSize, int leftSize) {
-        super(resultColumnSize, leftSize);
-        _leftRow = leftRow;
-        _rightRow = rightRow;
-      }
-
-      @Override
-      public Object get(int i) {
-        return i < _leftSize ? _leftRow[i] : _rightRow[i - _leftSize];
-      }
-
-      @Override
-      public Object[] toArray() {
-        Object[] resultRow = new Object[_size];
-        System.arraycopy(_leftRow, 0, resultRow, 0, _leftSize);
-        System.arraycopy(_rightRow, 0, resultRow, _leftSize, _rightRow.length);
-        return resultRow;
-      }
-    }
-
-    private static final class RightNotNullView extends JoinedRowView {
-      private final Object[] _rightRow;
-
-      public RightNotNullView(Object[] rightRow, int resultColumnSize, int leftSize) {
-        super(resultColumnSize, leftSize);
-        _rightRow = rightRow;
-      }
-
-      @Override
-      public Object get(int i) {
-        return i < _leftSize ? null : _rightRow[i - _leftSize];
-      }
-
-      @Override
-      public Object[] toArray() {
-        Object[] resultRow = new Object[_size];
-        System.arraycopy(_rightRow, 0, resultRow, _leftSize, _rightRow.length);
-        return resultRow;
-      }
-    }
-
-    private static final class LeftNotNullView extends JoinedRowView {
-      private final Object[] _leftRow;
-
-      public LeftNotNullView(Object[] leftRow, int resultColumnSize, int leftSize) {
-        super(resultColumnSize, leftSize);
-        _leftRow = leftRow;
-      }
-
-      @Override
-      public Object get(int i) {
-        return i < _leftSize ? _leftRow[i] : null;
-      }
-
-      @Override
-      public Object[] toArray() {
-        Object[] resultRow = new Object[_size];
-        System.arraycopy(_leftRow, 0, resultRow, 0, _leftSize);
-        return resultRow;
-      }
-    }
-
-    public static JoinedRowView of(@Nullable Object[] leftRow, @Nullable Object[] rightRow, int resultColumnSize,
-        int leftSize) {
-      if (leftRow == null && rightRow == null) {
-        throw new IllegalStateException("both left and right side of join are null");
-      }
-      if (leftRow == null) {
-        return new RightNotNullView(rightRow, resultColumnSize, leftSize);
-      }
-      if (rightRow == null) {
-        return new LeftNotNullView(leftRow, resultColumnSize, leftSize);
-      }
-      return new BothNotNullView(leftRow, rightRow, resultColumnSize, leftSize);
-    }
-
-    @Override
-    public int size() {
-      return _size;
-    }
-  }
 }

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/EnrichedHashJoinOperator.java

@@ -27,6 +27,7 @@
 import org.apache.pinot.common.utils.DataSchema;
 import org.apache.pinot.query.planner.plannode.EnrichedJoinNode;
 import org.apache.pinot.query.runtime.blocks.MseBlock;
+import org.apache.pinot.query.runtime.operator.join.JoinedRowView;
 import org.apache.pinot.query.runtime.operator.operands.TransformOperand;
 import org.apache.pinot.query.runtime.operator.operands.TransformOperandFactory;
 import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/LookupJoinOperator.java

@@ -33,6 +33,7 @@
 import org.apache.pinot.query.planner.plannode.JoinNode;
 import org.apache.pinot.query.runtime.blocks.MseBlock;
 import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock;
+import org.apache.pinot.query.runtime.operator.join.JoinedRowView;
 import org.apache.pinot.query.runtime.operator.operands.TransformOperand;
 import org.apache.pinot.query.runtime.operator.operands.TransformOperandFactory;
 import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
@@ -58,6 +59,7 @@ public class LookupJoinOperator extends MultiStageOperator {
       Set.of(JoinRelType.INNER, JoinRelType.LEFT, JoinRelType.SEMI, JoinRelType.ANTI);
 
   private final MultiStageOperator _leftInput;
+  private final int _leftColumnSize;
   private final LeafOperator _rightInput;
   private final JoinRelType _joinType;
   private final int[] _leftKeyIds;
@@ -68,10 +70,11 @@ public class LookupJoinOperator extends MultiStageOperator {
   private final List<TransformOperand> _nonEquiEvaluators;
   private final StatMap<StatKey> _statMap = new StatMap<>(StatKey.class);
 
-  public LookupJoinOperator(OpChainExecutionContext context, MultiStageOperator leftInput,
+  public LookupJoinOperator(OpChainExecutionContext context, MultiStageOperator leftInput, DataSchema leftSchema,
       MultiStageOperator rightInput, JoinNode node) {
     super(context);
     _leftInput = leftInput;
+    _leftColumnSize = leftSchema.size();
     Preconditions.checkState(rightInput instanceof LeafOperator, "Right input must be leaf operator");
     _rightInput = (LeafOperator) rightInput;
     _joinType = node.getJoinType();
@@ -169,11 +172,11 @@ private List<Object[]> buildJoinedDataBlockDefault(MseBlock.Data leftBlock) {
       PrimaryKey key = getKey(leftRow);
       Object[] rightRow = _rightTable.lookupValues(key, _rightColumns);
       if (rightRow != null) {
-        // TODO: Optimize this to avoid unnecessary object copy.
-        Object[] resultRow = joinRow(leftRow, rightRow);
+        List<Object> resultRow = JoinedRowView.of(leftRow, rightRow, _resultColumnSize, _leftColumnSize);
         if (_nonEquiEvaluators.isEmpty() || _nonEquiEvaluators.stream()
             .allMatch(evaluator -> BooleanUtils.isTrueInternalValue(evaluator.apply(resultRow)))) {
-          rows.add(resultRow);
+          // defer copying of the content until row matches
+          rows.add(resultRow.toArray());
           continue;
         }
       }

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/factory/DefaultJoinOperatorFactory.java

@@ -49,7 +49,7 @@ public MultiStageOperator createJoinOperator(OpChainExecutionContext context, Mu
           return new HashJoinOperator(context, leftOperator, leftSchema, rightOperator, joinNode);
         }
       case LOOKUP:
-        return new LookupJoinOperator(context, leftOperator, rightOperator, joinNode);
+        return new LookupJoinOperator(context, leftOperator, leftSchema, rightOperator, joinNode);
       case ASOF:
         return new AsofJoinOperator(context, leftOperator, leftSchema, rightOperator, joinNode);
       default:

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/JoinedRowView.java

@@ -0,0 +1,124 @@
+/**
+ * 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.join;
+
+import java.util.AbstractList;
+import java.util.List;
+import javax.annotation.Nullable;
+
+
+/**
+ * This util class is a view over the left and right row joined together.
+ * Currently, this is used for filtering and input of projection. So if the joined
+ * tuple doesn't pass the predicate, the join result is not materialized into {@code Object[]}.
+ */
+public abstract class JoinedRowView extends AbstractList<Object> implements List<Object> {
+  protected final int _leftSize;
+  protected final int _size;
+
+  protected JoinedRowView(int resultColumnSize, int leftSize) {
+    _leftSize = leftSize;
+    _size = resultColumnSize;
+  }
+
+  private static final class BothNotNullView extends JoinedRowView {
+    private final Object[] _leftRow;
+    private final Object[] _rightRow;
+
+    private BothNotNullView(Object[] leftRow, Object[] rightRow, int resultColumnSize, int leftSize) {
+      super(resultColumnSize, leftSize);
+      _leftRow = leftRow;
+      _rightRow = rightRow;
+    }
+
+    @Override
+    public Object get(int i) {
+      return i < _leftSize ? _leftRow[i] : _rightRow[i - _leftSize];
+    }
+
+    @Override
+    public Object[] toArray() {
+      Object[] resultRow = new Object[_size];
+      System.arraycopy(_leftRow, 0, resultRow, 0, _leftSize);
+      System.arraycopy(_rightRow, 0, resultRow, _leftSize, _rightRow.length);
+      return resultRow;
+    }
+  }
+
+  private static final class RightNotNullView extends JoinedRowView {
+    private final Object[] _rightRow;
+
+    public RightNotNullView(Object[] rightRow, int resultColumnSize, int leftSize) {
+      super(resultColumnSize, leftSize);
+      _rightRow = rightRow;
+    }
+
+    @Override
+    public Object get(int i) {
+      return i < _leftSize ? null : _rightRow[i - _leftSize];
+    }
+
+    @Override
+    public Object[] toArray() {
+      Object[] resultRow = new Object[_size];
+      System.arraycopy(_rightRow, 0, resultRow, _leftSize, _rightRow.length);
+      return resultRow;
+    }
+  }
+
+  private static final class LeftNotNullView extends JoinedRowView {
+    private final Object[] _leftRow;
+
+    public LeftNotNullView(Object[] leftRow, int resultColumnSize, int leftSize) {
+      super(resultColumnSize, leftSize);
+      _leftRow = leftRow;
+    }
+
+    @Override
+    public Object get(int i) {
+      return i < _leftSize ? _leftRow[i] : null;
+    }
+
+    @Override
+    public Object[] toArray() {
+      Object[] resultRow = new Object[_size];
+      System.arraycopy(_leftRow, 0, resultRow, 0, _leftSize);
+      return resultRow;
+    }
+  }
+
+  public static JoinedRowView of(@Nullable Object[] leftRow, @Nullable Object[] rightRow, int resultColumnSize,
+      int leftSize) {
+    if (leftRow == null && rightRow == null) {
+      throw new IllegalStateException("both left and right side of join are null");
+    }
+    if (leftRow == null) {
+      return new RightNotNullView(rightRow, resultColumnSize, leftSize);
+    }
+    if (rightRow == null) {
+      return new LeftNotNullView(leftRow, resultColumnSize, leftSize);
+    }
+    return new BothNotNullView(leftRow, rightRow, resultColumnSize, leftSize);
+  }
+
+  @Override
+  public int size() {
+    return _size;
+  }
+}