Skip to content

Commit 52e4501

Browse files
authored
ShimCursor: Fix dimension selectors on non-strings. (apache#18371)
Follow-up to apache#18305. Fixes a bug where calling ShimCursor's makeDimensionSelector on a non-string type would throw an error. Now, it has behavior similar to QueryableIndexCursor: - Uses ValueTypes.makeNumericWrappingDimensionSelector on numeric types, casting the numbers to strings. - Returns a nil selector for array and complex types. The new tests verify selector compatibility more extensively.
1 parent 69bbb1e commit 52e4501

2 files changed

Lines changed: 423 additions & 159 deletions

File tree

processing/src/main/java/org/apache/druid/segment/shim/ShimColumnSelectorFactory.java

Lines changed: 37 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import org.apache.druid.segment.NilColumnValueSelector;
3030
import org.apache.druid.segment.column.ColumnCapabilities;
3131
import org.apache.druid.segment.column.ValueType;
32+
import org.apache.druid.segment.column.ValueTypes;
3233
import org.apache.druid.segment.vector.MultiValueDimensionVectorSelector;
3334
import org.apache.druid.segment.vector.SingleValueDimensionVectorSelector;
3435
import org.apache.druid.segment.vector.VectorObjectSelector;
@@ -57,28 +58,49 @@ public DimensionSelector makeDimensionSelector(DimensionSpec dimensionSpec)
5758
return dimensionSelectors.computeIfAbsent(
5859
dimensionSpec,
5960
spec -> {
60-
if (spec.mustDecorate()) {
61+
if (spec.mustDecorate() || spec.getExtractionFn() != null) {
6162
throw DruidException.defensive("Only non-decorated dimensions can be vectorized.");
6263
}
63-
final ColumnCapabilities capabilities = cursor.vectorColumnSelectorFactory
64-
.getColumnCapabilities(dimensionSpec.getDimension());
64+
65+
final String columnName = dimensionSpec.getDimension();
66+
final ColumnCapabilities capabilities = cursor.vectorColumnSelectorFactory.getColumnCapabilities(columnName);
67+
6568
if (capabilities == null) {
6669
return DimensionSelector.nilSelector();
67-
} else if (ColumnProcessors.useDictionaryEncodedSelector(capabilities)) {
68-
if (capabilities.hasMultipleValues().isMaybeTrue()) {
69-
final MultiValueDimensionVectorSelector vectorSelector =
70-
cursor.vectorColumnSelectorFactory.makeMultiValueDimensionSelector(spec);
71-
return new ShimMultiValueDimensionSelector(cursor, vectorSelector);
70+
} else if (capabilities.is(ValueType.STRING)) {
71+
if (ColumnProcessors.useDictionaryEncodedSelector(capabilities)) {
72+
// Dictionary-encoded string column.
73+
if (capabilities.hasMultipleValues().isMaybeTrue()) {
74+
final MultiValueDimensionVectorSelector vectorSelector =
75+
cursor.vectorColumnSelectorFactory.makeMultiValueDimensionSelector(spec);
76+
return new ShimMultiValueDimensionSelector(cursor, vectorSelector);
77+
} else {
78+
final SingleValueDimensionVectorSelector vectorSelector =
79+
cursor.vectorColumnSelectorFactory.makeSingleValueDimensionSelector(spec);
80+
return new ShimSingleValueDimensionSelector(cursor, vectorSelector);
81+
}
7282
} else {
73-
final SingleValueDimensionVectorSelector vectorSelector =
74-
cursor.vectorColumnSelectorFactory.makeSingleValueDimensionSelector(spec);
75-
return new ShimSingleValueDimensionSelector(cursor, vectorSelector);
83+
// Non-dictionary encoded string column. Possibly an expression virtual column.
84+
final VectorObjectSelector vectorObjectSelector =
85+
cursor.vectorColumnSelectorFactory.makeObjectSelector(spec.getDimension());
86+
return new ShimVectorObjectDimSelector(
87+
cursor,
88+
vectorObjectSelector,
89+
capabilities.hasMultipleValues().isMaybeTrue()
90+
);
7691
}
92+
} else if (capabilities.isNumeric()) {
93+
// Caller requested a dimension selector on top of a numeric column.
94+
return ValueTypes.makeNumericWrappingDimensionSelector(
95+
capabilities.getType(),
96+
makeColumnValueSelector(columnName),
97+
null /* No extractionFn; we checked above that extractionFn is not present. */
98+
);
7799
} else {
78-
// Non-dictionary encoded column, like virtual columns.
79-
VectorObjectSelector vectorObjectSelector =
80-
cursor.vectorColumnSelectorFactory.makeObjectSelector(spec.getDimension());
81-
return new ShimVectorObjectDimSelector(cursor, vectorObjectSelector, capabilities.hasMultipleValues().isMaybeTrue());
100+
// Array or complex. Callers should be calling makeColumnValueSelector instead for these types. If they do
101+
// call makeDimensionSelector for some reason, give them a column full of nulls, since that's what
102+
// QueryableIndexColumnSelectorFactory would do.
103+
return DimensionSelector.nilSelector();
82104
}
83105
}
84106
);

0 commit comments

Comments
 (0)