fix bug when rewriting sql virtual column registry (#12718)

This commit is contained in:
Clint Wylie 2022-07-01 02:24:00 -07:00 committed by GitHub
parent c09b5a2294
commit d30efb1c1e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 67 additions and 2 deletions

View File

@ -29,12 +29,14 @@ import org.apache.druid.sql.calcite.planner.Calcites;
import org.apache.druid.sql.calcite.planner.PlannerContext;
import javax.annotation.Nullable;
import java.util.ArrayDeque;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Queue;
import java.util.stream.Collectors;
/**
@ -198,13 +200,15 @@ public class VirtualColumnRegistry
public void visitAllSubExpressions(DruidExpression.DruidExpressionShuttle shuttle)
{
for (Map.Entry<String, ExpressionAndTypeHint> entry : virtualColumnsByName.entrySet()) {
final Queue<Map.Entry<String, ExpressionAndTypeHint>> toVisit = new ArrayDeque<>(virtualColumnsByName.entrySet());
while (!toVisit.isEmpty()) {
final Map.Entry<String, ExpressionAndTypeHint> entry = toVisit.poll();
final String key = entry.getKey();
final ExpressionAndTypeHint wrapped = entry.getValue();
virtualColumnsByExpression.remove(wrapped);
final List<DruidExpression> newArgs = shuttle.visitAll(wrapped.getExpression().getArguments());
final ExpressionAndTypeHint newWrapped = wrap(wrapped.getExpression().withArguments(newArgs), wrapped.getTypeHint());
virtualColumnsByName.put(key, newWrapped);
virtualColumnsByExpression.remove(wrapped);
virtualColumnsByExpression.put(newWrapped, key);
}
}

View File

@ -1428,6 +1428,67 @@ public class CalciteMultiValueStringQueryTest extends BaseCalciteQueryTest
);
}
@Test
public void testMultiValueListFilterComposedMultipleExpressions() throws Exception
{
// Cannot vectorize due to usage of expressions.
cannotVectorize();
testQuery(
"SELECT MV_LENGTH(MV_FILTER_ONLY(dim3, ARRAY['b'])), MV_LENGTH(dim3), SUM(cnt) FROM druid.numfoo GROUP BY 1,2 ORDER BY 3 DESC",
ImmutableList.of(
GroupByQuery.builder()
.setDataSource(CalciteTests.DATASOURCE3)
.setInterval(querySegmentSpec(Filtration.eternity()))
.setGranularity(Granularities.ALL)
.setVirtualColumns(
expressionVirtualColumn(
"v0",
"array_length(\"v2\")",
ColumnType.LONG
),
expressionVirtualColumn(
"v1",
"array_length(\"dim3\")",
ColumnType.LONG
),
new ListFilteredVirtualColumn(
"v2",
DefaultDimensionSpec.of("dim3"),
ImmutableSet.of("b"),
true
)
)
.setDimensions(
dimensions(
new DefaultDimensionSpec("v0", "_d0", ColumnType.LONG),
new DefaultDimensionSpec("v1", "_d1", ColumnType.LONG)
)
)
.setAggregatorSpecs(aggregators(new LongSumAggregatorFactory("a0", "cnt")))
.setLimitSpec(new DefaultLimitSpec(
ImmutableList.of(new OrderByColumnSpec(
"a0",
OrderByColumnSpec.Direction.DESCENDING,
StringComparators.NUMERIC
)),
Integer.MAX_VALUE
))
.setContext(QUERY_CONTEXT_DEFAULT)
.build()
),
useDefault ? ImmutableList.of(
new Object[]{0, 0, 3L},
new Object[]{1, 2, 2L},
new Object[]{0, 1, 1L}
) : ImmutableList.of(
new Object[]{null, null, 2L},
new Object[]{null, 1, 2L},
new Object[]{1, 2, 2L}
)
);
}
@Test
public void testFilterOnMultiValueListFilterNoMatch() throws Exception
{