Skip to content

Commit

Permalink
fix expression column capabilities to not report dictionary encoded u…
Browse files Browse the repository at this point in the history 
…nless input is string (#16577)
  • Loading branch information
@clintropolis
clintropolis authored Jun 8, 2024
1 parent 40ba429 commit 3fb6ba2
Show file tree
Hide file tree
Showing 5 changed files with 157 additions and 23 deletions.
4 changes: 0 additions & 4 deletions ...ssing/src/main/java/org/apache/druid/segment/serde/NestedCommonFormatColumnPartSerde.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -303,10 +303,6 @@ public void read(ByteBuffer buffer, ColumnBuilder builder, ColumnConfig columnCo
bitmapSerdeFactory,
byteOrder
);
ColumnCapabilitiesImpl capabilitiesBuilder = builder.getCapabilitiesBuilder();
capabilitiesBuilder.setDictionaryEncoded(true);
capabilitiesBuilder.setDictionaryValuesSorted(true);
capabilitiesBuilder.setDictionaryValuesUnique(true);
ColumnType simpleType = supplier.getLogicalType();
ColumnType logicalType = simpleType == null ? ColumnType.NESTED_DATA : simpleType;
builder.setType(logicalType);
Expand Down
16 changes: 11 additions & 5 deletions processing/src/main/java/org/apache/druid/segment/virtual/ExpressionPlan.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -274,12 +274,18 @@ public ColumnCapabilities inferColumnCapabilities(@Nullable ColumnType outputTyp
// since we don't know if the expression is 1:1 or if it retains ordering we can only piggy back only
// report as dictionary encoded, but it still allows us to use algorithms which work with dictionaryIds
// to create a dictionary encoded selector instead of an object selector to defer expression evaluation
// until query time
return ColumnCapabilitiesImpl.copyOf(underlyingCapabilities)
// until query time. However, currently dictionary encodedness is tied to string selectors and sad stuff
// happens if the input type isn't string, so we also limit this to string input types
final boolean useDictionary = underlyingCapabilities.isDictionaryEncoded().isTrue() &&
underlyingCapabilities.is(ValueType.STRING);
return ColumnCapabilitiesImpl.createDefault()
.setType(ColumnType.STRING)
.setDictionaryValuesSorted(false)
.setDictionaryValuesUnique(false)
.setHasBitmapIndexes(false)
.setDictionaryEncoded(useDictionary)
.setHasBitmapIndexes(underlyingCapabilities.hasBitmapIndexes())
.setHasMultipleValues(underlyingCapabilities.hasMultipleValues())
.setHasSpatialIndexes(underlyingCapabilities.hasSpatialIndexes())
// we aren't sure if the expression might produce null values or not, so always
// set hasNulls to true
.setHasNulls(true);
}
}
Expand Down
126 changes: 113 additions & 13 deletions processing/src/test/java/org/apache/druid/segment/virtual/ExpressionPlannerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,13 @@

package org.apache.druid.segment.virtual;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import org.apache.druid.common.config.NullHandling;
import org.apache.druid.error.DruidException;
import org.apache.druid.math.expr.Expr;
import org.apache.druid.math.expr.ExprEval;
import org.apache.druid.math.expr.ExprMacroTable;
import org.apache.druid.math.expr.ExpressionType;
import org.apache.druid.math.expr.Parser;
import org.apache.druid.query.expression.TestExprMacroTable;
Expand All @@ -36,11 +41,13 @@
import org.junit.rules.ExpectedException;

import javax.annotation.Nullable;
import java.util.List;
import java.util.Map;

public class ExpressionPlannerTest extends InitializedNullHandlingTest
{
public static final ColumnInspector SYNTHETIC_INSPECTOR = new ColumnInspector()
private static ColumnType DICTIONARY_COMPLEX = ColumnType.ofComplex("dictionaryComplex");
private static final ColumnInspector SYNTHETIC_INSPECTOR = new ColumnInspector()
{
private final Map<String, ColumnCapabilities> capabilitiesMap =
ImmutableMap.<String, ColumnCapabilities>builder()
Expand Down Expand Up @@ -141,6 +148,12 @@ public class ExpressionPlannerTest extends InitializedNullHandlingTest
"double_array_2",
ColumnCapabilitiesImpl.createSimpleArrayColumnCapabilities(ColumnType.DOUBLE_ARRAY)
)
.put(
"dictionary_complex",
ColumnCapabilitiesImpl.createDefault()
.setDictionaryEncoded(true)
.setType(DICTIONARY_COMPLEX)
)
.build();

@Nullable
Expand All @@ -151,6 +164,8 @@ public ColumnCapabilities getColumnCapabilities(String column)
}
};

private static final TestMacroTable MACRO_TABLE = new TestMacroTable();

@Rule
public ExpectedException expectedException = ExpectedException.none();

Expand Down Expand Up @@ -369,7 +384,7 @@ public void testScalarStringDictionaryEncoded()
Assert.assertFalse(inferred.areDictionaryValuesSorted().isMaybeTrue());
Assert.assertFalse(inferred.areDictionaryValuesUnique().isMaybeTrue());
Assert.assertFalse(inferred.hasMultipleValues().isMaybeTrue());
Assert.assertFalse(inferred.hasBitmapIndexes());
Assert.assertTrue(inferred.hasBitmapIndexes());
Assert.assertFalse(inferred.hasSpatialIndexes());

// multiple input columns
Expand Down Expand Up @@ -463,7 +478,7 @@ public void testMultiValueStringDictionaryEncoded()
Assert.assertFalse(inferred.areDictionaryValuesSorted().isMaybeTrue());
Assert.assertFalse(inferred.areDictionaryValuesUnique().isMaybeTrue());
Assert.assertTrue(inferred.hasMultipleValues().isTrue());
Assert.assertFalse(inferred.hasBitmapIndexes());
Assert.assertTrue(inferred.hasBitmapIndexes());
Assert.assertFalse(inferred.hasSpatialIndexes());

thePlan = plan("concat(scalar_string, multi_dictionary_string_nonunique)");
Expand Down Expand Up @@ -599,7 +614,7 @@ public void testIncompleteString()
)
);
Assert.assertFalse(
thePlan.is(
thePlan.any(
ExpressionPlan.Trait.SINGLE_INPUT_SCALAR,
ExpressionPlan.Trait.SINGLE_INPUT_MAPPABLE,
ExpressionPlan.Trait.UNKNOWN_INPUTS,
Expand Down Expand Up @@ -671,7 +686,7 @@ public void testArrayOutput()
)
);
Assert.assertFalse(
thePlan.is(
thePlan.any(
ExpressionPlan.Trait.SINGLE_INPUT_SCALAR,
ExpressionPlan.Trait.SINGLE_INPUT_MAPPABLE,
ExpressionPlan.Trait.UNKNOWN_INPUTS,
Expand Down Expand Up @@ -719,7 +734,22 @@ public void testScalarOutputMultiValueInput()

// what about a multi-valued input
thePlan = plan("array_to_string(array_append(scalar_string, multi_dictionary_string), ',')");
assertArrayInput(thePlan);
Assert.assertTrue(
thePlan.is(
ExpressionPlan.Trait.NON_SCALAR_INPUTS,
ExpressionPlan.Trait.NEEDS_APPLIED
)
);
Assert.assertFalse(
thePlan.any(
ExpressionPlan.Trait.SINGLE_INPUT_SCALAR,
ExpressionPlan.Trait.SINGLE_INPUT_MAPPABLE,
ExpressionPlan.Trait.NON_SCALAR_OUTPUT,
ExpressionPlan.Trait.INCOMPLETE_INPUTS,
ExpressionPlan.Trait.UNKNOWN_INPUTS,
ExpressionPlan.Trait.VECTORIZABLE
)
);

Assert.assertEquals(
"array_to_string(map((\"multi_dictionary_string\") -> array_append(\"scalar_string\", \"multi_dictionary_string\"), \"multi_dictionary_string\"), ',')",
Expand Down Expand Up @@ -789,7 +819,7 @@ public void testArrayConstruction()
)
);
Assert.assertFalse(
thePlan.is(
thePlan.any(
ExpressionPlan.Trait.SINGLE_INPUT_SCALAR,
ExpressionPlan.Trait.SINGLE_INPUT_MAPPABLE,
ExpressionPlan.Trait.UNKNOWN_INPUTS,
Expand All @@ -812,15 +842,14 @@ public void testNestedColumnExpression()
{
ExpressionPlan thePlan = plan("json_object('long1', long1, 'long2', long2)");
Assert.assertFalse(
thePlan.is(
thePlan.any(
ExpressionPlan.Trait.NON_SCALAR_OUTPUT,
ExpressionPlan.Trait.SINGLE_INPUT_SCALAR,
ExpressionPlan.Trait.SINGLE_INPUT_MAPPABLE,
ExpressionPlan.Trait.UNKNOWN_INPUTS,
ExpressionPlan.Trait.INCOMPLETE_INPUTS,
ExpressionPlan.Trait.NEEDS_APPLIED,
ExpressionPlan.Trait.NON_SCALAR_INPUTS,
ExpressionPlan.Trait.VECTORIZABLE
ExpressionPlan.Trait.NON_SCALAR_INPUTS
)
);
Assert.assertEquals(ExpressionType.NESTED_DATA, thePlan.getOutputType());
Expand All @@ -837,9 +866,40 @@ public void testNestedColumnExpression()
);
}

@Test
public void testDictionaryComplexStringOutput()
{
ExpressionPlan thePlan = plan("dict_complex_to_string(dictionary_complex)");
Assert.assertFalse(
thePlan.any(
ExpressionPlan.Trait.NON_SCALAR_OUTPUT,
ExpressionPlan.Trait.SINGLE_INPUT_MAPPABLE,
ExpressionPlan.Trait.UNKNOWN_INPUTS,
ExpressionPlan.Trait.INCOMPLETE_INPUTS,
ExpressionPlan.Trait.NEEDS_APPLIED,
ExpressionPlan.Trait.NON_SCALAR_INPUTS
)
);
Assert.assertTrue(
thePlan.is(
ExpressionPlan.Trait.SINGLE_INPUT_SCALAR,
ExpressionPlan.Trait.VECTORIZABLE
)
);
Assert.assertEquals(ExpressionType.STRING, thePlan.getOutputType());
ColumnCapabilities inferred = thePlan.inferColumnCapabilities(
ExpressionType.toColumnType(thePlan.getOutputType())
);
Assert.assertEquals(
ColumnType.STRING.getType(),
inferred.getType()
);
Assert.assertFalse(inferred.isDictionaryEncoded().isMaybeTrue());
}

private static ExpressionPlan plan(String expression)
{
return ExpressionPlanner.plan(SYNTHETIC_INSPECTOR, Parser.parse(expression, TestExprMacroTable.INSTANCE));
return ExpressionPlanner.plan(SYNTHETIC_INSPECTOR, Parser.parse(expression, MACRO_TABLE));
}

private static void assertArrayInput(ExpressionPlan thePlan)
Expand All @@ -850,7 +910,7 @@ private static void assertArrayInput(ExpressionPlan thePlan)
)
);
Assert.assertFalse(
thePlan.is(
thePlan.any(
ExpressionPlan.Trait.SINGLE_INPUT_SCALAR,
ExpressionPlan.Trait.SINGLE_INPUT_MAPPABLE,
ExpressionPlan.Trait.NON_SCALAR_OUTPUT,
Expand All @@ -871,7 +931,7 @@ private static void assertArrayInAndOut(ExpressionPlan thePlan)
)
);
Assert.assertFalse(
thePlan.is(
thePlan.any(
ExpressionPlan.Trait.SINGLE_INPUT_SCALAR,
ExpressionPlan.Trait.SINGLE_INPUT_MAPPABLE,
ExpressionPlan.Trait.INCOMPLETE_INPUTS,
Expand All @@ -881,4 +941,44 @@ private static void assertArrayInAndOut(ExpressionPlan thePlan)
)
);
}

private static class TestMacroTable extends ExprMacroTable
{
public TestMacroTable()
{
super(
ImmutableList.<ExprMacroTable.ExprMacro>builder()
.addAll(TestExprMacroTable.INSTANCE.getMacros())
.add(new ExprMacroTable.ExprMacro()
{
@Override
public Expr apply(List<Expr> args)
{
return new ExprMacroTable.BaseScalarMacroFunctionExpr(this, args)
{
@Override
public ExprEval eval(ObjectBinding bindings)
{
throw DruidException.defensive("just for planner test");
}

@Nullable
@Override
public ExpressionType getOutputType(InputBindingInspector inspector)
{
return ExpressionType.STRING;
}
};
}

@Override
public String name()
{
return "dict_complex_to_string";
}
})
.build()
);
}
}
}
2 changes: 1 addition & 1 deletion processing/src/test/java/org/apache/druid/segment/virtual/ExpressionVirtualColumnTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -115,7 +115,7 @@ public class ExpressionVirtualColumnTest extends InitializedNullHandlingTest
ImmutableMap.of(
"x", 3L,
"y", 4L,
"b", Arrays.asList(new String[]{"3", null, "5"})
"b", Arrays.asList("3", null, "5")
)
);

Expand Down
32 changes: 32 additions & 0 deletions sql/src/test/java/org/apache/druid/sql/calcite/CalciteNestedDataQueryTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7558,4 +7558,36 @@ public void testGroupByAutoDouble()
.build()
);
}

@Test
public void testToJsonString()
{
testQuery(
"SELECT TO_JSON_STRING(nester) FROM druid.nested GROUP BY 1",
ImmutableList.of(
GroupByQuery.builder()
.setDataSource(DATA_SOURCE)
.setInterval(querySegmentSpec(Filtration.eternity()))
.setGranularity(Granularities.ALL)
.setDimensions(
dimensions(
new DefaultDimensionSpec("v0", "d0", ColumnType.STRING)
)
)
.setVirtualColumns(
expressionVirtualColumn("v0", "to_json_string(\"nester\")", ColumnType.STRING)
)
.setContext(QUERY_CONTEXT_DEFAULT)
.build()
),
ImmutableList.of(
new Object[]{NullHandling.defaultStringValue()},
new Object[]{"\"hello\""},
new Object[]{"2"},
new Object[]{"{\"array\":[\"a\",\"b\"],\"n\":{\"x\":\"hello\"}}"},
new Object[]{"{\"array\":[\"a\",\"b\"],\"n\":{\"x\":1}}"}
),
RowSignature.builder().add("EXPR$0", ColumnType.STRING).build()
);
}
}

0 comments on commit 3fb6ba2

Please sign in to comment.