From ca73bf800ee3bfff0685d7b39b7f3f6fae7d88ca Mon Sep 17 00:00:00 2001 From: Gian Merlino Date: Sun, 29 May 2022 12:50:44 -0700 Subject: [PATCH 1/2] Add RowIdSupplier to ColumnSelectorFactory. This enables virtual columns to cache their outputs in case they are called multiple times on the same underlying row. This is common for numeric selectors, where the common pattern is to call isNull() and then follow with getLong(), getFloat(), or getDouble(). Here, output caching reduces the number of expression evals by half. --- .../druid/segment/ColumnSelectorFactory.java | 11 ++ .../QueryableIndexColumnSelectorFactory.java | 15 ++- .../RowBasedColumnSelectorFactory.java | 28 +++-- .../apache/druid/segment/RowIdSupplier.java | 43 +++++++ ...IncrementalIndexColumnSelectorFactory.java | 16 ++- .../druid/segment/join/HashJoinEngine.java | 31 ++++- .../BaseExpressionColumnValueSelector.java | 115 ++++++++++++++++++ .../ExpressionColumnValueSelector.java | 56 ++------- .../segment/virtual/ExpressionSelectors.java | 13 +- ...RowBasedExpressionColumnValueSelector.java | 24 +++- ...tCachingExpressionColumnValueSelector.java | 44 ++----- ...tCachingExpressionColumnValueSelector.java | 45 ++----- .../VirtualizedColumnSelectorFactory.java | 10 ++ 13 files changed, 311 insertions(+), 140 deletions(-) create mode 100644 processing/src/main/java/org/apache/druid/segment/RowIdSupplier.java create mode 100644 processing/src/main/java/org/apache/druid/segment/virtual/BaseExpressionColumnValueSelector.java diff --git a/processing/src/main/java/org/apache/druid/segment/ColumnSelectorFactory.java b/processing/src/main/java/org/apache/druid/segment/ColumnSelectorFactory.java index f99f0eade621..82d9b0a5149a 100644 --- a/processing/src/main/java/org/apache/druid/segment/ColumnSelectorFactory.java +++ b/processing/src/main/java/org/apache/druid/segment/ColumnSelectorFactory.java @@ -53,4 +53,15 @@ public interface ColumnSelectorFactory extends ColumnInspector @Override @Nullable ColumnCapabilities getColumnCapabilities(String column); + + /** + * Returns a {@link RowIdSupplier} that allows callers to detect whether the selectors returned by this + * factory have moved or not. Useful for selectors that wrap other selectors, such as virtual columns, + * as it allows them to cache their outputs. + */ + @Nullable + default RowIdSupplier getRowIdSupplier() + { + return null; + } } diff --git a/processing/src/main/java/org/apache/druid/segment/QueryableIndexColumnSelectorFactory.java b/processing/src/main/java/org/apache/druid/segment/QueryableIndexColumnSelectorFactory.java index 72859024ec6e..bc146f019a4f 100644 --- a/processing/src/main/java/org/apache/druid/segment/QueryableIndexColumnSelectorFactory.java +++ b/processing/src/main/java/org/apache/druid/segment/QueryableIndexColumnSelectorFactory.java @@ -39,7 +39,7 @@ * It's counterpart for incremental index is {@link * org.apache.druid.segment.incremental.IncrementalIndexColumnSelectorFactory}. */ -public class QueryableIndexColumnSelectorFactory implements ColumnSelectorFactory +public class QueryableIndexColumnSelectorFactory implements ColumnSelectorFactory, RowIdSupplier { private final QueryableIndex index; private final VirtualColumns virtualColumns; @@ -193,6 +193,19 @@ private T getCachedColumn(final String columnName, final } } + @Nullable + @Override + public RowIdSupplier getRowIdSupplier() + { + return this; + } + + @Override + public long getRowId() + { + return offset.getOffset(); + } + @Override @Nullable public ColumnCapabilities getColumnCapabilities(String columnName) diff --git a/processing/src/main/java/org/apache/druid/segment/RowBasedColumnSelectorFactory.java b/processing/src/main/java/org/apache/druid/segment/RowBasedColumnSelectorFactory.java index 4fd9860e15a8..a55f18856c7c 100644 --- a/processing/src/main/java/org/apache/druid/segment/RowBasedColumnSelectorFactory.java +++ b/processing/src/main/java/org/apache/druid/segment/RowBasedColumnSelectorFactory.java @@ -40,7 +40,6 @@ import java.util.List; import java.util.Objects; import java.util.function.Function; -import java.util.function.LongSupplier; import java.util.function.Supplier; import java.util.function.ToLongFunction; @@ -49,12 +48,10 @@ */ public class RowBasedColumnSelectorFactory implements ColumnSelectorFactory { - private static final long NO_ID = -1; - private final Supplier rowSupplier; @Nullable - private final LongSupplier rowIdSupplier; + private final RowIdSupplier rowIdSupplier; private final RowAdapter adapter; private final ColumnInspector columnInspector; private final boolean throwParseExceptions; @@ -65,7 +62,7 @@ public class RowBasedColumnSelectorFactory implements ColumnSelectorFactory */ RowBasedColumnSelectorFactory( final Supplier rowSupplier, - @Nullable final LongSupplier rowIdSupplier, + @Nullable final RowIdSupplier rowIdSupplier, final RowAdapter adapter, final ColumnInspector columnInspector, final boolean throwParseExceptions @@ -167,7 +164,7 @@ private DimensionSelector makeDimensionSelectorUndecorated(DimensionSpec dimensi return new BaseSingleValueDimensionSelector() { - private long currentId = NO_ID; + private long currentId = RowIdSupplier.INIT; private String currentValue; @Override @@ -186,11 +183,11 @@ public void inspectRuntimeShape(RuntimeShapeInspector inspector) private void updateCurrentValue() { - if (rowIdSupplier == null || rowIdSupplier.getAsLong() != currentId) { + if (rowIdSupplier == null || rowIdSupplier.getRowId() != currentId) { currentValue = extractionFn.apply(timestampFunction.applyAsLong(rowSupplier.get())); if (rowIdSupplier != null) { - currentId = rowIdSupplier.getAsLong(); + currentId = rowIdSupplier.getRowId(); } } } @@ -200,7 +197,7 @@ private void updateCurrentValue() return new DimensionSelector() { - private long currentId = NO_ID; + private long currentId = RowIdSupplier.INIT; private List dimensionValues; private final RangeIndexedInts indexedInts = new RangeIndexedInts(); @@ -331,7 +328,7 @@ public void inspectRuntimeShape(RuntimeShapeInspector inspector) private void updateCurrentValues() { - if (rowIdSupplier == null || rowIdSupplier.getAsLong() != currentId) { + if (rowIdSupplier == null || rowIdSupplier.getRowId() != currentId) { try { final Object rawValue = dimFunction.apply(rowSupplier.get()); @@ -377,12 +374,12 @@ private void updateCurrentValues() } } catch (Throwable e) { - currentId = NO_ID; + currentId = RowIdSupplier.INIT; throw e; } if (rowIdSupplier != null) { - currentId = rowIdSupplier.getAsLong(); + currentId = rowIdSupplier.getRowId(); } } } @@ -491,6 +488,13 @@ private Number getCurrentValueAsNumber() } } + @Nullable + @Override + public RowIdSupplier getRowIdSupplier() + { + return rowIdSupplier; + } + @Nullable @Override public ColumnCapabilities getColumnCapabilities(String columnName) diff --git a/processing/src/main/java/org/apache/druid/segment/RowIdSupplier.java b/processing/src/main/java/org/apache/druid/segment/RowIdSupplier.java new file mode 100644 index 000000000000..c2c2c857bd3d --- /dev/null +++ b/processing/src/main/java/org/apache/druid/segment/RowIdSupplier.java @@ -0,0 +1,43 @@ +/* + * 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.druid.segment; + +/** + * Returned by {@link ColumnSelectorFactory#getRowIdSupplier()}. Allows users of {@link ColumnSelectorFactory} + * to cache objects returned by their selectors. + */ +public interface RowIdSupplier +{ + /** + * A number that will never be returned from {@link #getRowId()}. Useful for initialization. + */ + long INIT = -1; + + /** + * Returns a number that uniquely identifies the current position of some underlying cursor. This is useful for + * caching: it is safe to assume nothing has changed in the selector as long as the row ID stays the same. + * + * Row IDs do not need to be contiguous or monotonic. They need not have any meaning. In particular: they may not + * be row *numbers* (row number 0 may have any arbitrary row ID). + * + * Valid row IDs are always nonnegative. + */ + long getRowId(); +} diff --git a/processing/src/main/java/org/apache/druid/segment/incremental/IncrementalIndexColumnSelectorFactory.java b/processing/src/main/java/org/apache/druid/segment/incremental/IncrementalIndexColumnSelectorFactory.java index 7cab1e314edb..92e80f3d7038 100644 --- a/processing/src/main/java/org/apache/druid/segment/incremental/IncrementalIndexColumnSelectorFactory.java +++ b/processing/src/main/java/org/apache/druid/segment/incremental/IncrementalIndexColumnSelectorFactory.java @@ -25,6 +25,7 @@ import org.apache.druid.segment.ColumnValueSelector; import org.apache.druid.segment.DimensionIndexer; import org.apache.druid.segment.DimensionSelector; +import org.apache.druid.segment.RowIdSupplier; import org.apache.druid.segment.SingleScanTimeDimensionSelector; import org.apache.druid.segment.VirtualColumns; import org.apache.druid.segment.column.ColumnCapabilities; @@ -37,7 +38,7 @@ * The basic implementation of {@link ColumnSelectorFactory} over an {@link IncrementalIndex}. It's counterpart for * historical segments is {@link org.apache.druid.segment.QueryableIndexColumnSelectorFactory}. */ -class IncrementalIndexColumnSelectorFactory implements ColumnSelectorFactory +class IncrementalIndexColumnSelectorFactory implements ColumnSelectorFactory, RowIdSupplier { private final IncrementalIndexStorageAdapter adapter; private final IncrementalIndex index; @@ -133,4 +134,17 @@ public ColumnCapabilities getColumnCapabilities(String columnName) // Use adapter.getColumnCapabilities instead of index.getCapabilities (see note in IncrementalIndexStorageAdapater) return adapter.getColumnCapabilities(columnName); } + + @Nullable + @Override + public RowIdSupplier getRowIdSupplier() + { + return this; + } + + @Override + public long getRowId() + { + return rowHolder.get().getRowIndex(); + } } diff --git a/processing/src/main/java/org/apache/druid/segment/join/HashJoinEngine.java b/processing/src/main/java/org/apache/druid/segment/join/HashJoinEngine.java index cdf80dcf6764..e6f67a1683e5 100644 --- a/processing/src/main/java/org/apache/druid/segment/join/HashJoinEngine.java +++ b/processing/src/main/java/org/apache/druid/segment/join/HashJoinEngine.java @@ -26,6 +26,7 @@ import org.apache.druid.segment.ColumnValueSelector; import org.apache.druid.segment.Cursor; import org.apache.druid.segment.DimensionSelector; +import org.apache.druid.segment.RowIdSupplier; import org.apache.druid.segment.column.ColumnCapabilities; import org.joda.time.DateTime; @@ -69,8 +70,10 @@ public static Cursor makeJoinCursor( closer ); - class JoinColumnSelectorFactory implements ColumnSelectorFactory + class JoinColumnSelectorFactory implements ColumnSelectorFactory, RowIdSupplier { + private long rowId = 0; + @Override public DimensionSelector makeDimensionSelector(DimensionSpec dimensionSpec) { @@ -116,6 +119,29 @@ public ColumnCapabilities getColumnCapabilities(String column) return leftColumnSelectorFactory.getColumnCapabilities(column); } } + + @Nullable + @Override + public RowIdSupplier getRowIdSupplier() + { + return this; + } + + @Override + public long getRowId() + { + return rowId; + } + + void advanceRowId() + { + rowId++; + } + + void resetRowId() + { + rowId = 0; + } } final JoinColumnSelectorFactory joinColumnSelectorFactory = new JoinColumnSelectorFactory(); @@ -171,6 +197,8 @@ private void matchCurrentPosition() @Override public void advanceUninterruptibly() { + joinColumnSelectorFactory.advanceRowId(); + if (joinMatcher.hasMatch()) { joinMatcher.nextMatch(); @@ -218,6 +246,7 @@ public void reset() { leftCursor.reset(); joinMatcher.reset(); + joinColumnSelectorFactory.resetRowId(); } } diff --git a/processing/src/main/java/org/apache/druid/segment/virtual/BaseExpressionColumnValueSelector.java b/processing/src/main/java/org/apache/druid/segment/virtual/BaseExpressionColumnValueSelector.java new file mode 100644 index 000000000000..868196c600b4 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/segment/virtual/BaseExpressionColumnValueSelector.java @@ -0,0 +1,115 @@ +/* + * 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.druid.segment.virtual; + +import org.apache.druid.math.expr.ExprEval; +import org.apache.druid.query.monomorphicprocessing.RuntimeShapeInspector; +import org.apache.druid.segment.ColumnValueSelector; +import org.apache.druid.segment.RowIdSupplier; + +import javax.annotation.Nullable; + +/** + * Base class for many (although not all) {@code ColumnValueSelector}. + */ +public abstract class BaseExpressionColumnValueSelector implements ColumnValueSelector +{ + @Nullable + private final RowIdSupplier rowIdSupplier; + private long currentRowId = RowIdSupplier.INIT; + private ExprEval currentEval; + + protected BaseExpressionColumnValueSelector(@Nullable RowIdSupplier rowIdSupplier) + { + this.rowIdSupplier = rowIdSupplier; + } + + @Override + public double getDouble() + { + // No assert for null handling as ExprEval already has it. + return computeCurrentEval().asDouble(); + } + + @Override + public float getFloat() + { + // No assert for null handling as ExprEval already has it. + return (float) computeCurrentEval().asDouble(); + } + + @Override + public long getLong() + { + // No assert for null handling as ExprEval already has it. + return computeCurrentEval().asLong(); + } + + @Override + public boolean isNull() + { + // It is possible for an expression to have a non-null String value, but be null when treated as long/float/double. + // Check specifically for numeric nulls, which matches the expected behavior of ColumnValueSelector.isNull. + return computeCurrentEval().isNumericNull(); + } + + @Nullable + @Override + public ExprEval getObject() + { + return computeCurrentEval(); + } + + @Override + public Class classOfObject() + { + return ExprEval.class; + } + + @Override + public void inspectRuntimeShape(final RuntimeShapeInspector inspector) + { + inspector.visit("rowIdSupplier", this); + } + + /** + * Implementations override this. + */ + protected abstract ExprEval eval(); + + /** + * Call {@link #eval()} or use {@code currentEval} as appropriate. + */ + private ExprEval computeCurrentEval() + { + if (rowIdSupplier == null) { + return eval(); + } else { + final long rowId = rowIdSupplier.getRowId(); + + if (currentRowId != rowId) { + currentEval = eval(); + currentRowId = rowId; + } + + return currentEval; + } + } +} diff --git a/processing/src/main/java/org/apache/druid/segment/virtual/ExpressionColumnValueSelector.java b/processing/src/main/java/org/apache/druid/segment/virtual/ExpressionColumnValueSelector.java index c0c3eafa1cac..fad91ded6ac2 100644 --- a/processing/src/main/java/org/apache/druid/segment/virtual/ExpressionColumnValueSelector.java +++ b/processing/src/main/java/org/apache/druid/segment/virtual/ExpressionColumnValueSelector.java @@ -24,54 +24,32 @@ import org.apache.druid.math.expr.ExprEval; import org.apache.druid.query.monomorphicprocessing.RuntimeShapeInspector; import org.apache.druid.segment.ColumnValueSelector; +import org.apache.druid.segment.RowIdSupplier; -import javax.annotation.Nonnull; +import javax.annotation.Nullable; /** * Basic expression {@link ColumnValueSelector}. Evaluates {@link Expr} into {@link ExprEval} against * {@link Expr.ObjectBinding} which are backed by the underlying expression input {@link ColumnValueSelector}s */ -public class ExpressionColumnValueSelector implements ColumnValueSelector +public class ExpressionColumnValueSelector extends BaseExpressionColumnValueSelector { - final Expr.ObjectBinding bindings; - final Expr expression; + private final Expr.ObjectBinding bindings; + private final Expr expression; - public ExpressionColumnValueSelector(Expr expression, Expr.ObjectBinding bindings) + public ExpressionColumnValueSelector( + Expr expression, + Expr.ObjectBinding bindings, + @Nullable RowIdSupplier rowIdSupplier + ) { + super(rowIdSupplier); this.bindings = Preconditions.checkNotNull(bindings, "bindings"); this.expression = Preconditions.checkNotNull(expression, "expression"); } @Override - public double getDouble() - { - // No Assert for null handling as ExprEval already have it. - return getObject().asDouble(); - } - - @Override - public float getFloat() - { - // No Assert for null handling as ExprEval already have it. - return (float) getObject().asDouble(); - } - - @Override - public long getLong() - { - // No Assert for null handling as ExprEval already have it. - return getObject().asLong(); - } - - @Override - public Class classOfObject() - { - return ExprEval.class; - } - - @Nonnull - @Override - public ExprEval getObject() + protected ExprEval eval() { return expression.eval(bindings); } @@ -79,16 +57,8 @@ public ExprEval getObject() @Override public void inspectRuntimeShape(RuntimeShapeInspector inspector) { + super.inspectRuntimeShape(inspector); inspector.visit("expression", expression); inspector.visit("bindings", bindings); } - - @Override - public boolean isNull() - { - // It is possible for an expression to have a non-null String value but it can return null when parsed - // as a primitive long/float/double. - // ExprEval.isNumericNull checks whether the parsed primitive value is null or not. - return getObject().isNumericNull(); - } } diff --git a/processing/src/main/java/org/apache/druid/segment/virtual/ExpressionSelectors.java b/processing/src/main/java/org/apache/druid/segment/virtual/ExpressionSelectors.java index a691f625c02b..76f40db6f9d1 100644 --- a/processing/src/main/java/org/apache/druid/segment/virtual/ExpressionSelectors.java +++ b/processing/src/main/java/org/apache/druid/segment/virtual/ExpressionSelectors.java @@ -40,6 +40,7 @@ import org.apache.druid.segment.ConstantExprEvalSelector; import org.apache.druid.segment.DimensionSelector; import org.apache.druid.segment.NilColumnValueSelector; +import org.apache.druid.segment.RowIdSupplier; import org.apache.druid.segment.column.ColumnCapabilities; import org.apache.druid.segment.column.ColumnHolder; import org.apache.druid.segment.column.ColumnType; @@ -143,6 +144,8 @@ public static ColumnValueSelector makeExprEvalSelector( ExpressionPlan plan ) { + final RowIdSupplier rowIdSupplier = columnSelectorFactory.getRowIdSupplier(); + if (plan.is(ExpressionPlan.Trait.SINGLE_INPUT_SCALAR)) { final String column = plan.getSingleInputName(); final ColumnType inputType = plan.getSingleInputType(); @@ -150,12 +153,14 @@ public static ColumnValueSelector makeExprEvalSelector( return new SingleLongInputCachingExpressionColumnValueSelector( columnSelectorFactory.makeColumnValueSelector(column), plan.getExpression(), - !ColumnHolder.TIME_COLUMN_NAME.equals(column) // __time doesn't need an LRU cache since it is sorted. + !ColumnHolder.TIME_COLUMN_NAME.equals(column), // __time doesn't need an LRU cache since it is sorted. + rowIdSupplier ); } else if (inputType.is(ValueType.STRING)) { return new SingleStringInputCachingExpressionColumnValueSelector( columnSelectorFactory.makeDimensionSelector(new DefaultDimensionSpec(column, column, ColumnType.STRING)), - plan.getExpression() + plan.getExpression(), + rowIdSupplier ); } } @@ -169,11 +174,11 @@ public static ColumnValueSelector makeExprEvalSelector( // if any unknown column input types, fall back to an expression selector that examines input bindings on a // per row basis if (plan.any(ExpressionPlan.Trait.UNKNOWN_INPUTS, ExpressionPlan.Trait.INCOMPLETE_INPUTS)) { - return new RowBasedExpressionColumnValueSelector(plan, bindings); + return new RowBasedExpressionColumnValueSelector(plan, bindings, rowIdSupplier); } // generic expression value selector for fully known input types - return new ExpressionColumnValueSelector(plan.getAppliedExpression(), bindings); + return new ExpressionColumnValueSelector(plan.getAppliedExpression(), bindings, rowIdSupplier); } /** diff --git a/processing/src/main/java/org/apache/druid/segment/virtual/RowBasedExpressionColumnValueSelector.java b/processing/src/main/java/org/apache/druid/segment/virtual/RowBasedExpressionColumnValueSelector.java index 9838103ca013..86e4aedb988e 100644 --- a/processing/src/main/java/org/apache/druid/segment/virtual/RowBasedExpressionColumnValueSelector.java +++ b/processing/src/main/java/org/apache/druid/segment/virtual/RowBasedExpressionColumnValueSelector.java @@ -24,7 +24,10 @@ import org.apache.druid.math.expr.Expr; import org.apache.druid.math.expr.ExprEval; import org.apache.druid.math.expr.Parser; +import org.apache.druid.segment.RowIdSupplier; +import javax.annotation.Nullable; +import java.util.ArrayList; import java.util.HashSet; import java.util.List; import java.util.Set; @@ -37,8 +40,10 @@ * Currently, string dimensions are the only bindings which might appear as a {@link String} or a {@link Object[]}, so * numbers are eliminated from the set of 'unknown' bindings to check as they are encountered. */ -public class RowBasedExpressionColumnValueSelector extends ExpressionColumnValueSelector +public class RowBasedExpressionColumnValueSelector extends BaseExpressionColumnValueSelector { + private final Expr.ObjectBinding bindings; + private final Expr expression; private final List unknownColumns; private final Expr.BindingAnalysis baseBindingAnalysis; private final Set ignoredColumns; @@ -46,10 +51,13 @@ public class RowBasedExpressionColumnValueSelector extends ExpressionColumnValue public RowBasedExpressionColumnValueSelector( ExpressionPlan plan, - Expr.ObjectBinding bindings + Expr.ObjectBinding bindings, + @Nullable RowIdSupplier rowIdSupplier ) { - super(plan.getAppliedExpression(), bindings); + super(rowIdSupplier); + this.bindings = bindings; + this.expression = plan.getAppliedExpression(); this.unknownColumns = plan.getUnknownInputs() .stream() .filter(x -> !plan.getAnalysis().getArrayBindings().contains(x)) @@ -60,10 +68,16 @@ public RowBasedExpressionColumnValueSelector( } @Override - public ExprEval getObject() + protected ExprEval eval() { // check to find any arrays for this row - List arrayBindings = unknownColumns.stream().filter(this::isBindingArray).collect(Collectors.toList()); + List arrayBindings = new ArrayList<>(); + + for (String unknownColumn : unknownColumns) { + if (isBindingArray(unknownColumn)) { + arrayBindings.add(unknownColumn); + } + } // eliminate anything that will never be an array if (ignoredColumns.size() > 0) { diff --git a/processing/src/main/java/org/apache/druid/segment/virtual/SingleLongInputCachingExpressionColumnValueSelector.java b/processing/src/main/java/org/apache/druid/segment/virtual/SingleLongInputCachingExpressionColumnValueSelector.java index d88c4fbc3c50..1da8be504046 100644 --- a/processing/src/main/java/org/apache/druid/segment/virtual/SingleLongInputCachingExpressionColumnValueSelector.java +++ b/processing/src/main/java/org/apache/druid/segment/virtual/SingleLongInputCachingExpressionColumnValueSelector.java @@ -27,6 +27,7 @@ import org.apache.druid.math.expr.ExpressionType; import org.apache.druid.query.monomorphicprocessing.RuntimeShapeInspector; import org.apache.druid.segment.ColumnValueSelector; +import org.apache.druid.segment.RowIdSupplier; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; import javax.annotation.Nonnull; @@ -36,7 +37,7 @@ * Like {@link ExpressionColumnValueSelector}, but caches the most recently computed value and re-uses it in the case * of runs in the underlying column. This is especially useful for the __time column, where we expect runs. */ -public class SingleLongInputCachingExpressionColumnValueSelector implements ColumnValueSelector +public class SingleLongInputCachingExpressionColumnValueSelector extends BaseExpressionColumnValueSelector { private static final int CACHE_SIZE = 1000; @@ -61,9 +62,12 @@ public class SingleLongInputCachingExpressionColumnValueSelector implements Colu public SingleLongInputCachingExpressionColumnValueSelector( final ColumnValueSelector selector, final Expr expression, - final boolean useLruCache + final boolean useLruCache, + @Nullable RowIdSupplier rowIdSupplier ) { + super(rowIdSupplier); + // Verify expression has just one binding. if (expression.analyzeInputs().getRequiredBindings().size() != 1) { throw new ISE("Expected expression with just one binding"); @@ -77,31 +81,14 @@ public SingleLongInputCachingExpressionColumnValueSelector( @Override public void inspectRuntimeShape(final RuntimeShapeInspector inspector) { + super.inspectRuntimeShape(inspector); inspector.visit("selector", selector); inspector.visit("expression", expression); } - @Override - public double getDouble() - { - return getObject().asDouble(); - } - - @Override - public float getFloat() - { - return (float) getObject().asDouble(); - } - - @Override - public long getLong() - { - return getObject().asLong(); - } - @Nonnull @Override - public ExprEval getObject() + protected ExprEval eval() { // things can still call this even when underlying selector is null (e.g. ExpressionColumnValueSelector#isNull) if (selector.isNull()) { @@ -129,21 +116,6 @@ public ExprEval getObject() return lastOutput; } - @Override - public Class classOfObject() - { - return ExprEval.class; - } - - @Override - public boolean isNull() - { - // It is possible for an expression to have a non-null String value but it can return null when parsed - // as a primitive long/float/double. - // ExprEval.isNumericNull checks whether the parsed primitive value is null or not. - return getObject().isNumericNull(); - } - public class LruEvalCache { private final Long2ObjectLinkedOpenHashMap m = new Long2ObjectLinkedOpenHashMap<>(); diff --git a/processing/src/main/java/org/apache/druid/segment/virtual/SingleStringInputCachingExpressionColumnValueSelector.java b/processing/src/main/java/org/apache/druid/segment/virtual/SingleStringInputCachingExpressionColumnValueSelector.java index 22700f20fa19..a3e535735573 100644 --- a/processing/src/main/java/org/apache/druid/segment/virtual/SingleStringInputCachingExpressionColumnValueSelector.java +++ b/processing/src/main/java/org/apache/druid/segment/virtual/SingleStringInputCachingExpressionColumnValueSelector.java @@ -28,9 +28,9 @@ import org.apache.druid.math.expr.ExpressionType; import org.apache.druid.math.expr.InputBindings; import org.apache.druid.query.monomorphicprocessing.RuntimeShapeInspector; -import org.apache.druid.segment.ColumnValueSelector; import org.apache.druid.segment.DimensionDictionarySelector; import org.apache.druid.segment.DimensionSelector; +import org.apache.druid.segment.RowIdSupplier; import org.apache.druid.segment.data.IndexedInts; import javax.annotation.Nullable; @@ -39,7 +39,7 @@ * Like {@link ExpressionColumnValueSelector}, but caches results for the first CACHE_SIZE dictionary IDs of * a string column. Must only be used on selectors with dictionaries. */ -public class SingleStringInputCachingExpressionColumnValueSelector implements ColumnValueSelector +public class SingleStringInputCachingExpressionColumnValueSelector extends BaseExpressionColumnValueSelector { private static final int CACHE_SIZE = 1000; @@ -53,9 +53,12 @@ public class SingleStringInputCachingExpressionColumnValueSelector implements Co public SingleStringInputCachingExpressionColumnValueSelector( final DimensionSelector selector, - final Expr expression + final Expr expression, + @Nullable final RowIdSupplier rowIdSupplier ) { + super(rowIdSupplier); + // Verify expression has just one binding. if (expression.analyzeInputs().getRequiredBindings().size() != 1) { throw new ISE("Expected expression with just one binding"); @@ -81,45 +84,13 @@ public SingleStringInputCachingExpressionColumnValueSelector( @Override public void inspectRuntimeShape(final RuntimeShapeInspector inspector) { + super.inspectRuntimeShape(inspector); inspector.visit("selector", selector); inspector.visit("expression", expression); } @Override - public double getDouble() - { - // No Assert for null handling as ExprEval already have it. - return eval().asDouble(); - } - - @Override - public float getFloat() - { - // No Assert for null handling as ExprEval already have it. - return (float) eval().asDouble(); - } - - @Override - public long getLong() - { - // No Assert for null handling as ExprEval already have it. - return eval().asLong(); - } - - @Nullable - @Override - public ExprEval getObject() - { - return eval(); - } - - @Override - public Class classOfObject() - { - return ExprEval.class; - } - - private ExprEval eval() + protected ExprEval eval() { final IndexedInts row = selector.getRow(); diff --git a/processing/src/main/java/org/apache/druid/segment/virtual/VirtualizedColumnSelectorFactory.java b/processing/src/main/java/org/apache/druid/segment/virtual/VirtualizedColumnSelectorFactory.java index 9dba14e91e4e..0bfcea8c2c29 100644 --- a/processing/src/main/java/org/apache/druid/segment/virtual/VirtualizedColumnSelectorFactory.java +++ b/processing/src/main/java/org/apache/druid/segment/virtual/VirtualizedColumnSelectorFactory.java @@ -24,8 +24,11 @@ import org.apache.druid.segment.ColumnSelectorFactory; import org.apache.druid.segment.ColumnValueSelector; import org.apache.druid.segment.DimensionSelector; +import org.apache.druid.segment.RowIdSupplier; import org.apache.druid.segment.VirtualColumns; +import javax.annotation.Nullable; + /** * {@link ColumnSelectorFactory} which can create selectors for both virtual and non-virtual columns */ @@ -61,4 +64,11 @@ public ColumnValueSelector makeColumnValueSelector(String columnName) return baseFactory.makeColumnValueSelector(columnName); } } + + @Nullable + @Override + public RowIdSupplier getRowIdSupplier() + { + return baseFactory.getRowIdSupplier(); + } } From 84eeebcab2c2e897b746fe53e17697f4d46de8a0 Mon Sep 17 00:00:00 2001 From: Gian Merlino Date: Sun, 29 May 2022 15:00:10 -0700 Subject: [PATCH 2/2] Fix tests. --- .../segment/virtual/VirtualColumnsTest.java | 68 +++++++++++++++---- 1 file changed, 54 insertions(+), 14 deletions(-) diff --git a/processing/src/test/java/org/apache/druid/segment/virtual/VirtualColumnsTest.java b/processing/src/test/java/org/apache/druid/segment/virtual/VirtualColumnsTest.java index cb46c634d012..607324487eda 100644 --- a/processing/src/test/java/org/apache/druid/segment/virtual/VirtualColumnsTest.java +++ b/processing/src/test/java/org/apache/druid/segment/virtual/VirtualColumnsTest.java @@ -56,6 +56,11 @@ import org.junit.Rule; import org.junit.Test; import org.junit.rules.ExpectedException; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.junit.MockitoJUnit; +import org.mockito.junit.MockitoRule; +import org.mockito.quality.Strictness; import javax.annotation.Nullable; import java.util.Arrays; @@ -68,6 +73,12 @@ public class VirtualColumnsTest extends InitializedNullHandlingTest @Rule public ExpectedException expectedException = ExpectedException.none(); + @Rule + public MockitoRule mockitoRule = MockitoJUnit.rule().strictness(Strictness.STRICT_STUBS); + + @Mock + public ColumnSelectorFactory baseColumnSelectorFactory; + @Test public void testExists() { @@ -197,24 +208,35 @@ public void testNonExistentSelector() expectedException.expect(IllegalArgumentException.class); expectedException.expectMessage("No such virtual column[bar]"); - virtualColumns.makeColumnValueSelector("bar", null); + virtualColumns.makeColumnValueSelector("bar", baseColumnSelectorFactory); } @Test public void testMakeSelectors() { + Mockito.when(baseColumnSelectorFactory.getRowIdSupplier()).thenReturn(null); + final VirtualColumns virtualColumns = makeVirtualColumns(); - final BaseObjectColumnValueSelector objectSelector = virtualColumns.makeColumnValueSelector("expr", null); + final BaseObjectColumnValueSelector objectSelector = virtualColumns.makeColumnValueSelector( + "expr", + baseColumnSelectorFactory + ); final DimensionSelector dimensionSelector = virtualColumns.makeDimensionSelector( new DefaultDimensionSpec("expr", "x"), - null + baseColumnSelectorFactory ); final DimensionSelector extractionDimensionSelector = virtualColumns.makeDimensionSelector( new ExtractionDimensionSpec("expr", "x", new BucketExtractionFn(1.0, 0.5)), - null + baseColumnSelectorFactory + ); + final BaseFloatColumnValueSelector floatSelector = virtualColumns.makeColumnValueSelector( + "expr", + baseColumnSelectorFactory + ); + final BaseLongColumnValueSelector longSelector = virtualColumns.makeColumnValueSelector( + "expr", + baseColumnSelectorFactory ); - final BaseFloatColumnValueSelector floatSelector = virtualColumns.makeColumnValueSelector("expr", null); - final BaseLongColumnValueSelector longSelector = virtualColumns.makeColumnValueSelector("expr", null); Assert.assertEquals(1L, objectSelector.getObject()); Assert.assertEquals("1", dimensionSelector.lookupName(dimensionSelector.getRow().get(0))); @@ -227,13 +249,22 @@ public void testMakeSelectors() public void testMakeSelectorsWithDotSupport() { final VirtualColumns virtualColumns = makeVirtualColumns(); - final BaseObjectColumnValueSelector objectSelector = virtualColumns.makeColumnValueSelector("foo.5", null); + final BaseObjectColumnValueSelector objectSelector = virtualColumns.makeColumnValueSelector( + "foo.5", + baseColumnSelectorFactory + ); final DimensionSelector dimensionSelector = virtualColumns.makeDimensionSelector( new DefaultDimensionSpec("foo.5", "x"), - null + baseColumnSelectorFactory + ); + final BaseFloatColumnValueSelector floatSelector = virtualColumns.makeColumnValueSelector( + "foo.5", + baseColumnSelectorFactory + ); + final BaseLongColumnValueSelector longSelector = virtualColumns.makeColumnValueSelector( + "foo.5", + baseColumnSelectorFactory ); - final BaseFloatColumnValueSelector floatSelector = virtualColumns.makeColumnValueSelector("foo.5", null); - final BaseLongColumnValueSelector longSelector = virtualColumns.makeColumnValueSelector("foo.5", null); Assert.assertEquals(5L, objectSelector.getObject()); Assert.assertEquals("5", dimensionSelector.lookupName(dimensionSelector.getRow().get(0))); @@ -245,13 +276,22 @@ public void testMakeSelectorsWithDotSupport() public void testMakeSelectorsWithDotSupportBaseNameOnly() { final VirtualColumns virtualColumns = makeVirtualColumns(); - final BaseObjectColumnValueSelector objectSelector = virtualColumns.makeColumnValueSelector("foo", null); + final BaseObjectColumnValueSelector objectSelector = virtualColumns.makeColumnValueSelector( + "foo", + baseColumnSelectorFactory + ); final DimensionSelector dimensionSelector = virtualColumns.makeDimensionSelector( new DefaultDimensionSpec("foo", "x"), - null + baseColumnSelectorFactory + ); + final BaseFloatColumnValueSelector floatSelector = virtualColumns.makeColumnValueSelector( + "foo", + baseColumnSelectorFactory + ); + final BaseLongColumnValueSelector longSelector = virtualColumns.makeColumnValueSelector( + "foo", + baseColumnSelectorFactory ); - final BaseFloatColumnValueSelector floatSelector = virtualColumns.makeColumnValueSelector("foo", null); - final BaseLongColumnValueSelector longSelector = virtualColumns.makeColumnValueSelector("foo", null); Assert.assertEquals(-1L, objectSelector.getObject()); Assert.assertEquals("-1", dimensionSelector.lookupName(dimensionSelector.getRow().get(0)));