Apache
/
druid
mirror of https://github.com/apache/druid.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix VirtualColumn related issues in window expressions (#15119) 

for some exotic queries like:

  SELECT
  	'_'||dim1,
    MIN(cast(0 as double)) OVER (),
    MIN(cast((cnt||cnt) as bigint)) OVER ()
  FROM foo
the compilation have resulted in NPE -s mostly because VirtualColumn -s were not handled properly
This commit is contained in:
Zoltan Haindrich 2023-10-23 10:35:59 +02:00 committed by GitHub
parent c8e458452d
commit b95035f183
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
23 changed files with 449 additions and 61 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
processing/src/main/java/org/apache/druid/frame/write/columnar/FrameColumnWriters.java
View File

@ -19,7 +19,6 @@
package org.apache.druid.frame.write.columnar;
import org.apache.druid.common.config.NullHandling;
import org.apache.druid.frame.allocation.MemoryAllocator;
import org.apache.druid.frame.write.UnsupportedColumnTypeException;
import org.apache.druid.java.util.common.ISE;
@ -167,9 +166,7 @@ public class FrameColumnWriters
private static boolean hasNullsForNumericWriter(final ColumnCapabilities capabilities)
{
if (NullHandling.replaceWithDefault()) {
return false;
} else if (capabilities == null) {
if (capabilities == null) {
return true;
} else if (capabilities.getType().isNumeric()) {
return capabilities.hasNulls().isMaybeTrue();

1
processing/src/main/java/org/apache/druid/query/operator/OperatorFactory.java
View File

@ -33,6 +33,7 @@ import org.apache.druid.query.operator.window.WindowOperatorFactory;
@JsonSubTypes.Type(name = "naivePartition", value = NaivePartitioningOperatorFactory.class),
@JsonSubTypes.Type(name = "naiveSort", value = NaiveSortOperatorFactory.class),
@JsonSubTypes.Type(name = "window", value = WindowOperatorFactory.class),
@JsonSubTypes.Type(name = "scan", value = ScanOperatorFactory.class),
})
public interface OperatorFactory
{

14
processing/src/main/java/org/apache/druid/query/operator/WindowOperatorQuery.java
View File

@ -196,7 +196,7 @@ public class WindowOperatorQuery extends BaseQuery<RowsAndColumns>
{
return new WindowOperatorQuery(
getDataSource(),
getQuerySegmentSpec(),
spec,
getContext(),
rowSignature,
operators,
@ -217,6 +217,18 @@ public class WindowOperatorQuery extends BaseQuery<RowsAndColumns>
);
}
public Query<RowsAndColumns> withOperators(List<OperatorFactory> operators)
{
return new WindowOperatorQuery(
getDataSource(),
getQuerySegmentSpec(),
getContext(),
rowSignature,
operators,
leafOperators
);
}
@Override
public boolean equals(Object o)
{

10
processing/src/main/java/org/apache/druid/query/operator/WindowOperatorQueryQueryToolChest.java
View File

@ -54,9 +54,17 @@ public class WindowOperatorQueryQueryToolChest extends QueryToolChest<RowsAndCol
(queryPlus, responseContext) -> {
final WindowOperatorQuery query = (WindowOperatorQuery) queryPlus.getQuery();
final List<OperatorFactory> opFactories = query.getOperators();
if (opFactories.isEmpty()) {
return runner.run(queryPlus, responseContext);
}
Supplier<Operator> opSupplier = () -> {
Operator retVal = new SequenceOperator(runner.run(queryPlus, responseContext));
Operator retVal = new SequenceOperator(
runner.run(
queryPlus.withQuery(query.withOperators(new ArrayList<OperatorFactory>())),
responseContext
)
);
for (OperatorFactory operatorFactory : opFactories) {
retVal = operatorFactory.wrap(retVal);
}

45
processing/src/main/java/org/apache/druid/query/rowsandcols/LazilyDecoratedRowsAndColumns.java
View File

@ -19,6 +19,7 @@
package org.apache.druid.query.rowsandcols;
import com.google.common.collect.ImmutableList;
import org.apache.druid.frame.Frame;
import org.apache.druid.frame.FrameType;
import org.apache.druid.frame.allocation.ArenaMemoryAllocatorFactory;
@ -101,6 +102,11 @@ public class LazilyDecoratedRowsAndColumns implements RowsAndColumns
return viewableColumns == null ? base.getColumnNames() : viewableColumns;
}
public RowsAndColumns getBase()
{
return base;
}
@Override
public int numRows()
{
@ -115,7 +121,6 @@ public class LazilyDecoratedRowsAndColumns implements RowsAndColumns
if (viewableColumns != null && !viewableColumns.contains(name)) {
return null;
}
maybeMaterialize();
return base.findColumn(name);
}
@ -158,7 +163,7 @@ public class LazilyDecoratedRowsAndColumns implements RowsAndColumns
private void maybeMaterialize()
{
if (!(interval == null && filter == null && limit == -1 && ordering == null)) {
if (needsMaterialization()) {
final Pair<byte[], RowSignature> thePair = materialize();
if (thePair == null) {
reset(new EmptyRowsAndColumns());
@ -168,6 +173,11 @@ public class LazilyDecoratedRowsAndColumns implements RowsAndColumns
}
}
private boolean needsMaterialization()
{
return interval != null || filter != null || limit != -1 || ordering != null || virtualColumns != null;
}
private Pair<byte[], RowSignature> materialize()
{
if (ordering != null) {
@ -180,7 +190,6 @@ public class LazilyDecoratedRowsAndColumns implements RowsAndColumns
} else {
return materializeStorageAdapter(as);
}
}
private void reset(RowsAndColumns rac)
@ -200,13 +209,26 @@ public class LazilyDecoratedRowsAndColumns implements RowsAndColumns
final Sequence<Cursor> cursors = as.makeCursors(
filter,
interval == null ? Intervals.ETERNITY : interval,
virtualColumns,
virtualColumns == null ? VirtualColumns.EMPTY : virtualColumns,
Granularities.ALL,
false,
null
);
Collection<String> cols = viewableColumns == null ? base.getColumnNames() : viewableColumns;
final Collection<String> cols;
if (viewableColumns != null) {
cols = viewableColumns;
} else {
if (virtualColumns == null) {
cols = base.getColumnNames();
} else {
cols = ImmutableList.<String>builder()
.addAll(base.getColumnNames())
.addAll(virtualColumns.getColumnNames())
.build();
}
}
AtomicReference<RowSignature> siggy = new AtomicReference<>(null);
FrameWriter writer = cursors.accumulate(null, (accumulated, in) -> {
@ -222,9 +244,18 @@ public class LazilyDecoratedRowsAndColumns implements RowsAndColumns
final RowSignature.Builder sigBob = RowSignature.builder();
for (String col : cols) {
final ColumnCapabilities capabilities = columnSelectorFactory.getColumnCapabilities(col);
ColumnCapabilities capabilities;
capabilities = columnSelectorFactory.getColumnCapabilities(col);
if (capabilities != null) {
sigBob.add(col, capabilities.toColumnType());
continue;
}
if (virtualColumns != null) {
capabilities = virtualColumns.getColumnCapabilities(columnSelectorFactory, col);
if (capabilities != null) {
sigBob.add(col, capabilities.toColumnType());
continue;
}
}
}
final RowSignature signature = sigBob.build();
@ -350,12 +381,12 @@ public class LazilyDecoratedRowsAndColumns implements RowsAndColumns
final RowSignature.Builder sigBob = RowSignature.builder();
final ArenaMemoryAllocatorFactory memFactory = new ArenaMemoryAllocatorFactory(200 << 20);
for (String column : columnsToGenerate) {
final Column racColumn = rac.findColumn(column);
if (racColumn == null) {
continue;
}
sigBob.add(column, racColumn.toAccessor().getType());
}

8
processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/FrameRowsAndColumns.java
View File

@ -21,10 +21,14 @@ package org.apache.druid.query.rowsandcols.concrete;
import org.apache.druid.frame.Frame;
import org.apache.druid.frame.FrameType;
import org.apache.druid.frame.read.FrameReader;
import org.apache.druid.frame.read.columnar.FrameColumnReaders;
import org.apache.druid.frame.segment.FrameStorageAdapter;
import org.apache.druid.java.util.common.ISE;
import org.apache.druid.java.util.common.Intervals;
import org.apache.druid.query.rowsandcols.RowsAndColumns;
import org.apache.druid.query.rowsandcols.column.Column;
import org.apache.druid.segment.StorageAdapter;
import org.apache.druid.segment.column.ColumnType;
import org.apache.druid.segment.column.RowSignature;
@ -77,10 +81,14 @@ public class FrameRowsAndColumns implements RowsAndColumns
}
@SuppressWarnings("unchecked")
@Nullable
@Override
public <T> T as(Class<T> clazz)
{
if (StorageAdapter.class.equals(clazz)) {
return (T) new FrameStorageAdapter(frame, FrameReader.create(signature), Intervals.ETERNITY);
}
return null;
}
}

14
processing/src/main/java/org/apache/druid/query/rowsandcols/semantic/DefaultColumnSelectorFactoryMaker.java
View File

@ -141,7 +141,7 @@ public class DefaultColumnSelectorFactoryMaker implements ColumnSelectorFactoryM
{
return withColumnAccessor(columnName, columnAccessor -> {
if (columnAccessor == null) {
return DimensionSelector.constant(null);
return DimensionSelector.nilSelector();
} else {
final ColumnType type = columnAccessor.getType();
switch (type.getType()) {
@ -160,16 +160,22 @@ public class DefaultColumnSelectorFactoryMaker implements ColumnSelectorFactoryM
@Override
public ColumnCapabilities getColumnCapabilities(String column)
{
return withColumnAccessor(column, columnAccessor ->
new ColumnCapabilitiesImpl()
return withColumnAccessor(column, columnAccessor -> {
if (columnAccessor == null) {
return null;
} else {
return new ColumnCapabilitiesImpl()
.setType(columnAccessor.getType())
.setHasMultipleValues(false)
.setDictionaryEncoded(false)
.setHasBitmapIndexes(false));
.setHasBitmapIndexes(false);
}
});
}
private <T> T withColumnAccessor(String column, Function<ColumnAccessor, T> fn)
{
@Nullable
ColumnAccessor retVal = accessorCache.get(column);
if (retVal == null) {
Column racColumn = rac.findColumn(column);

17
processing/src/main/java/org/apache/druid/segment/VirtualColumns.java
View File

@ -48,10 +48,12 @@ import org.apache.druid.segment.virtual.VirtualizedColumnSelectorFactory;
import javax.annotation.Nullable;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
/**
* Class allowing lookup and usage of virtual columns.
@ -112,6 +114,11 @@ public class VirtualColumns implements Cacheable
return new VirtualColumns(ImmutableList.copyOf(virtualColumns), withDotSupport, withoutDotSupport);
}
public static VirtualColumns create(VirtualColumn... virtualColumns)
{
return create(Arrays.asList(virtualColumns));
}
public static VirtualColumns nullToEmpty(@Nullable VirtualColumns virtualColumns)
{
return virtualColumns == null ? EMPTY : virtualColumns;
@ -519,4 +526,14 @@ public class VirtualColumns implements Cacheable
((VirtualColumns) obj).virtualColumns.isEmpty();
}
}
public boolean isEmpty()
{
return virtualColumns.isEmpty();
}
public List<String> getColumnNames()
{
return virtualColumns.stream().map(v -> v.getOutputName()).collect(Collectors.toList());
}
}

9
processing/src/test/java/org/apache/druid/error/DruidExceptionMatcher.java
View File

@ -45,6 +45,15 @@ public class DruidExceptionMatcher extends DiagnosingMatcher<Throwable>
return invalidInput().expectContext("sourceType", "sql");
}
public static DruidExceptionMatcher defensive()
{
return new DruidExceptionMatcher(
DruidException.Persona.DEVELOPER,
DruidException.Category.DEFENSIVE,
"general"
);
}
private final AllOf<DruidException> delegate;
private final ArrayList<Matcher<? super DruidException>> matcherList;

21
processing/src/test/java/org/apache/druid/query/operator/WindowOperatorQueryTest.java
View File

@ -19,18 +19,23 @@
package org.apache.druid.query.operator;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import org.apache.druid.java.util.common.Intervals;
import org.apache.druid.query.InlineDataSource;
import org.apache.druid.query.QueryContext;
import org.apache.druid.query.TableDataSource;
import org.apache.druid.query.spec.LegacySegmentSpec;
import org.apache.druid.query.spec.MultipleIntervalSegmentSpec;
import org.apache.druid.query.spec.QuerySegmentSpec;
import org.apache.druid.segment.column.RowSignature;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Set;
/**
@ -107,6 +112,22 @@ public class WindowOperatorQueryTest
Assert.assertSame(newDs, query.withDataSource(newDs).getDataSource());
}
@Test
public void withQuerySpec()
{
QuerySegmentSpec spec = new MultipleIntervalSegmentSpec(Collections.emptyList());
Assert.assertSame(spec, ((WindowOperatorQuery) query.withQuerySegmentSpec(spec)).getQuerySegmentSpec());
}
@Test
public void withOperators()
{
List<OperatorFactory> operators = ImmutableList.<OperatorFactory>builder()
.add(new NaivePartitioningOperatorFactory(Collections.singletonList("some")))
.build();
Assert.assertSame(operators, ((WindowOperatorQuery) query.withOperators(operators)).getOperators());
}
@Test
public void testEquals()
{

5
processing/src/test/java/org/apache/druid/query/rowsandcols/RowsAndColumnsTestBase.java
View File

@ -22,6 +22,8 @@ package org.apache.druid.query.rowsandcols;
import com.google.common.collect.Lists;
import org.apache.druid.common.config.NullHandling;
import org.apache.druid.java.util.common.ISE;
import org.apache.druid.query.rowsandcols.concrete.FrameRowsAndColumns;
import org.apache.druid.query.rowsandcols.concrete.FrameRowsAndColumnsTest;
import org.junit.Assert;
import org.junit.Test;
@ -63,7 +65,8 @@ public abstract class RowsAndColumnsTestBase
new Object[]{MapOfColumnsRowsAndColumns.class, Function.identity()},
new Object[]{ArrayListRowsAndColumns.class, ArrayListRowsAndColumnsTest.MAKER},
new Object[]{ConcatRowsAndColumns.class, ConcatRowsAndColumnsTest.MAKER},
new Object[]{RearrangedRowsAndColumns.class, RearrangedRowsAndColumnsTest.MAKER}
new Object[]{RearrangedRowsAndColumns.class, RearrangedRowsAndColumnsTest.MAKER},
new Object[]{FrameRowsAndColumns.class, FrameRowsAndColumnsTest.MAKER}
);
}

47
processing/src/test/java/org/apache/druid/query/rowsandcols/concrete/FrameRowsAndColumnsTest.java Normal file
View File

@ -0,0 +1,47 @@
/*
* 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.query.rowsandcols.concrete;
import org.apache.druid.query.rowsandcols.LazilyDecoratedRowsAndColumns;
import org.apache.druid.query.rowsandcols.MapOfColumnsRowsAndColumns;
import org.apache.druid.query.rowsandcols.RowsAndColumnsTestBase;
import java.util.function.Function;
public class FrameRowsAndColumnsTest extends RowsAndColumnsTestBase
{
public FrameRowsAndColumnsTest()
{
super(FrameRowsAndColumns.class);
}
public static Function<MapOfColumnsRowsAndColumns, FrameRowsAndColumns> MAKER = input -> {
return buildFrame(input);
};
private static FrameRowsAndColumns buildFrame(MapOfColumnsRowsAndColumns input)
{
LazilyDecoratedRowsAndColumns rac = new LazilyDecoratedRowsAndColumns(input, null, null, null, Integer.MAX_VALUE, null, null);
rac.numRows(); // materialize
return (FrameRowsAndColumns) rac.getBase();
}
}

30
processing/src/test/java/org/apache/druid/query/rowsandcols/semantic/CombinedSemanticInterfacesTest.java
View File

@ -30,6 +30,8 @@ import org.apache.druid.query.rowsandcols.RowsAndColumns;
import org.apache.druid.query.rowsandcols.column.ColumnAccessor;
import org.apache.druid.query.rowsandcols.column.IntArrayColumn;
import org.apache.druid.query.rowsandcols.column.ObjectArrayColumn;
import org.apache.druid.segment.ColumnSelectorFactory;
import org.apache.druid.segment.DimensionSelector;
import org.apache.druid.segment.column.ColumnType;
import org.junit.Assert;
import org.junit.Test;
@ -38,8 +40,12 @@ import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNull;
/**
* Place where tests can live that are testing the interactions of multiple semantic interfaces
*/
@ -54,6 +60,30 @@ public class CombinedSemanticInterfacesTest extends SemanticTestBase
super(name, fn);
}
@Test
public void testColumnSelectorFactoryMakeColumnValueSelectorNonExistentColumn()
{
RowsAndColumns rac = make(MapOfColumnsRowsAndColumns.fromMap(
ImmutableMap.of(
"some", new IntArrayColumn(new int[] {3, 54, 21, 1, 5, 54, 2, 3, 92}))));
AtomicInteger currRow = new AtomicInteger();
ColumnSelectorFactory csfm = ColumnSelectorFactoryMaker.fromRAC(rac).make(currRow);
assertEquals(DimensionSelector.nilSelector(), csfm.makeColumnValueSelector("nonexistent"));
}
@Test
public void testColumnSelectorFactoryGetColumnCapabilitiesNonExistentColumn()
{
RowsAndColumns rac = make(MapOfColumnsRowsAndColumns.fromMap(
ImmutableMap.of(
"some", new IntArrayColumn(new int[] {3, 54, 21, 1, 5, 54, 2, 3, 92}))));
AtomicInteger currRow = new AtomicInteger();
ColumnSelectorFactory csfm = ColumnSelectorFactoryMaker.fromRAC(rac).make(currRow);
assertNull(csfm.getColumnCapabilities("nonexistent"));
}
/**
* Tests a relatively common series of operations for window functions: partition -> aggregate -> sort
*/

92
processing/src/test/java/org/apache/druid/query/rowsandcols/semantic/TestVirtualColumnEvaluationRowsAndColumnsTest.java Normal file
View File

@ -0,0 +1,92 @@
/*
* 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.query.rowsandcols.semantic;
import com.google.common.collect.Lists;
import org.apache.druid.query.expression.TestExprMacroTable;
import org.apache.druid.query.operator.window.RowsAndColumnsHelper;
import org.apache.druid.query.rowsandcols.LazilyDecoratedRowsAndColumns;
import org.apache.druid.query.rowsandcols.MapOfColumnsRowsAndColumns;
import org.apache.druid.query.rowsandcols.RowsAndColumns;
import org.apache.druid.segment.StorageAdapter;
import org.apache.druid.segment.VirtualColumns;
import org.apache.druid.segment.column.ColumnType;
import org.apache.druid.segment.column.RowSignature;
import org.apache.druid.segment.virtual.ExpressionVirtualColumn;
import org.junit.Test;
import java.util.function.Function;
import static org.junit.Assert.assertEquals;
import static org.junit.Assume.assumeNotNull;
public class TestVirtualColumnEvaluationRowsAndColumnsTest extends SemanticTestBase
{
public TestVirtualColumnEvaluationRowsAndColumnsTest(String name, Function<MapOfColumnsRowsAndColumns, RowsAndColumns> fn)
{
super(name, fn);
}
@Test
public void testMaterializeVirtualColumns()
{
Object[][] vals = new Object[][] {
{1L, "a", 123L, 0L},
{2L, "a", 456L, 1L},
{3L, "b", 789L, 2L},
{4L, "b", 123L, 3L},
};
RowSignature siggy = RowSignature.builder()
.add("__time", ColumnType.LONG)
.add("dim", ColumnType.STRING)
.add("val", ColumnType.LONG)
.add("arrayIndex", ColumnType.LONG)
.build();
final RowsAndColumns base = make(MapOfColumnsRowsAndColumns.fromRowObjects(vals, siggy));
assumeNotNull("skipping: StorageAdapter not supported", base.as(StorageAdapter.class));
LazilyDecoratedRowsAndColumns ras = new LazilyDecoratedRowsAndColumns(
base,
null,
null,
VirtualColumns.create(new ExpressionVirtualColumn(
"expr",
"val * 2",
ColumnType.LONG,
TestExprMacroTable.INSTANCE)),
Integer.MAX_VALUE,
null,
null);
// do the materialziation
ras.numRows();
assertEquals(Lists.newArrayList("__time", "dim", "val", "arrayIndex", "expr"), ras.getColumnNames());
new RowsAndColumnsHelper()
.expectColumn("expr", new long[] {123 * 2, 456L * 2, 789 * 2, 123 * 2})
.validate(ras);
}
}

4
processing/src/test/java/org/apache/druid/segment/virtual/ExpressionVirtualColumnTest.java
View File

@ -716,8 +716,8 @@ public class ExpressionVirtualColumnTest extends InitializedNullHandlingTest
CURRENT_ROW.set(ROW0);
Assert.assertEquals(DateTimes.of("2000-01-01").getMillis(), selector.getLong());
Assert.assertEquals((float) DateTimes.of("2000-01-01").getMillis(), selector.getFloat(), 0.0f);
Assert.assertEquals((double) DateTimes.of("2000-01-01").getMillis(), selector.getDouble(), 0.0d);
Assert.assertEquals(DateTimes.of("2000-01-01").getMillis(), selector.getFloat(), 0.0f);
Assert.assertEquals(DateTimes.of("2000-01-01").getMillis(), selector.getDouble(), 0.0d);
Assert.assertEquals(DateTimes.of("2000-01-01").getMillis(), selector.getObject());
CURRENT_ROW.set(ROW1);

24
processing/src/test/java/org/apache/druid/segment/virtual/VirtualColumnsTest.java
View File

@ -64,8 +64,12 @@ import org.mockito.quality.Strictness;
import javax.annotation.Nullable;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
public class VirtualColumnsTest extends InitializedNullHandlingTest
{
private static final String REAL_COLUMN_NAME = "real_column";
@ -90,6 +94,26 @@ public class VirtualColumnsTest extends InitializedNullHandlingTest
Assert.assertFalse(virtualColumns.exists("bar"));
}
@Test
public void testIsEmpty()
{
assertTrue(VirtualColumns.EMPTY.isEmpty());
assertTrue(VirtualColumns.create(Collections.emptyList()).isEmpty());
}
@Test
public void testGetColumnNames()
{
final VirtualColumns virtualColumns = makeVirtualColumns();
List<String> colNames = ImmutableList.<String>builder()
.add("expr")
.add("expr2i")
.add("expr2")
.add("foo")
.build();
assertEquals(colNames, virtualColumns.getColumnNames());
}
@Test
public void testGetColumnCapabilitiesNilBase()
{

25
sql/src/main/java/org/apache/druid/sql/calcite/rel/DruidQuery.java
View File

@ -67,6 +67,8 @@ import org.apache.druid.query.groupby.GroupByQuery;
import org.apache.druid.query.groupby.having.DimFilterHavingSpec;
import org.apache.druid.query.groupby.orderby.DefaultLimitSpec;
import org.apache.druid.query.groupby.orderby.OrderByColumnSpec;
import org.apache.druid.query.operator.OperatorFactory;
import org.apache.druid.query.operator.ScanOperatorFactory;
import org.apache.druid.query.operator.WindowOperatorQuery;
import org.apache.druid.query.ordering.StringComparator;
import org.apache.druid.query.scan.ScanQuery;
@ -282,7 +284,8 @@ public class DruidQuery
partialQuery,
plannerContext,
sourceRowSignature, // Plans immediately after Scan, so safe to use the row signature from scan
rexBuilder
rexBuilder,
virtualColumnRegistry
)
);
} else {
@ -1442,12 +1445,30 @@ public class DruidQuery
return null;
}
// all virtual cols are needed - these columns are only referenced from the aggregates
VirtualColumns virtualColumns = virtualColumnRegistry.build(Collections.emptySet());
final List<OperatorFactory> operators;
if (virtualColumns.isEmpty()) {
operators = windowing.getOperators();
} else {
operators = ImmutableList.<OperatorFactory>builder()
.add(new ScanOperatorFactory(
null,
null,
null,
null,
virtualColumns,
null))
.addAll(windowing.getOperators())
.build();
}
return new WindowOperatorQuery(
dataSource,
new LegacySegmentSpec(Intervals.ETERNITY),
plannerContext.queryContextMap(),
windowing.getSignature(),
windowing.getOperators(),
operators,
null
);
}

27
sql/src/main/java/org/apache/druid/sql/calcite/rel/VirtualColumnRegistry.java
View File

@ -22,6 +22,7 @@ package org.apache.druid.sql.calcite.rel;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.sql.type.SqlTypeName;
import org.apache.druid.segment.VirtualColumn;
import org.apache.druid.segment.VirtualColumns;
import org.apache.druid.segment.column.ColumnCapabilities;
import org.apache.druid.segment.column.ColumnType;
import org.apache.druid.segment.column.RowSignature;
@ -32,13 +33,17 @@ import org.apache.druid.sql.calcite.planner.PlannerContext;
import javax.annotation.Nullable;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
import java.util.Optional;
import java.util.Queue;
import java.util.Set;
import java.util.stream.Collectors;
/**
@ -88,6 +93,11 @@ public class VirtualColumnRegistry
);
}
public boolean isEmpty()
{
return virtualColumnsByExpression.isEmpty();
}
/**
* Check if a {@link VirtualColumn} is defined by column name
*/
@ -350,4 +360,21 @@ public class VirtualColumnRegistry
return Objects.hash(expression, typeHint);
}
}
public VirtualColumns build(Set<String> exclude)
{
List<VirtualColumn> columns = new ArrayList<>();
if (virtualColumnsByName == null) {
return VirtualColumns.EMPTY;
}
for (Entry<String, ExpressionAndTypeHint> entry : virtualColumnsByName.entrySet()) {
if (exclude.contains(entry.getKey())) {
continue;
}
columns.add(getVirtualColumn(entry.getKey()));
}
columns.sort(Comparator.comparing(VirtualColumn::getOutputName));
return VirtualColumns.create(columns);
}
}

6
sql/src/main/java/org/apache/druid/sql/calcite/rel/Windowing.java
View File

@ -115,7 +115,8 @@ public class Windowing
final PartialDruidQuery partialQuery,
final PlannerContext plannerContext,
final RowSignature sourceRowSignature,
final RexBuilder rexBuilder
final RexBuilder rexBuilder,
final VirtualColumnRegistry virtualColumnRegistry
)
{
final Window window = Preconditions.checkNotNull(partialQuery.getWindow(), "window");
@ -172,10 +173,11 @@ public class Windowing
ProcessorMaker maker = KNOWN_WINDOW_FNS.get(aggregateCall.getAggregation().getName());
if (maker == null) {
final Aggregation aggregation = GroupByRules.translateAggregateCall(
plannerContext,
sourceRowSignature,
null,
virtualColumnRegistry,
rexBuilder,
InputAccessor.buildFor(
rexBuilder,

3
sql/src/test/java/org/apache/druid/sql/calcite/BaseCalciteQueryTest.java
View File

@ -1335,8 +1335,7 @@ public class BaseCalciteQueryTest extends CalciteTestBase
return queryJsonMapper.treeToValue(newQueryNode, Query.class);
}
catch (Exception e) {
Assert.fail(e.getMessage());
return null;
throw new RuntimeException(e);
}
}

68
sql/src/test/java/org/apache/druid/sql/calcite/DrillWindowQueryTest.java
View File

@ -4691,7 +4691,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("aggregates/winFnQry_12")
@Test
public void test_aggregates_winFnQry_12()
@ -4699,7 +4699,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("aggregates/winFnQry_13")
@Test
public void test_aggregates_winFnQry_13()
@ -4707,7 +4707,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("aggregates/winFnQry_20")
@Test
public void test_aggregates_winFnQry_20()
@ -4715,7 +4715,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("aggregates/winFnQry_21")
@Test
public void test_aggregates_winFnQry_21()
@ -4731,7 +4731,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/defaultFrame/RBUPACR_chr_1")
@Test
public void test_frameclause_defaultFrame_RBUPACR_chr_1()
@ -4739,7 +4739,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/defaultFrame/RBUPACR_chr_2")
@Test
public void test_frameclause_defaultFrame_RBUPACR_chr_2()
@ -4747,7 +4747,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/defaultFrame/RBUPACR_vchr_1")
@Test
public void test_frameclause_defaultFrame_RBUPACR_vchr_1()
@ -4755,7 +4755,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/defaultFrame/RBUPACR_vchr_2")
@Test
public void test_frameclause_defaultFrame_RBUPACR_vchr_2()
@ -4763,7 +4763,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/multipl_wnwds/max_mulwds")
@Test
public void test_frameclause_multipl_wnwds_max_mulwds()
@ -4771,7 +4771,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/multipl_wnwds/min_mulwds")
@Test
public void test_frameclause_multipl_wnwds_min_mulwds()
@ -4779,7 +4779,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/RBCRACR/RBCRACR_char_1")
@Test
public void test_frameclause_RBCRACR_RBCRACR_char_1()
@ -4787,7 +4787,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/RBCRACR/RBCRACR_char_2")
@Test
public void test_frameclause_RBCRACR_RBCRACR_char_2()
@ -4795,7 +4795,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/RBCRACR/RBCRACR_vchar_1")
@Test
public void test_frameclause_RBCRACR_RBCRACR_vchar_1()
@ -4803,7 +4803,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/RBCRACR/RBCRACR_vchar_2")
@Test
public void test_frameclause_RBCRACR_RBCRACR_vchar_2()
@ -4811,7 +4811,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/RBUPACR/RBUPACR_chr_1")
@Test
public void test_frameclause_RBUPACR_RBUPACR_chr_1()
@ -4819,7 +4819,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/RBUPACR/RBUPACR_chr_2")
@Test
public void test_frameclause_RBUPACR_RBUPACR_chr_2()
@ -4827,7 +4827,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/RBUPACR/RBUPACR_vchr_1")
@Test
public void test_frameclause_RBUPACR_RBUPACR_vchr_1()
@ -4835,7 +4835,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/RBUPACR/RBUPACR_vchr_2")
@Test
public void test_frameclause_RBUPACR_RBUPACR_vchr_2()
@ -4843,7 +4843,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/RBUPAUF/RBUPAUF_char_1")
@Test
public void test_frameclause_RBUPAUF_RBUPAUF_char_1()
@ -4851,7 +4851,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/RBUPAUF/RBUPAUF_char_2")
@Test
public void test_frameclause_RBUPAUF_RBUPAUF_char_2()
@ -4859,7 +4859,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/RBUPAUF/RBUPAUF_vchar_1")
@Test
public void test_frameclause_RBUPAUF_RBUPAUF_vchar_1()
@ -4867,7 +4867,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/RBUPAUF/RBUPAUF_vchar_2")
@Test
public void test_frameclause_RBUPAUF_RBUPAUF_vchar_2()
@ -4875,7 +4875,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/subQueries/frmInSubQry_22")
@Test
public void test_frameclause_subQueries_frmInSubQry_22()
@ -4883,7 +4883,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/subQueries/frmInSubQry_23")
@Test
public void test_frameclause_subQueries_frmInSubQry_23()
@ -4891,7 +4891,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/subQueries/frmInSubQry_24")
@Test
public void test_frameclause_subQueries_frmInSubQry_24()
@ -4899,7 +4899,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/subQueries/frmInSubQry_41")
@Test
public void test_frameclause_subQueries_frmInSubQry_41()
@ -4907,7 +4907,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/subQueries/frmInSubQry_42")
@Test
public void test_frameclause_subQueries_frmInSubQry_42()
@ -4915,7 +4915,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/subQueries/frmInSubQry_43")
@Test
public void test_frameclause_subQueries_frmInSubQry_43()
@ -4923,7 +4923,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/subQueries/frmInSubQry_44")
@Test
public void test_frameclause_subQueries_frmInSubQry_44()
@ -4931,7 +4931,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/subQueries/frmInSubQry_45")
@Test
public void test_frameclause_subQueries_frmInSubQry_45()
@ -4939,7 +4939,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("frameclause/subQueries/frmInSubQry_46")
@Test
public void test_frameclause_subQueries_frmInSubQry_46()
@ -4963,7 +4963,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("nestedAggs/basic_10")
@Test
public void test_nestedAggs_basic_10()
@ -4971,7 +4971,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE)
@NotYetSupported(Modes.AGGREGATION_NOT_SUPPORT_TYPE)
@DrillTest("nestedAggs/cte_win_01")
@Test
public void test_nestedAggs_cte_win_01()
@ -4979,7 +4979,7 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
windowQueryTest();
}
@NotYetSupported(Modes.NPE_PLAIN)
@NotYetSupported(Modes.RESULT_MISMATCH)
@DrillTest("aggregates/winFnQry_7")
@Test
public void test_aggregates_winFnQry_7()

2
sql/src/test/java/org/apache/druid/sql/calcite/NotYetSupported.java
View File

@ -79,7 +79,7 @@ public @interface NotYetSupported
BIGINT_TO_DATE(DruidException.class, "BIGINT to type (DATE|TIME)"),
NPE_PLAIN(NullPointerException.class, "java.lang.NullPointerException"),
NPE(DruidException.class, "java.lang.NullPointerException"),
AGGREGATION_NOT_SUPPORT_TYPE(DruidException.class, "Aggregation \\[(MIN|MAX)\\] does not support type"),
AGGREGATION_NOT_SUPPORT_TYPE(DruidException.class, "Aggregation \\[(MIN|MAX)\\] does not support type \\[STRING\\]"),
CANNOT_APPLY_VIRTUAL_COL(UOE.class, "apply virtual columns"),
MISSING_DESC(DruidException.class, "function signature DESC"),
RESULT_COUNT_MISMATCH(AssertionError.class, "result count:"),

33
sql/src/test/resources/calcite/tests/window/virtualColumns.sqlTest Normal file
View File

@ -0,0 +1,33 @@
type: "operatorValidation"
sql: |
SELECT
'_'||dim1,
MIN(cast(42 as double)) OVER (),
MIN(cast((cnt||cnt) as bigint)) OVER ()
FROM foo
expectedOperators:
- type: "scan"
limit: -1
virtualColumns:
- type: "expression"
name: "_v0"
expression: "42.0"
outputType: "DOUBLE"
- { type: "naivePartition", partitionColumns: [ ] }
- type: "window"
processor:
type: "framedAgg"
frame: { peerType: "ROWS", lowUnbounded: true, lowOffset: 0, uppUnbounded: true, uppOffset: 0 }
aggregations:
- { type: "doubleMin", name: "w0", fieldName: "_v0" }
- { type: "longMin", name: "w1", fieldName: "v1" }
expectedResults:
- ["_",42.0,11]
- ["_10.1",42.0,11]
- ["_2",42.0,11]
- ["_1",42.0,11]
- ["_def",42.0,11]
- ["_abc",42.0,11]