diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/table/RowSignatures.java b/sql/src/main/java/org/apache/druid/sql/calcite/table/RowSignatures.java index 528a0a4a0b5..0234d6b5319 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/table/RowSignatures.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/table/RowSignatures.java @@ -96,8 +96,8 @@ public class RowSignatures } /** - * 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) { @@ -105,8 +105,8 @@ public class RowSignatures } /** - * 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, @@ -126,44 +126,7 @@ public class RowSignatures 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); @@ -172,6 +135,49 @@ public class RowSignatures 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 diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteArraysQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteArraysQueryTest.java index 2c165ffe3c3..67c04a9cf91 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteArraysQueryTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteArraysQueryTest.java @@ -7327,4 +7327,164 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest ) ); } + + @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>[]", + 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>[]", + 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} + ) + ); + } }