From 53da6108b866b8f4cd800438008685277766cbbb Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Fri, 5 Jan 2024 20:46:35 -0800 Subject: [PATCH 1/3] fix bugs with expression virtual column indexes for expression virtual columns which refer to other virtual columns changes: * ColumnIndexSelector now extends ColumnSelector. The only real implementation of ColumnIndexSelector, ColumnSelectorColumnIndexSelector, already has a ColumnSelector, so this isn't very disruptive * removed getColumnNames from ColumnSelector since it was not used * VirtualColumns and VirtualColumn getIndexSupplier method now needs argument of ColumnIndexSelector instead of ColumnSelector, which allows expression virtual columns to correctly recognize other virtual columns, fixing an issue which would incorrectly handle other virtual columns as non-existent columns instead * fixed a bug with sql planner incorrectly not using expression filter for equality filters on columns with extractionFn and no virtual column registry --- .../java/org/apache/druid/math/expr/Expr.java | 17 +++---- .../druid/math/expr/IdentifierExpr.java | 15 +++--- .../query/filter/ColumnIndexSelector.java | 4 +- .../org/apache/druid/segment/ColumnCache.java | 7 --- .../apache/druid/segment/ColumnSelector.java | 7 --- .../ColumnSelectorColumnIndexSelector.java | 16 +++++-- ...eprecatedQueryableIndexColumnSelector.java | 7 --- .../apache/druid/segment/VirtualColumn.java | 12 +++-- .../apache/druid/segment/VirtualColumns.java | 8 +++- .../virtual/ExpressionVirtualColumn.java | 9 ++-- .../virtual/FallbackVirtualColumn.java | 9 ++-- .../virtual/ListFilteredVirtualColumn.java | 17 +++---- .../virtual/NestedFieldVirtualColumn.java | 5 +- .../druid/segment/TestColumnSelector.java | 8 ---- .../druid/segment/filter/BaseFilterTest.java | 3 ++ .../segment/filter/EqualityFilterTests.java | 43 +++++++++++++++++ .../filter/ExtractionDimFilterTest.java | 8 ++++ .../druid/segment/filter/NullFilterTests.java | 27 +++++++++++ .../virtual/DummyStringVirtualColumn.java | 10 ++-- .../virtual/FallbackVirtualColumnTest.java | 16 +++++-- .../sql/calcite/expression/Expressions.java | 15 ++++-- .../sql/calcite/CalciteArraysQueryTest.java | 46 +++++++++++++++++++ 22 files changed, 220 insertions(+), 89 deletions(-) diff --git a/processing/src/main/java/org/apache/druid/math/expr/Expr.java b/processing/src/main/java/org/apache/druid/math/expr/Expr.java index 3511abb758a2..4dec1fa74a91 100644 --- a/processing/src/main/java/org/apache/druid/math/expr/Expr.java +++ b/processing/src/main/java/org/apache/druid/math/expr/Expr.java @@ -28,9 +28,8 @@ import org.apache.druid.java.util.common.ISE; import org.apache.druid.math.expr.vector.ExprVectorProcessor; import org.apache.druid.query.cache.CacheKeyBuilder; -import org.apache.druid.segment.ColumnSelector; +import org.apache.druid.query.filter.ColumnIndexSelector; import org.apache.druid.segment.column.ColumnCapabilities; -import org.apache.druid.segment.column.ColumnHolder; import org.apache.druid.segment.column.ColumnIndexSupplier; import org.apache.druid.segment.column.ColumnType; import org.apache.druid.segment.index.semantic.DictionaryEncodedValueIndex; @@ -194,7 +193,10 @@ default ExprVectorProcessor asVectorProcessor(VectorInputBindingInspector } @Nullable - default ColumnIndexSupplier asColumnIndexSupplier(ColumnSelector columnSelector, @Nullable ColumnType outputType) + default ColumnIndexSupplier asColumnIndexSupplier( + ColumnIndexSelector columnIndexSelector, + @Nullable ColumnType outputType + ) { final Expr.BindingAnalysis details = analyzeInputs(); if (details.getRequiredBindings().size() == 1) { @@ -202,13 +204,8 @@ default ColumnIndexSupplier asColumnIndexSupplier(ColumnSelector columnSelector, // map over the values of the index. final String column = Iterables.getOnlyElement(details.getRequiredBindings()); - final ColumnHolder holder = columnSelector.getColumnHolder(column); - if (holder == null) { - // column doesn't exist, no index supplier - return null; - } - final ColumnCapabilities capabilities = holder.getCapabilities(); - final ColumnIndexSupplier delegateIndexSupplier = holder.getIndexSupplier(); + final ColumnCapabilities capabilities = columnIndexSelector.getColumnCapabilities(column); + final ColumnIndexSupplier delegateIndexSupplier = columnIndexSelector.getIndexSupplier(column); if (delegateIndexSupplier == null) { return null; } diff --git a/processing/src/main/java/org/apache/druid/math/expr/IdentifierExpr.java b/processing/src/main/java/org/apache/druid/math/expr/IdentifierExpr.java index 9a16acbe92b9..ef222cda61f9 100644 --- a/processing/src/main/java/org/apache/druid/math/expr/IdentifierExpr.java +++ b/processing/src/main/java/org/apache/druid/math/expr/IdentifierExpr.java @@ -23,8 +23,7 @@ import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.math.expr.vector.ExprVectorProcessor; import org.apache.druid.math.expr.vector.VectorProcessors; -import org.apache.druid.segment.ColumnSelector; -import org.apache.druid.segment.column.ColumnHolder; +import org.apache.druid.query.filter.ColumnIndexSelector; import org.apache.druid.segment.column.ColumnIndexSupplier; import org.apache.druid.segment.column.ColumnType; @@ -158,15 +157,13 @@ public ExprVectorProcessor asVectorProcessor(VectorInputBindingInspector insp @Nullable @Override - public ColumnIndexSupplier asColumnIndexSupplier(ColumnSelector columnSelector, @Nullable ColumnType outputType) + public ColumnIndexSupplier asColumnIndexSupplier( + ColumnIndexSelector indexSelector, + @Nullable ColumnType outputType + ) { // identifier just wraps a column, we can return its index supplier directly if the column exists - final ColumnHolder holder = columnSelector.getColumnHolder(binding); - if (holder == null) { - // column doesn't exist, no index supplier - return null; - } - return holder.getIndexSupplier(); + return indexSelector.getIndexSupplier(binding); } @Override diff --git a/processing/src/main/java/org/apache/druid/query/filter/ColumnIndexSelector.java b/processing/src/main/java/org/apache/druid/query/filter/ColumnIndexSelector.java index c0eb6a4ac491..2a9fea71870d 100644 --- a/processing/src/main/java/org/apache/druid/query/filter/ColumnIndexSelector.java +++ b/processing/src/main/java/org/apache/druid/query/filter/ColumnIndexSelector.java @@ -20,14 +20,14 @@ package org.apache.druid.query.filter; import org.apache.druid.collections.bitmap.BitmapFactory; -import org.apache.druid.segment.ColumnInspector; +import org.apache.druid.segment.ColumnSelector; import org.apache.druid.segment.column.ColumnIndexSupplier; import javax.annotation.Nullable; /** */ -public interface ColumnIndexSelector extends ColumnInspector +public interface ColumnIndexSelector extends ColumnSelector { int getNumRows(); diff --git a/processing/src/main/java/org/apache/druid/segment/ColumnCache.java b/processing/src/main/java/org/apache/druid/segment/ColumnCache.java index fb285fa282a5..8ec7ced4ae79 100644 --- a/processing/src/main/java/org/apache/druid/segment/ColumnCache.java +++ b/processing/src/main/java/org/apache/druid/segment/ColumnCache.java @@ -28,7 +28,6 @@ import javax.annotation.Nullable; import java.util.HashMap; -import java.util.List; public class ColumnCache implements ColumnSelector { @@ -45,12 +44,6 @@ public ColumnCache(QueryableIndex index, Closer closer) } - @Override - public List getColumnNames() - { - return index.getColumnNames(); - } - @Nullable @Override public ColumnHolder getColumnHolder(String columnName) diff --git a/processing/src/main/java/org/apache/druid/segment/ColumnSelector.java b/processing/src/main/java/org/apache/druid/segment/ColumnSelector.java index 32a3e95fb403..49319a6713a1 100644 --- a/processing/src/main/java/org/apache/druid/segment/ColumnSelector.java +++ b/processing/src/main/java/org/apache/druid/segment/ColumnSelector.java @@ -23,18 +23,11 @@ import org.apache.druid.segment.column.ColumnHolder; import javax.annotation.Nullable; -import java.util.List; /** */ public interface ColumnSelector extends ColumnInspector { - /** - * This method is apparently no longer used anymore, so deprecating it. - */ - @Deprecated - List getColumnNames(); - @Nullable ColumnHolder getColumnHolder(String columnName); diff --git a/processing/src/main/java/org/apache/druid/segment/ColumnSelectorColumnIndexSelector.java b/processing/src/main/java/org/apache/druid/segment/ColumnSelectorColumnIndexSelector.java index 184138b0ea85..b509d8d1a8bc 100644 --- a/processing/src/main/java/org/apache/druid/segment/ColumnSelectorColumnIndexSelector.java +++ b/processing/src/main/java/org/apache/druid/segment/ColumnSelectorColumnIndexSelector.java @@ -77,7 +77,7 @@ public ColumnIndexSupplier getIndexSupplier(String column) { final ColumnIndexSupplier indexSupplier; if (isVirtualColumn(column)) { - indexSupplier = virtualColumns.getIndexSupplier(column, columnSelector); + indexSupplier = virtualColumns.getIndexSupplier(column, this); } else { final ColumnHolder columnHolder = columnSelector.getColumnHolder(column); // for missing columns we return null here. This allows callers to fabricate an 'all true' or 'all false' @@ -90,9 +90,14 @@ public ColumnIndexSupplier getIndexSupplier(String column) return indexSupplier; } - private boolean isVirtualColumn(final String columnName) + @Nullable + @Override + public ColumnHolder getColumnHolder(String columnName) { - return virtualColumns.getVirtualColumn(columnName) != null; + if (isVirtualColumn(columnName)) { + return null; + } + return columnSelector.getColumnHolder(columnName); } @Nullable @@ -101,4 +106,9 @@ public ColumnCapabilities getColumnCapabilities(String column) { return virtualColumns.getColumnCapabilitiesWithFallback(columnSelector, column); } + + private boolean isVirtualColumn(final String columnName) + { + return virtualColumns.getVirtualColumn(columnName) != null; + } } diff --git a/processing/src/main/java/org/apache/druid/segment/DeprecatedQueryableIndexColumnSelector.java b/processing/src/main/java/org/apache/druid/segment/DeprecatedQueryableIndexColumnSelector.java index 0eb0180ea0c2..73863eef27c8 100644 --- a/processing/src/main/java/org/apache/druid/segment/DeprecatedQueryableIndexColumnSelector.java +++ b/processing/src/main/java/org/apache/druid/segment/DeprecatedQueryableIndexColumnSelector.java @@ -22,7 +22,6 @@ import org.apache.druid.segment.column.ColumnHolder; import javax.annotation.Nullable; -import java.util.List; /** * It likely looks weird that we are creating a new instance of ColumnSelector here that begins its life deprecated @@ -51,12 +50,6 @@ public DeprecatedQueryableIndexColumnSelector(QueryableIndex index) this.index = index; } - @Override - public List getColumnNames() - { - return index.getColumnNames(); - } - @Nullable @Override public ColumnHolder getColumnHolder(String columnName) diff --git a/processing/src/main/java/org/apache/druid/segment/VirtualColumn.java b/processing/src/main/java/org/apache/druid/segment/VirtualColumn.java index e79baa00b1bf..3698a8a731bf 100644 --- a/processing/src/main/java/org/apache/druid/segment/VirtualColumn.java +++ b/processing/src/main/java/org/apache/druid/segment/VirtualColumn.java @@ -23,6 +23,7 @@ import com.fasterxml.jackson.annotation.JsonTypeInfo; import org.apache.druid.java.util.common.Cacheable; import org.apache.druid.query.dimension.DimensionSpec; +import org.apache.druid.query.filter.ColumnIndexSelector; import org.apache.druid.segment.column.ColumnCapabilities; import org.apache.druid.segment.column.ColumnIndexSupplier; import org.apache.druid.segment.data.ReadableOffset; @@ -299,13 +300,16 @@ default ColumnCapabilities capabilities(ColumnInspector inspector, String column /** * Get the {@link ColumnIndexSupplier} for the specified virtual column, with the assistance of a - * {@link ColumnSelector} to allow reading things from segments. If the virtual column has no indexes, this method - * will return null, or may also return a non-null supplier whose methods may return null values - having a supplier - * is no guarantee that the column has indexes. + * {@link ColumnIndexSelector} to allow reading things from segments. If the virtual column has no indexes, this + * method will return null, or may also return a non-null supplier whose methods may return null values - having a + * supplier is no guarantee that the column has indexes. */ @SuppressWarnings("unused") @Nullable - default ColumnIndexSupplier getIndexSupplier(String columnName, ColumnSelector columnSelector) + default ColumnIndexSupplier getIndexSupplier( + String columnName, + ColumnIndexSelector columnIndexSelector + ) { return NoIndexesColumnIndexSupplier.getInstance(); } diff --git a/processing/src/main/java/org/apache/druid/segment/VirtualColumns.java b/processing/src/main/java/org/apache/druid/segment/VirtualColumns.java index 5560d10d4edf..583f0425c2a6 100644 --- a/processing/src/main/java/org/apache/druid/segment/VirtualColumns.java +++ b/processing/src/main/java/org/apache/druid/segment/VirtualColumns.java @@ -33,6 +33,7 @@ import org.apache.druid.query.Query; import org.apache.druid.query.cache.CacheKeyBuilder; import org.apache.druid.query.dimension.DimensionSpec; +import org.apache.druid.query.filter.ColumnIndexSelector; import org.apache.druid.segment.column.ColumnCapabilities; import org.apache.druid.segment.column.ColumnHolder; import org.apache.druid.segment.column.ColumnIndexSupplier; @@ -185,10 +186,13 @@ public VirtualColumn getVirtualColumn(String columnName) * is no guarantee that the column has indexes. */ @Nullable - public ColumnIndexSupplier getIndexSupplier(String columnName, ColumnSelector columnSelector) + public ColumnIndexSupplier getIndexSupplier( + String columnName, + ColumnIndexSelector columnIndexSelector + ) { final VirtualColumn virtualColumn = getVirtualColumnForSelector(columnName); - return virtualColumn.getIndexSupplier(columnName, columnSelector); + return virtualColumn.getIndexSupplier(columnName, columnIndexSelector); } /** diff --git a/processing/src/main/java/org/apache/druid/segment/virtual/ExpressionVirtualColumn.java b/processing/src/main/java/org/apache/druid/segment/virtual/ExpressionVirtualColumn.java index f8ace2ca2f26..589520612398 100644 --- a/processing/src/main/java/org/apache/druid/segment/virtual/ExpressionVirtualColumn.java +++ b/processing/src/main/java/org/apache/druid/segment/virtual/ExpressionVirtualColumn.java @@ -34,8 +34,8 @@ import org.apache.druid.math.expr.Parser; import org.apache.druid.query.cache.CacheKeyBuilder; import org.apache.druid.query.dimension.DimensionSpec; +import org.apache.druid.query.filter.ColumnIndexSelector; import org.apache.druid.segment.ColumnInspector; -import org.apache.druid.segment.ColumnSelector; import org.apache.druid.segment.ColumnSelectorFactory; import org.apache.druid.segment.ColumnValueSelector; import org.apache.druid.segment.DimensionSelector; @@ -244,9 +244,12 @@ public VectorObjectSelector makeVectorObjectSelector(String columnName, VectorCo @Nullable @Override - public ColumnIndexSupplier getIndexSupplier(String columnName, ColumnSelector columnSelector) + public ColumnIndexSupplier getIndexSupplier( + String columnName, + ColumnIndexSelector columnIndexSelector + ) { - return getParsedExpression().get().asColumnIndexSupplier(columnSelector, outputType); + return getParsedExpression().get().asColumnIndexSupplier(columnIndexSelector, outputType); } @Override diff --git a/processing/src/main/java/org/apache/druid/segment/virtual/FallbackVirtualColumn.java b/processing/src/main/java/org/apache/druid/segment/virtual/FallbackVirtualColumn.java index 461a6b68dff7..26ce21b98ee6 100644 --- a/processing/src/main/java/org/apache/druid/segment/virtual/FallbackVirtualColumn.java +++ b/processing/src/main/java/org/apache/druid/segment/virtual/FallbackVirtualColumn.java @@ -25,8 +25,8 @@ import org.apache.druid.java.util.common.Pair; import org.apache.druid.query.cache.CacheKeyBuilder; import org.apache.druid.query.dimension.DimensionSpec; +import org.apache.druid.query.filter.ColumnIndexSelector; import org.apache.druid.segment.ColumnInspector; -import org.apache.druid.segment.ColumnSelector; import org.apache.druid.segment.ColumnSelectorFactory; import org.apache.druid.segment.ColumnValueSelector; import org.apache.druid.segment.DimensionSelector; @@ -189,9 +189,12 @@ public boolean usesDotNotation() @Nullable @Override - public ColumnIndexSupplier getIndexSupplier(String columnName, ColumnSelector columnSelector) + public ColumnIndexSupplier getIndexSupplier( + String columnName, + ColumnIndexSelector indexSelector + ) { - final ColumnHolder columnHolder = columnSelector.getColumnHolder(columnToUse(columnSelector).getDimension()); + final ColumnHolder columnHolder = indexSelector.getColumnHolder(columnToUse(indexSelector).getDimension()); if (columnHolder == null) { return null; } diff --git a/processing/src/main/java/org/apache/druid/segment/virtual/ListFilteredVirtualColumn.java b/processing/src/main/java/org/apache/druid/segment/virtual/ListFilteredVirtualColumn.java index 741db0019612..9262e8a17303 100644 --- a/processing/src/main/java/org/apache/druid/segment/virtual/ListFilteredVirtualColumn.java +++ b/processing/src/main/java/org/apache/druid/segment/virtual/ListFilteredVirtualColumn.java @@ -34,11 +34,11 @@ import org.apache.druid.query.dimension.DefaultDimensionSpec; import org.apache.druid.query.dimension.DimensionSpec; import org.apache.druid.query.dimension.ListFilteredDimensionSpec; +import org.apache.druid.query.filter.ColumnIndexSelector; import org.apache.druid.query.filter.DruidObjectPredicate; import org.apache.druid.query.filter.DruidPredicateFactory; import org.apache.druid.query.ordering.StringComparators; import org.apache.druid.segment.ColumnInspector; -import org.apache.druid.segment.ColumnSelector; import org.apache.druid.segment.ColumnSelectorFactory; import org.apache.druid.segment.ColumnValueSelector; import org.apache.druid.segment.DimensionSelector; @@ -183,7 +183,10 @@ public boolean usesDotNotation() @Nullable @Override - public ColumnIndexSupplier getIndexSupplier(String columnName, ColumnSelector columnSelector) + public ColumnIndexSupplier getIndexSupplier( + String columnName, + ColumnIndexSelector indexSelector + ) { if (delegate.getExtractionFn() != null) { return NoIndexesColumnIndexSupplier.getInstance(); @@ -194,13 +197,7 @@ public ColumnIndexSupplier getIndexSupplier(String columnName, ColumnSelector co @Override public T as(Class clazz) { - - final ColumnHolder holder = columnSelector.getColumnHolder(delegate.getDimension()); - if (holder == null) { - return null; - } - - ColumnIndexSupplier indexSupplier = holder.getIndexSupplier(); + ColumnIndexSupplier indexSupplier = indexSelector.getIndexSupplier(delegate.getDimension()); if (indexSupplier == null) { return null; } @@ -229,7 +226,7 @@ public T as(Class clazz) } // someday maybe we can have a better way to get row count.. - final ColumnHolder time = columnSelector.getColumnHolder(ColumnHolder.TIME_COLUMN_NAME); + final ColumnHolder time = indexSelector.getColumnHolder(ColumnHolder.TIME_COLUMN_NAME); final int numRows = time.getLength(); final Supplier nullValueBitmapSupplier = Suppliers.memoize( () -> underlyingIndex.getBitmapFactory().complement( diff --git a/processing/src/main/java/org/apache/druid/segment/virtual/NestedFieldVirtualColumn.java b/processing/src/main/java/org/apache/druid/segment/virtual/NestedFieldVirtualColumn.java index 0a282f5e4bc6..63b8598ef634 100644 --- a/processing/src/main/java/org/apache/druid/segment/virtual/NestedFieldVirtualColumn.java +++ b/processing/src/main/java/org/apache/druid/segment/virtual/NestedFieldVirtualColumn.java @@ -34,6 +34,7 @@ import org.apache.druid.query.cache.CacheKeyBuilder; import org.apache.druid.query.dimension.DimensionSpec; import org.apache.druid.query.extraction.ExtractionFn; +import org.apache.druid.query.filter.ColumnIndexSelector; import org.apache.druid.query.filter.DruidPredicateFactory; import org.apache.druid.query.filter.ValueMatcher; import org.apache.druid.query.monomorphicprocessing.RuntimeShapeInspector; @@ -1158,10 +1159,10 @@ private void computeVectorsIfNeeded() @Override public ColumnIndexSupplier getIndexSupplier( String columnName, - ColumnSelector selector + ColumnIndexSelector indexSelector ) { - ColumnHolder holder = selector.getColumnHolder(this.columnName); + ColumnHolder holder = indexSelector.getColumnHolder(this.columnName); if (holder == null) { return null; } diff --git a/processing/src/test/java/org/apache/druid/segment/TestColumnSelector.java b/processing/src/test/java/org/apache/druid/segment/TestColumnSelector.java index ea8f85391f5b..244213588e82 100644 --- a/processing/src/test/java/org/apache/druid/segment/TestColumnSelector.java +++ b/processing/src/test/java/org/apache/druid/segment/TestColumnSelector.java @@ -19,14 +19,12 @@ package org.apache.druid.segment; -import com.google.common.collect.Lists; import org.apache.druid.java.util.common.UOE; import org.apache.druid.segment.column.ColumnCapabilities; import org.apache.druid.segment.column.ColumnHolder; import javax.annotation.Nullable; import java.util.LinkedHashMap; -import java.util.List; import java.util.Map; public class TestColumnSelector implements ColumnSelector @@ -46,12 +44,6 @@ public TestColumnSelector addCapabilities(String name, ColumnCapabilities capabi return this; } - @Override - public List getColumnNames() - { - return Lists.newArrayList(holders.keySet()); - } - @Nullable @Override public ColumnHolder getColumnHolder(String columnName) diff --git a/processing/src/test/java/org/apache/druid/segment/filter/BaseFilterTest.java b/processing/src/test/java/org/apache/druid/segment/filter/BaseFilterTest.java index 6bab72f2cdb4..906c5cedc726 100644 --- a/processing/src/test/java/org/apache/druid/segment/filter/BaseFilterTest.java +++ b/processing/src/test/java/org/apache/druid/segment/filter/BaseFilterTest.java @@ -143,6 +143,9 @@ public abstract class BaseFilterTest extends InitializedNullHandlingTest new ExpressionVirtualColumn("vd0-add-sub", "d0 + (d0 - d0)", ColumnType.DOUBLE, TestExprMacroTable.INSTANCE), new ExpressionVirtualColumn("vf0-add-sub", "f0 + (f0 - f0)", ColumnType.FLOAT, TestExprMacroTable.INSTANCE), new ExpressionVirtualColumn("vl0-add-sub", "l0 + (l0 - l0)", ColumnType.LONG, TestExprMacroTable.INSTANCE), + new ExpressionVirtualColumn("double-vd0-add-sub", "vd0 + (vd0 - vd0)", ColumnType.DOUBLE, TestExprMacroTable.INSTANCE), + new ExpressionVirtualColumn("double-vf0-add-sub", "vf0 + (vf0 - vf0)", ColumnType.FLOAT, TestExprMacroTable.INSTANCE), + new ExpressionVirtualColumn("double-vl0-add-sub", "vl0 + (vl0 - vl0)", ColumnType.LONG, TestExprMacroTable.INSTANCE), new ExpressionVirtualColumn("vdim3-concat", "dim3 + dim3", ColumnType.LONG, TestExprMacroTable.INSTANCE), new ExpressionVirtualColumn("nestedArrayLong", "array(arrayLong)", ColumnType.ofArray(ColumnType.LONG_ARRAY), TestExprMacroTable.INSTANCE), new ListFilteredVirtualColumn("allow-dim0", DefaultDimensionSpec.of("dim0"), ImmutableSet.of("3", "4"), true), diff --git a/processing/src/test/java/org/apache/druid/segment/filter/EqualityFilterTests.java b/processing/src/test/java/org/apache/druid/segment/filter/EqualityFilterTests.java index 01a7b2d43edc..3a141d63758f 100644 --- a/processing/src/test/java/org/apache/druid/segment/filter/EqualityFilterTests.java +++ b/processing/src/test/java/org/apache/druid/segment/filter/EqualityFilterTests.java @@ -649,6 +649,23 @@ public void testVirtualNumericColumnNullsAndDefaults() NotDimFilter.of(new EqualityFilter("vd0-add-sub", ColumnType.DOUBLE, 0.0, null)), ImmutableList.of("1", "3", "4", "5") ); + + // virtual column that refers to another virtual column + assertFilterMatches(new EqualityFilter("double-vf0-add-sub", ColumnType.STRING, "0", null), ImmutableList.of("0", "4")); + assertFilterMatches(new EqualityFilter("double-vd0-add-sub", ColumnType.STRING, "0", null), ImmutableList.of("0", "2")); + assertFilterMatches(new EqualityFilter("double-vl0-add-sub", ColumnType.STRING, "0", null), ImmutableList.of("0", "3")); + + assertFilterMatches(new EqualityFilter("double-vf0-add-sub", ColumnType.FLOAT, 0f, null), ImmutableList.of("0", "4")); + assertFilterMatches( + NotDimFilter.of(new EqualityFilter("double-vf0-add-sub", ColumnType.FLOAT, 0f, null)), + ImmutableList.of("1", "2", "3", "5") + ); + assertFilterMatches(new EqualityFilter("double-vd0-add-sub", ColumnType.DOUBLE, 0.0, null), ImmutableList.of("0", "2")); + assertFilterMatches( + NotDimFilter.of(new EqualityFilter("double-vd0-add-sub", ColumnType.DOUBLE, 0.0, null)), + ImmutableList.of("1", "3", "4", "5") + ); + assertFilterMatches(new EqualityFilter("vl0", ColumnType.LONG, 0L, null), ImmutableList.of("0", "3")); assertFilterMatches( NotDimFilter.of(new EqualityFilter("vl0", ColumnType.LONG, 0L, null)), @@ -713,6 +730,32 @@ public void testVirtualNumericColumnNullsAndDefaults() assertFilterMatches(new EqualityFilter("vf0-add-sub", ColumnType.STRING, "0", null), ImmutableList.of("0")); assertFilterMatches(new EqualityFilter("vd0-add-sub", ColumnType.STRING, "0", null), ImmutableList.of("0")); assertFilterMatches(new EqualityFilter("vl0-add-sub", ColumnType.STRING, "0", null), ImmutableList.of("0")); + + assertFilterMatches(new EqualityFilter("double-vf0-add-sub", ColumnType.FLOAT, 0f, null), ImmutableList.of("0")); + assertFilterMatches( + NotDimFilter.of(new EqualityFilter("double-vf0-add-sub", ColumnType.FLOAT, 0f, null)), + NullHandling.sqlCompatible() + ? ImmutableList.of("1", "2", "3", "5") + : ImmutableList.of("1", "2", "3", "4", "5") + ); + assertFilterMatches(new EqualityFilter("double-vd0-add-sub", ColumnType.DOUBLE, 0.0, null), ImmutableList.of("0")); + assertFilterMatches( + NotDimFilter.of(new EqualityFilter("double-vd0-add-sub", ColumnType.DOUBLE, 0.0, null)), + NullHandling.sqlCompatible() + ? ImmutableList.of("1", "3", "4", "5") + : ImmutableList.of("1", "2", "3", "4", "5") + ); + assertFilterMatches(new EqualityFilter("double-vl0-add-sub", ColumnType.LONG, 0L, null), ImmutableList.of("0")); + assertFilterMatches( + NotDimFilter.of(new EqualityFilter("double-vl0-add-sub", ColumnType.LONG, 0L, null)), + NullHandling.sqlCompatible() + ? ImmutableList.of("1", "2", "4", "5") + : ImmutableList.of("1", "2", "3", "4", "5") + ); + + assertFilterMatches(new EqualityFilter("double-vf0-add-sub", ColumnType.STRING, "0", null), ImmutableList.of("0")); + assertFilterMatches(new EqualityFilter("double-vd0-add-sub", ColumnType.STRING, "0", null), ImmutableList.of("0")); + assertFilterMatches(new EqualityFilter("double-vl0-add-sub", ColumnType.STRING, "0", null), ImmutableList.of("0")); } } } diff --git a/processing/src/test/java/org/apache/druid/segment/filter/ExtractionDimFilterTest.java b/processing/src/test/java/org/apache/druid/segment/filter/ExtractionDimFilterTest.java index 692a2d499f40..b890e666d36a 100644 --- a/processing/src/test/java/org/apache/druid/segment/filter/ExtractionDimFilterTest.java +++ b/processing/src/test/java/org/apache/druid/segment/filter/ExtractionDimFilterTest.java @@ -36,6 +36,7 @@ import org.apache.druid.query.filter.SelectorDimFilter; import org.apache.druid.segment.column.ColumnCapabilities; import org.apache.druid.segment.column.ColumnCapabilitiesImpl; +import org.apache.druid.segment.column.ColumnHolder; import org.apache.druid.segment.column.ColumnIndexSupplier; import org.apache.druid.segment.column.ColumnType; import org.apache.druid.segment.data.BitmapSerdeFactory; @@ -88,6 +89,13 @@ public ExtractionDimFilterTest(BitmapFactory bitmapFactory, BitmapSerdeFactory b private final ColumnIndexSelector BITMAP_INDEX_SELECTOR = new ColumnIndexSelector() { + @Nullable + @Override + public ColumnHolder getColumnHolder(String columnName) + { + return null; + } + @Nullable @Override public ColumnCapabilities getColumnCapabilities(String column) diff --git a/processing/src/test/java/org/apache/druid/segment/filter/NullFilterTests.java b/processing/src/test/java/org/apache/druid/segment/filter/NullFilterTests.java index 443f655baaed..9cb64f75d780 100644 --- a/processing/src/test/java/org/apache/druid/segment/filter/NullFilterTests.java +++ b/processing/src/test/java/org/apache/druid/segment/filter/NullFilterTests.java @@ -248,6 +248,24 @@ public void testVirtualNumericColumnNullsAndDefaults() NotDimFilter.of(NullFilter.forColumn("vl0-add-sub")), ImmutableList.of("0", "1", "2", "3", "4", "5") ); + + assertFilterMatches(NullFilter.forColumn("double-vf0-add-sub"), ImmutableList.of()); + assertFilterMatches( + NotDimFilter.of(NullFilter.forColumn("double-vf0-add-sub")), + ImmutableList.of("0", "1", "2", "3", "4", "5") + ); + + assertFilterMatches(NullFilter.forColumn("double-vd0-add-sub"), ImmutableList.of()); + assertFilterMatches( + NotDimFilter.of(NullFilter.forColumn("double-vd0-add-sub")), + ImmutableList.of("0", "1", "2", "3", "4", "5") + ); + + assertFilterMatches(NullFilter.forColumn("double-vl0-add-sub"), ImmutableList.of()); + assertFilterMatches( + NotDimFilter.of(NullFilter.forColumn("double-vl0-add-sub")), + ImmutableList.of("0", "1", "2", "3", "4", "5") + ); } else { assertFilterMatches(NullFilter.forColumn("vf0"), ImmutableList.of("4")); assertFilterMatches(NotDimFilter.of(NullFilter.forColumn("vf0")), ImmutableList.of("0", "1", "2", "3", "5")); @@ -269,6 +287,15 @@ public void testVirtualNumericColumnNullsAndDefaults() assertFilterMatches(NullFilter.forColumn("vl0-add-sub"), ImmutableList.of("3")); assertFilterMatches(NotDimFilter.of(NullFilter.forColumn("vl0-add-sub")), ImmutableList.of("0", "1", "2", "4", "5")); + + assertFilterMatches(NullFilter.forColumn("double-vf0-add-sub"), ImmutableList.of("4")); + assertFilterMatches(NotDimFilter.of(NullFilter.forColumn("double-vf0-add-sub")), ImmutableList.of("0", "1", "2", "3", "5")); + + assertFilterMatches(NullFilter.forColumn("double-vd0-add-sub"), ImmutableList.of("2")); + assertFilterMatches(NotDimFilter.of(NullFilter.forColumn("double-vd0-add-sub")), ImmutableList.of("0", "1", "3", "4", "5")); + + assertFilterMatches(NullFilter.forColumn("vl0-add-sub"), ImmutableList.of("3")); + assertFilterMatches(NotDimFilter.of(NullFilter.forColumn("double-vl0-add-sub")), ImmutableList.of("0", "1", "2", "4", "5")); } } } diff --git a/processing/src/test/java/org/apache/druid/segment/virtual/DummyStringVirtualColumn.java b/processing/src/test/java/org/apache/druid/segment/virtual/DummyStringVirtualColumn.java index 27f5113eea29..817fb8446a39 100644 --- a/processing/src/test/java/org/apache/druid/segment/virtual/DummyStringVirtualColumn.java +++ b/processing/src/test/java/org/apache/druid/segment/virtual/DummyStringVirtualColumn.java @@ -22,6 +22,7 @@ import org.apache.druid.java.util.common.IAE; import org.apache.druid.query.dimension.DefaultDimensionSpec; import org.apache.druid.query.dimension.DimensionSpec; +import org.apache.druid.query.filter.ColumnIndexSelector; import org.apache.druid.query.filter.DruidPredicateFactory; import org.apache.druid.query.filter.ValueMatcher; import org.apache.druid.query.monomorphicprocessing.RuntimeShapeInspector; @@ -169,7 +170,7 @@ public ColumnValueSelector makeColumnValueSelector( @Override public ColumnIndexSupplier getIndexSupplier( String columnName, - ColumnSelector columnSelector + ColumnIndexSelector indexSelector ) { return new ColumnIndexSupplier() @@ -180,12 +181,11 @@ public ColumnIndexSupplier getIndexSupplier( public T as(Class clazz) { if (enableBitmaps) { - ColumnHolder holder = columnSelector.getColumnHolder(baseColumnName); - if (holder == null) { + ColumnIndexSupplier supplier = indexSelector.getIndexSupplier(baseColumnName); + if (supplier == null) { return null; } - - return holder.getIndexSupplier().as(clazz); + return supplier.as(clazz); } else { return null; } diff --git a/processing/src/test/java/org/apache/druid/segment/virtual/FallbackVirtualColumnTest.java b/processing/src/test/java/org/apache/druid/segment/virtual/FallbackVirtualColumnTest.java index 72de5466ff64..9a04b906e91d 100644 --- a/processing/src/test/java/org/apache/druid/segment/virtual/FallbackVirtualColumnTest.java +++ b/processing/src/test/java/org/apache/druid/segment/virtual/FallbackVirtualColumnTest.java @@ -19,16 +19,19 @@ package org.apache.druid.segment.virtual; +import org.apache.druid.collections.bitmap.RoaringBitmapFactory; import org.apache.druid.common.config.NullHandling; import org.apache.druid.query.cache.CacheKeyBuilder; import org.apache.druid.query.dimension.DefaultDimensionSpec; import org.apache.druid.query.dimension.DimensionSpec; import org.apache.druid.query.extraction.ExtractionFn; +import org.apache.druid.segment.ColumnSelectorColumnIndexSelector; import org.apache.druid.segment.ConstantDimensionSelector; import org.apache.druid.segment.DimensionSelector; import org.apache.druid.segment.IdLookup; import org.apache.druid.segment.TestColumnSelector; import org.apache.druid.segment.TestColumnSelectorFactory; +import org.apache.druid.segment.VirtualColumns; import org.apache.druid.segment.column.BaseColumn; import org.apache.druid.segment.column.ColumnCapabilities; import org.apache.druid.segment.column.ColumnCapabilitiesImpl; @@ -317,17 +320,22 @@ public void testGetIndexSupplier() .addCapabilities("colA", ColumnCapabilitiesImpl.createDefault()) .addCapabilities("colB", ColumnCapabilitiesImpl.createDefault()) .addCapabilities("colC", ColumnCapabilitiesImpl.createDefault()); + final ColumnSelectorColumnIndexSelector columnIndexSelector = new ColumnSelectorColumnIndexSelector( + RoaringBitmapFactory.INSTANCE, + VirtualColumns.EMPTY, + selectorFactory + ); - Assert.assertSame(colA, col.getIndexSupplier("abcd", selectorFactory)); + Assert.assertSame(colA, col.getIndexSupplier("abcd", columnIndexSelector)); selectorFactory.addCapabilities("colA", null); - Assert.assertSame(colB, col.getIndexSupplier("abcd", selectorFactory)); + Assert.assertSame(colB, col.getIndexSupplier("abcd", columnIndexSelector)); selectorFactory.addCapabilities("colB", null); - Assert.assertSame(colC, col.getIndexSupplier("abcd", selectorFactory)); + Assert.assertSame(colC, col.getIndexSupplier("abcd", columnIndexSelector)); selectorFactory.addCapabilities("colC", null); - Assert.assertSame(colA, col.getIndexSupplier("abcd", selectorFactory)); + Assert.assertSame(colA, col.getIndexSupplier("abcd", columnIndexSelector)); } diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/expression/Expressions.java b/sql/src/main/java/org/apache/druid/sql/calcite/expression/Expressions.java index 538663395d0f..3da2b5fc2f28 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/expression/Expressions.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/expression/Expressions.java @@ -562,9 +562,17 @@ private static DimFilter toSimpleLeafFilter( druidExpression.getSimpleExtraction().getExtractionFn() ); } else { - if (virtualColumnRegistry != null && druidExpression.getSimpleExtraction().getExtractionFn() != null) { - String column = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(druidExpression, druidExpression.getDruidType()); - equalFilter = NullFilter.forColumn(column); + if (druidExpression.getSimpleExtraction().getExtractionFn() != null) { + if (virtualColumnRegistry != null) { + String column = virtualColumnRegistry.getOrCreateVirtualColumnForExpression( + druidExpression, + druidExpression.getDruidType() + ); + equalFilter = NullFilter.forColumn(column); + } else { + // virtual column registry unavailable, fallback to expression filter + return null; + } } else { equalFilter = NullFilter.forColumn(druidExpression.getDirectColumn()); } @@ -770,6 +778,7 @@ private static DimFilter toSimpleLeafFilter( ); } else { // if this happens for some reason, bail and use an expression filter + return null; } } diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteArraysQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteArraysQueryTest.java index 54e28f0b01f5..344ea269121e 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteArraysQueryTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteArraysQueryTest.java @@ -7291,4 +7291,50 @@ public void testArrayToMvPostaggInline() .build() ); } + + @Test + public void testUnnestExtractionFn() + { + // This tells the test to skip generating (vectorize = force) path + // Generates only 1 native query with vectorize = false + skipVectorize(); + // This tells that both vectorize = force and vectorize = false takes the same path of non vectorization + // Generates 2 native queries with 2 different values of vectorize + cannotVectorize(); + testQuery( + "SELECT substring(d3,1) FROM druid.numfoo, UNNEST(MV_TO_ARRAY(dim3)) as unnested (d3) WHERE substring(d3,1) <> 'b'", + QUERY_CONTEXT_UNNEST, + ImmutableList.of( + Druids.newScanQueryBuilder() + .dataSource(UnnestDataSource.create( + new TableDataSource(CalciteTests.DATASOURCE3), + expressionVirtualColumn("j0.unnest", "\"dim3\"", ColumnType.STRING), + NullHandling.sqlCompatible() + ? expressionFilter("(substring(\"j0.unnest\", 0, -1) != 'b')") + : not(selector("j0.unnest", "b", new SubstringDimExtractionFn(0, null))) + )) + .virtualColumns(expressionVirtualColumn("v0", "substring(\"j0.unnest\", 0, -1)", ColumnType.STRING)) + .intervals(querySegmentSpec(Filtration.eternity())) + .resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST) + .legacy(false) + .context(QUERY_CONTEXT_UNNEST) + .columns(ImmutableList.of("v0")) + .build() + ), + useDefault ? + ImmutableList.of( + new Object[]{"a"}, + new Object[]{"c"}, + new Object[]{"d"}, + new Object[]{""}, + new Object[]{""}, + new Object[]{""} + ) : + ImmutableList.of( + new Object[]{"a"}, + new Object[]{"c"}, + new Object[]{"d"} + ) + ); + } } From 09cf3559fd37521a76a589da3ac065d99847167a Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Fri, 5 Jan 2024 21:07:42 -0800 Subject: [PATCH 2/3] more test --- .../sql/calcite/CalciteArraysQueryTest.java | 39 +++++++++++++++++-- 1 file changed, 35 insertions(+), 4 deletions(-) diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteArraysQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteArraysQueryTest.java index 344ea269121e..e4f609ef9ffc 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteArraysQueryTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteArraysQueryTest.java @@ -7295,11 +7295,7 @@ public void testArrayToMvPostaggInline() @Test public void testUnnestExtractionFn() { - // This tells the test to skip generating (vectorize = force) path - // Generates only 1 native query with vectorize = false skipVectorize(); - // This tells that both vectorize = force and vectorize = false takes the same path of non vectorization - // Generates 2 native queries with 2 different values of vectorize cannotVectorize(); testQuery( "SELECT substring(d3,1) FROM druid.numfoo, UNNEST(MV_TO_ARRAY(dim3)) as unnested (d3) WHERE substring(d3,1) <> 'b'", @@ -7337,4 +7333,39 @@ public void testUnnestExtractionFn() ) ); } + + @Test + public void testUnnestExtractionFnNull() + { + skipVectorize(); + cannotVectorize(); + testQuery( + "SELECT substring(d3,1) FROM druid.numfoo, UNNEST(MV_TO_ARRAY(dim3)) as unnested (d3) WHERE substring(d3,1) is not null", + QUERY_CONTEXT_UNNEST, + ImmutableList.of( + Druids.newScanQueryBuilder() + .dataSource(UnnestDataSource.create( + new TableDataSource(CalciteTests.DATASOURCE3), + expressionVirtualColumn("j0.unnest", "\"dim3\"", ColumnType.STRING), + NullHandling.sqlCompatible() + ? expressionFilter("notnull(substring(\"j0.unnest\", 0, -1))") + : not(selector("j0.unnest", null, new SubstringDimExtractionFn(0, null))) + )) + .virtualColumns(expressionVirtualColumn("v0", "substring(\"j0.unnest\", 0, -1)", ColumnType.STRING)) + .intervals(querySegmentSpec(Filtration.eternity())) + .resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST) + .legacy(false) + .context(QUERY_CONTEXT_UNNEST) + .columns(ImmutableList.of("v0")) + .build() + ), + ImmutableList.of( + new Object[]{"a"}, + new Object[]{"b"}, + new Object[]{"b"}, + new Object[]{"c"}, + new Object[]{"d"} + ) + ); + } } From 0f3ce8aebe7768dbe7f98288806942ae24a89e0a Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Fri, 5 Jan 2024 22:44:07 -0800 Subject: [PATCH 3/3] fix compile --- .../apache/druid/benchmark/MockColumnIndexSelector.java | 8 ++++++++ .../src/main/java/org/apache/druid/math/expr/Expr.java | 2 +- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/benchmarks/src/test/java/org/apache/druid/benchmark/MockColumnIndexSelector.java b/benchmarks/src/test/java/org/apache/druid/benchmark/MockColumnIndexSelector.java index 9e6bdc57eefb..33b824649bf8 100644 --- a/benchmarks/src/test/java/org/apache/druid/benchmark/MockColumnIndexSelector.java +++ b/benchmarks/src/test/java/org/apache/druid/benchmark/MockColumnIndexSelector.java @@ -22,6 +22,7 @@ import org.apache.druid.collections.bitmap.BitmapFactory; import org.apache.druid.query.filter.ColumnIndexSelector; import org.apache.druid.segment.column.ColumnCapabilities; +import org.apache.druid.segment.column.ColumnHolder; import org.apache.druid.segment.column.ColumnIndexSupplier; import javax.annotation.Nullable; @@ -58,6 +59,13 @@ public ColumnIndexSupplier getIndexSupplier(String column) return indexSupplier; } + @Nullable + @Override + public ColumnHolder getColumnHolder(String columnName) + { + return null; + } + @Nullable @Override public ColumnCapabilities getColumnCapabilities(String column) diff --git a/processing/src/main/java/org/apache/druid/math/expr/Expr.java b/processing/src/main/java/org/apache/druid/math/expr/Expr.java index 4dec1fa74a91..9471bb24eb5f 100644 --- a/processing/src/main/java/org/apache/druid/math/expr/Expr.java +++ b/processing/src/main/java/org/apache/druid/math/expr/Expr.java @@ -204,7 +204,6 @@ default ColumnIndexSupplier asColumnIndexSupplier( // map over the values of the index. final String column = Iterables.getOnlyElement(details.getRequiredBindings()); - final ColumnCapabilities capabilities = columnIndexSelector.getColumnCapabilities(column); final ColumnIndexSupplier delegateIndexSupplier = columnIndexSelector.getIndexSupplier(column); if (delegateIndexSupplier == null) { return null; @@ -213,6 +212,7 @@ default ColumnIndexSupplier asColumnIndexSupplier( DictionaryEncodedValueIndex.class ); + final ColumnCapabilities capabilities = columnIndexSelector.getColumnCapabilities(column); final ExpressionType inputType = ExpressionType.fromColumnTypeStrict(capabilities); final ColumnType outType; if (outputType == null) {