Skip to content

fix ColumnType to RelDataType conversion for nested arrays #16138

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gianm merged 2 commits into from
Mar 19, 2024
Merged

fix ColumnType to RelDataType conversion for nested arrays #16138

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
339a213
fix ColumnType to RelDataType conversion for nested arrays
clintropolis Mar 15, 2024
5866b66
fix test
clintropolis Mar 15, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 48 additions & 42 deletions sql/src/main/java/org/apache/druid/sql/calcite/table/RowSignatures.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,17 +96,17 @@ public static StringComparator getNaturalStringComparator(
}

/**
* Returns a Calcite RelDataType corresponding to a row signature. It will typecast __time column to TIMESTAMP
* irrespective of the type present in the row signature
* Returns a Calcite {@link RelDataType} corresponding to a {@link RowSignature}. It will typecast __time column to
* TIMESTAMP irrespective of the type present in the row signature
*/
public static RelDataType toRelDataType(final RowSignature rowSignature, final RelDataTypeFactory typeFactory)
{
return toRelDataType(rowSignature, typeFactory, true);
}

/**
* Returns a Calcite RelDataType corresponding to a row signature.
* For columns that are named "__time", it automatically casts it to TIMESTAMP if typecastTimeColumn is set to true
* Returns a Calcite {@link RelDataType} corresponding to a {@link RowSignature}. For columns that are named
* "__time", it automatically casts it to TIMESTAMP if typecastTimeColumn is set to true
*/
public static RelDataType toRelDataType(
final RowSignature rowSignature,
Expand All @@ -126,44 +126,7 @@ public static RelDataType toRelDataType(
rowSignature.getColumnType(columnName)
.orElseThrow(() -> new ISE("Encountered null type for column[%s]", columnName));

switch (columnType.getType()) {
case STRING:
// Note that there is no attempt here to handle multi-value in any special way. Maybe one day...
type = Calcites.createSqlTypeWithNullability(typeFactory, SqlTypeName.VARCHAR, true);
break;
case LONG:
type = Calcites.createSqlTypeWithNullability(typeFactory, SqlTypeName.BIGINT, nullNumeric);
break;
case FLOAT:
type = Calcites.createSqlTypeWithNullability(typeFactory, SqlTypeName.FLOAT, nullNumeric);
break;
case DOUBLE:
type = Calcites.createSqlTypeWithNullability(typeFactory, SqlTypeName.DOUBLE, nullNumeric);
break;
case ARRAY:
switch (columnType.getElementType().getType()) {
case STRING:
type = Calcites.createSqlArrayTypeWithNullability(typeFactory, SqlTypeName.VARCHAR, true);
break;
case LONG:
type = Calcites.createSqlArrayTypeWithNullability(typeFactory, SqlTypeName.BIGINT, nullNumeric);
break;
case DOUBLE:
type = Calcites.createSqlArrayTypeWithNullability(typeFactory, SqlTypeName.DOUBLE, nullNumeric);
break;
case FLOAT:
type = Calcites.createSqlArrayTypeWithNullability(typeFactory, SqlTypeName.FLOAT, nullNumeric);
break;
default:
throw new ISE("valueType[%s] not translatable", columnType);
}
break;
case COMPLEX:
type = makeComplexType(typeFactory, columnType, true);
break;
default:
throw new ISE("valueType[%s] not translatable", columnType);
}
type = columnTypeToRelDataType(typeFactory, columnType, nullNumeric);
}

builder.add(columnName, type);
Expand All @@ -172,6 +135,49 @@ public static RelDataType toRelDataType(
return builder.build();
}

/**
* Returns a Calcite {@link RelDataType} corresponding to a {@link ColumnType}
*/
public static RelDataType columnTypeToRelDataType(
RelDataTypeFactory typeFactory,
ColumnType columnType,
boolean nullNumeric
)
{
final RelDataType type;
switch (columnType.getType()) {
case STRING:
// Note that there is no attempt here to handle multi-value in any special way. Maybe one day...
type = Calcites.createSqlTypeWithNullability(typeFactory, SqlTypeName.VARCHAR, true);
break;
case LONG:
type = Calcites.createSqlTypeWithNullability(typeFactory, SqlTypeName.BIGINT, nullNumeric);
break;
case FLOAT:
type = Calcites.createSqlTypeWithNullability(typeFactory, SqlTypeName.FLOAT, nullNumeric);
break;
case DOUBLE:
type = Calcites.createSqlTypeWithNullability(typeFactory, SqlTypeName.DOUBLE, nullNumeric);
break;
case ARRAY:
final RelDataType elementType = columnTypeToRelDataType(
typeFactory,
(ColumnType) columnType.getElementType(),
nullNumeric
);
type = typeFactory.createTypeWithNullability(
typeFactory.createArrayType(elementType, -1),
true
);
break;
case COMPLEX:
type = makeComplexType(typeFactory, columnType, true);
break;
default:
throw new ISE("valueType[%s] not translatable", columnType);
} return type;
}

/**
* Creates a {@link ComplexSqlType} using the supplied {@link RelDataTypeFactory} to ensure that the
* {@link ComplexSqlType} is interned. This is important because Calcite checks that the references are equal
Expand Down
160 changes: 160 additions & 0 deletions sql/src/test/java/org/apache/druid/sql/calcite/CalciteArraysQueryTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7327,4 +7327,164 @@ public void testBooleanConstExprArray()
)
);
}

@Test
public void testGroupByNestedArrayInline()
{
cannotVectorize();
// msq does not support nested arrays currently
msqIncompatible();
testQuery(
"SELECT c1, ARRAY_PREPEND('1', ARRAY_AGG(ARRAY[1,c2], 100000)) c5 \n"
+ "FROM (VALUES (1,1),(2,2),(3,3)) t(c1,c2)\n"
+ "GROUP BY 1 \n"
+ "HAVING ARRAY_PREPEND('1', ARRAY_AGG(ARRAY[1,c2], 100000)) <> ARRAY_PREPEND('0', ARRAY_AGG(ARRAY[1,c2], 100000))",
QUERY_CONTEXT_NO_STRINGIFY_ARRAY,
ImmutableList.of(
GroupByQuery.builder()
.setDataSource(
InlineDataSource.fromIterable(
ImmutableList.of(
new Object[]{1L, 1L},
new Object[]{2L, 2L},
new Object[]{3L, 3L}
),
RowSignature.builder()
.add("c1", ColumnType.LONG)
.add("c2", ColumnType.LONG)
.build()
)
)
.setInterval(querySegmentSpec(Filtration.eternity()))
.setGranularity(Granularities.ALL)
.setVirtualColumns(
expressionVirtualColumn(
"v0",
"array(1,\"c2\")",
ColumnType.LONG_ARRAY
)
)
.setDimensions(new DefaultDimensionSpec("c1", "d0", ColumnType.LONG))
.setAggregatorSpecs(
new ExpressionLambdaAggregatorFactory(
"a0",
ImmutableSet.of("v0"),
"__acc",
"ARRAY<ARRAY<LONG>>[]",
"ARRAY<ARRAY<LONG>>[]",
true,
true,
false,
"array_append(\"__acc\", \"v0\")",
"array_concat(\"__acc\", \"a0\")",
null,
null,
HumanReadableBytes.valueOf(100000),
TestExprMacroTable.INSTANCE
)
)
.setPostAggregatorSpecs(
expressionPostAgg(
"p0",
"array_prepend('1',\"a0\")",
ColumnType.ofArray(ColumnType.LONG_ARRAY)
)
)
.setHavingSpec(
new DimFilterHavingSpec(
expressionFilter("(array_prepend('1',\"a0\") != array_prepend('0',\"a0\"))"),
true
)
)
.setContext(QUERY_CONTEXT_DEFAULT)
.build()
),
ImmutableList.of(
new Object[]{1, ImmutableList.of(ImmutableList.of(1L), ImmutableList.of(1L, 1L))},
new Object[]{2, ImmutableList.of(ImmutableList.of(1L), ImmutableList.of(1L, 2L))},
new Object[]{3, ImmutableList.of(ImmutableList.of(1L), ImmutableList.of(1L, 3L))}
)
);
}

@Test
public void testGroupByNestedArrayInlineCount()
{
cannotVectorize();
// msq does not support nested arrays currently
msqIncompatible();
testQuery(
"SELECT COUNT(*) c FROM (\n"
+ "SELECT c1, ARRAY_PREPEND('1', ARRAY_AGG(ARRAY[1,c2], 100000)) c5 \n"
+ "FROM (VALUES (1,1),(2,2),(3,3)) t(c1,c2)\n"
+ "GROUP BY 1 \n"
+ "HAVING ARRAY_PREPEND('1', ARRAY_AGG(ARRAY[1,c2], 100000)) <> ARRAY_PREPEND('0', ARRAY_AGG(ARRAY[1,c2], 100000))\n"
+ ")",
QUERY_CONTEXT_NO_STRINGIFY_ARRAY,
ImmutableList.of(
GroupByQuery.builder()
.setDataSource(
GroupByQuery.builder()
.setDataSource(
InlineDataSource.fromIterable(
ImmutableList.of(
new Object[]{1L, 1L},
new Object[]{2L, 2L},
new Object[]{3L, 3L}
),
RowSignature.builder()
.add("c1", ColumnType.LONG)
.add("c2", ColumnType.LONG)
.build()
)
)
.setInterval(querySegmentSpec(Filtration.eternity()))
.setGranularity(Granularities.ALL)
.setVirtualColumns(
expressionVirtualColumn(
"v0",
"array(1,\"c2\")",
ColumnType.LONG_ARRAY
)
)
.setDimensions(new DefaultDimensionSpec("c1", "d0", ColumnType.LONG))
.setAggregatorSpecs(
new ExpressionLambdaAggregatorFactory(
"a0",
ImmutableSet.of("v0"),
"__acc",
"ARRAY<ARRAY<LONG>>[]",
"ARRAY<ARRAY<LONG>>[]",
true,
true,
false,
"array_append(\"__acc\", \"v0\")",
"array_concat(\"__acc\", \"a0\")",
null,
null,
HumanReadableBytes.valueOf(100000),
TestExprMacroTable.INSTANCE
)
)
.setHavingSpec(
new DimFilterHavingSpec(
expressionFilter(
"(array_prepend('1',\"a0\") != array_prepend('0',\"a0\"))"),
true
)
)
.setContext(QUERY_CONTEXT_DEFAULT)
.build()
)
.setInterval(querySegmentSpec(Filtration.eternity()))
.setGranularity(Granularities.ALL)
.setAggregatorSpecs(new CountAggregatorFactory("_a0"))
.setContext(QUERY_CONTEXT_DEFAULT)
.build()
),
ImmutableList.of(
new Object[]{3L}
)
);
}
}