From 46469e0c3098a3063a66179a8e757858228251a3 Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Thu, 7 Dec 2023 23:42:38 -0800 Subject: [PATCH] add JSON_QUERY_ARRAY function to pluck ARRAY> out of COMPLEX --- .../apache/druid/guice/ExpressionModule.java | 1 + .../expression/NestedDataExpressions.java | 117 ++++++++++ .../segment/column/ColumnTypeFactory.java | 2 +- .../virtual/NestedFieldVirtualColumn.java | 7 + .../expression/NestedDataExpressionsTest.java | 34 +++ .../NestedDataOperatorConversions.java | 32 +++ .../calcite/planner/DruidOperatorTable.java | 1 + .../calcite/planner/DruidSqlValidator.java | 25 ++ .../calcite/schema/RootSchemaProvider.java | 7 + .../calcite/CalciteNestedDataQueryTest.java | 214 ++++++++++++++++++ 10 files changed, 439 insertions(+), 1 deletion(-) diff --git a/processing/src/main/java/org/apache/druid/guice/ExpressionModule.java b/processing/src/main/java/org/apache/druid/guice/ExpressionModule.java index db5aeaf17c216..5387daead4bbf 100644 --- a/processing/src/main/java/org/apache/druid/guice/ExpressionModule.java +++ b/processing/src/main/java/org/apache/druid/guice/ExpressionModule.java @@ -83,6 +83,7 @@ public class ExpressionModule implements Module .add(NestedDataExpressions.JsonPathsExprMacro.class) .add(NestedDataExpressions.JsonValueExprMacro.class) .add(NestedDataExpressions.JsonQueryExprMacro.class) + .add(NestedDataExpressions.JsonQueryArrayExprMacro.class) .add(NestedDataExpressions.ToJsonStringExprMacro.class) .add(NestedDataExpressions.ParseJsonExprMacro.class) .add(NestedDataExpressions.TryParseJsonExprMacro.class) diff --git a/processing/src/main/java/org/apache/druid/query/expression/NestedDataExpressions.java b/processing/src/main/java/org/apache/druid/query/expression/NestedDataExpressions.java index fec16ad99b1e3..f7476adf599e9 100644 --- a/processing/src/main/java/org/apache/druid/query/expression/NestedDataExpressions.java +++ b/processing/src/main/java/org/apache/druid/query/expression/NestedDataExpressions.java @@ -28,6 +28,7 @@ import org.apache.druid.math.expr.ExprMacroTable; import org.apache.druid.math.expr.ExprType; import org.apache.druid.math.expr.ExpressionType; +import org.apache.druid.math.expr.ExpressionTypeFactory; import org.apache.druid.math.expr.NamedFunction; import org.apache.druid.segment.nested.NestedPathFinder; import org.apache.druid.segment.nested.NestedPathPart; @@ -44,6 +45,8 @@ public class NestedDataExpressions { + private static ExpressionType JSON_ARRAY = ExpressionTypeFactory.getInstance().ofArray(ExpressionType.NESTED_DATA); + public static class JsonObjectExprMacro implements ExprMacroTable.ExprMacro { public static final String NAME = "json_object"; @@ -591,6 +594,120 @@ public ExpressionType getOutputType(InputBindingInspector inspector) } } + public static class JsonQueryArrayExprMacro implements ExprMacroTable.ExprMacro + { + public static final String NAME = "json_query_array"; + + @Override + public String name() + { + return NAME; + } + + @Override + public Expr apply(List args) + { + if (args.get(1).isLiteral()) { + return new JsonQueryArrayExpr(args); + } else { + return new JsonQueryArrayDynamicExpr(args); + } + } + + final class JsonQueryArrayExpr extends ExprMacroTable.BaseScalarMacroFunctionExpr + { + private final List parts; + + public JsonQueryArrayExpr(List args) + { + super(name(), args); + this.parts = getJsonPathPartsFromLiteral(JsonQueryArrayExprMacro.this, args.get(1)); + } + + @Override + public ExprEval eval(ObjectBinding bindings) + { + ExprEval input = args.get(0).eval(bindings); + final Object value = NestedPathFinder.find(unwrap(input), parts); + if (value instanceof List) { + return ExprEval.ofArray( + JSON_ARRAY, + ExprEval.bestEffortArray((List) value).asArray() + ); + } + return ExprEval.ofArray( + JSON_ARRAY, + ExprEval.bestEffortOf(value).asArray() + ); + } + + @Override + public Expr visit(Shuttle shuttle) + { + List newArgs = args.stream().map(x -> x.visit(shuttle)).collect(Collectors.toList()); + if (newArgs.get(1).isLiteral()) { + return shuttle.visit(new JsonQueryArrayExpr(newArgs)); + } else { + return shuttle.visit(new JsonQueryArrayDynamicExpr(newArgs)); + } + } + + @Nullable + @Override + public ExpressionType getOutputType(InputBindingInspector inspector) + { + // call all the output JSON typed + return ExpressionType.NESTED_DATA; + } + } + + final class JsonQueryArrayDynamicExpr extends ExprMacroTable.BaseScalarMacroFunctionExpr + { + public JsonQueryArrayDynamicExpr(List args) + { + super(name(), args); + } + + @Override + public ExprEval eval(ObjectBinding bindings) + { + ExprEval input = args.get(0).eval(bindings); + ExprEval path = args.get(1).eval(bindings); + final List parts = NestedPathFinder.parseJsonPath(path.asString()); + final Object value = NestedPathFinder.find(unwrap(input), parts); + if (value instanceof List) { + return ExprEval.ofArray( + JSON_ARRAY, + ExprEval.bestEffortArray((List) value).asArray() + ); + } + return ExprEval.ofArray( + JSON_ARRAY, + ExprEval.bestEffortOf(value).asArray() + ); + } + + @Override + public Expr visit(Shuttle shuttle) + { + List newArgs = args.stream().map(x -> x.visit(shuttle)).collect(Collectors.toList()); + if (newArgs.get(1).isLiteral()) { + return shuttle.visit(new JsonQueryArrayExpr(newArgs)); + } else { + return shuttle.visit(new JsonQueryArrayDynamicExpr(newArgs)); + } + } + + @Nullable + @Override + public ExpressionType getOutputType(InputBindingInspector inspector) + { + // call all the output ARRAY> typed + return JSON_ARRAY; + } + } + } + public static class JsonPathsExprMacro implements ExprMacroTable.ExprMacro { public static final String NAME = "json_paths"; diff --git a/processing/src/main/java/org/apache/druid/segment/column/ColumnTypeFactory.java b/processing/src/main/java/org/apache/druid/segment/column/ColumnTypeFactory.java index 49793fba4163a..c9b153bf12112 100644 --- a/processing/src/main/java/org/apache/druid/segment/column/ColumnTypeFactory.java +++ b/processing/src/main/java/org/apache/druid/segment/column/ColumnTypeFactory.java @@ -64,7 +64,7 @@ public static ColumnType ofType(TypeSignature type) case STRING: return ColumnType.STRING_ARRAY; default: - throw new ISE("Unsupported expression type[%s]", type.asTypeString()); + return ColumnType.ofArray(ofType(type.getElementType())); } case COMPLEX: return INTERNER.intern(new ColumnType(ValueType.COMPLEX, type.getComplexTypeName(), null)); diff --git a/processing/src/main/java/org/apache/druid/segment/virtual/NestedFieldVirtualColumn.java b/processing/src/main/java/org/apache/druid/segment/virtual/NestedFieldVirtualColumn.java index e486e7b773889..0a282f5e4bc62 100644 --- a/processing/src/main/java/org/apache/druid/segment/virtual/NestedFieldVirtualColumn.java +++ b/processing/src/main/java/org/apache/druid/segment/virtual/NestedFieldVirtualColumn.java @@ -1228,6 +1228,13 @@ public ColumnCapabilities capabilities(String columnName) public ColumnCapabilities capabilities(ColumnInspector inspector, String columnName) { if (processFromRaw) { + if (expectedType != null && expectedType.isArray() && ColumnType.NESTED_DATA.equals(expectedType.getElementType())) { + // arrays of objects! + return ColumnCapabilitiesImpl.createDefault() + .setType(ColumnType.ofArray(ColumnType.NESTED_DATA)) + .setHasMultipleValues(false) + .setHasNulls(true); + } // JSON_QUERY always returns a StructuredData return ColumnCapabilitiesImpl.createDefault() .setType(ColumnType.NESTED_DATA) diff --git a/processing/src/test/java/org/apache/druid/query/expression/NestedDataExpressionsTest.java b/processing/src/test/java/org/apache/druid/query/expression/NestedDataExpressionsTest.java index d8bd2d93841b0..b4959932ea1a0 100644 --- a/processing/src/test/java/org/apache/druid/query/expression/NestedDataExpressionsTest.java +++ b/processing/src/test/java/org/apache/druid/query/expression/NestedDataExpressionsTest.java @@ -29,6 +29,7 @@ import org.apache.druid.math.expr.ExprMacroTable; import org.apache.druid.math.expr.ExpressionProcessingException; import org.apache.druid.math.expr.ExpressionType; +import org.apache.druid.math.expr.ExpressionTypeFactory; import org.apache.druid.math.expr.InputBindings; import org.apache.druid.math.expr.Parser; import org.apache.druid.segment.nested.StructuredData; @@ -37,6 +38,7 @@ import org.junit.Test; import java.util.Arrays; +import java.util.List; import java.util.Map; public class NestedDataExpressionsTest extends InitializedNullHandlingTest @@ -49,6 +51,7 @@ public class NestedDataExpressionsTest extends InitializedNullHandlingTest new NestedDataExpressions.JsonObjectExprMacro(), new NestedDataExpressions.JsonValueExprMacro(), new NestedDataExpressions.JsonQueryExprMacro(), + new NestedDataExpressions.JsonQueryArrayExprMacro(), new NestedDataExpressions.ToJsonStringExprMacro(JSON_MAPPER), new NestedDataExpressions.ParseJsonExprMacro(JSON_MAPPER), new NestedDataExpressions.TryParseJsonExprMacro(JSON_MAPPER) @@ -329,6 +332,37 @@ public void testJsonQueryExpression() Assert.assertEquals(ExpressionType.NESTED_DATA, eval.type()); } + @Test + public void testJsonQueryArrayExpression() + { + final ExpressionType NESTED_ARRAY = ExpressionTypeFactory.getInstance().ofArray(ExpressionType.NESTED_DATA); + + Expr expr = Parser.parse("json_query_array(nest, '$.x')", MACRO_TABLE); + ExprEval eval = expr.eval(inputBindings); + Assert.assertArrayEquals(new Object[]{100L}, (Object[]) eval.value()); + Assert.assertEquals(NESTED_ARRAY, eval.type()); + + expr = Parser.parse("json_query_array(nester, '$.x')", MACRO_TABLE); + eval = expr.eval(inputBindings); + Assert.assertArrayEquals(((List) NESTER.get("x")).toArray(), (Object[]) eval.value()); + Assert.assertEquals(NESTED_ARRAY, eval.type()); + + expr = Parser.parse("json_query_array(nester, array_offset(json_paths(nester), 0))", MACRO_TABLE); + eval = expr.eval(inputBindings); + Assert.assertArrayEquals(((List) NESTER.get("x")).toArray(), (Object[]) eval.value()); + Assert.assertEquals(NESTED_ARRAY, eval.type()); + + expr = Parser.parse("json_query_array(nesterer, '$.y')", MACRO_TABLE); + eval = expr.eval(inputBindings); + Assert.assertArrayEquals(((List) NESTERER.get("y")).toArray(), (Object[]) eval.value()); + Assert.assertEquals(NESTED_ARRAY, eval.type()); + + expr = Parser.parse("array_length(json_query_array(nesterer, '$.y'))", MACRO_TABLE); + eval = expr.eval(inputBindings); + Assert.assertEquals(3L, eval.value()); + Assert.assertEquals(ExpressionType.LONG, eval.type()); + } + @Test public void testParseJsonTryParseJson() throws JsonProcessingException { diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/NestedDataOperatorConversions.java b/sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/NestedDataOperatorConversions.java index bfa8f47e56ecb..a51a3d713757b 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/NestedDataOperatorConversions.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/NestedDataOperatorConversions.java @@ -48,11 +48,13 @@ import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.math.expr.Expr; import org.apache.druid.math.expr.InputBindings; +import org.apache.druid.query.expression.NestedDataExpressions; import org.apache.druid.segment.column.ColumnType; import org.apache.druid.segment.column.RowSignature; import org.apache.druid.segment.nested.NestedPathFinder; import org.apache.druid.segment.nested.NestedPathPart; import org.apache.druid.segment.virtual.NestedFieldVirtualColumn; +import org.apache.druid.sql.calcite.expression.DirectOperatorConversion; import org.apache.druid.sql.calcite.expression.DruidExpression; import org.apache.druid.sql.calcite.expression.Expressions; import org.apache.druid.sql.calcite.expression.OperatorConversions; @@ -78,6 +80,16 @@ public class NestedDataOperatorConversions true ); + public static final SqlReturnTypeInference NESTED_ARRAY_RETURN_TYPE_INFERENCE = opBinding -> + opBinding.getTypeFactory().createArrayType( + RowSignatures.makeComplexType( + opBinding.getTypeFactory(), + ColumnType.NESTED_DATA, + true + ), + -1 + ); + public static class JsonPathsOperatorConversion implements SqlOperatorConversion { private static final SqlFunction SQL_FUNCTION = OperatorConversions @@ -231,6 +243,26 @@ public DruidExpression toDruidExpression( } } + public static class JsonQueryArrayOperatorConversion extends DirectOperatorConversion + { + private static final SqlFunction SQL_FUNCTION = OperatorConversions + .operatorBuilder(StringUtils.toUpperCase(NestedDataExpressions.JsonQueryArrayExprMacro.NAME)) + .operandTypeChecker( + OperandTypes.family( + SqlTypeFamily.ANY, + SqlTypeFamily.CHARACTER + ) + ) + .returnTypeInference(NESTED_ARRAY_RETURN_TYPE_INFERENCE) + .functionCategory(SqlFunctionCategory.SYSTEM) + .build(); + + public JsonQueryArrayOperatorConversion() + { + super(SQL_FUNCTION, NestedDataExpressions.JsonQueryArrayExprMacro.NAME); + } + } + /** * The {@link org.apache.calcite.sql2rel.StandardConvertletTable} converts json_value(.. RETURNING type) into * cast(json_value_any(..), type). diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidOperatorTable.java b/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidOperatorTable.java index 6c432c800d898..cc9779a30c63c 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidOperatorTable.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidOperatorTable.java @@ -333,6 +333,7 @@ public class DruidOperatorTable implements SqlOperatorTable .add(new NestedDataOperatorConversions.JsonKeysOperatorConversion()) .add(new NestedDataOperatorConversions.JsonPathsOperatorConversion()) .add(new NestedDataOperatorConversions.JsonQueryOperatorConversion()) + .add(new NestedDataOperatorConversions.JsonQueryArrayOperatorConversion()) .add(new NestedDataOperatorConversions.JsonValueAnyOperatorConversion()) .add(new NestedDataOperatorConversions.JsonValueBigintOperatorConversion()) .add(new NestedDataOperatorConversions.JsonValueDoubleOperatorConversion()) diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidSqlValidator.java b/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidSqlValidator.java index 34f3c7410f380..a29ed92e95132 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidSqlValidator.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidSqlValidator.java @@ -19,6 +19,7 @@ package org.apache.druid.sql.calcite.planner; +import com.google.common.collect.ImmutableSet; import org.apache.calcite.adapter.java.JavaTypeFactory; import org.apache.calcite.prepare.BaseDruidSqlValidator; import org.apache.calcite.prepare.CalciteCatalogReader; @@ -29,10 +30,14 @@ import org.apache.calcite.sql.SqlNode; import org.apache.calcite.sql.SqlOperatorTable; import org.apache.calcite.sql.parser.SqlParserPos; +import org.apache.calcite.sql.type.SqlTypeMappingRule; +import org.apache.calcite.sql.type.SqlTypeName; import org.apache.calcite.sql.validate.SqlValidatorScope; import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.sql.calcite.run.EngineFeature; +import java.util.Map; + /** * Druid extended SQL validator. (At present, it doesn't actually * have any extensions yet, but it will soon.) @@ -80,6 +85,26 @@ public void validateCall(SqlCall call, SqlValidatorScope scope) super.validateCall(call, scope); } + @Override + public SqlTypeMappingRule getTypeMappingRule() + { + SqlTypeMappingRule base = super.getTypeMappingRule(); + return new SqlTypeMappingRule() + { + @Override + public Map> getTypeMapping() + { + return base.getTypeMapping(); + } + + @Override + public boolean canApplyFrom(SqlTypeName to, SqlTypeName from) + { + return SqlTypeMappingRule.super.canApplyFrom(to, from); + } + }; + } + private CalciteContextException buildCalciteContextException(String message, SqlCall call) { SqlParserPos pos = call.getParserPosition(); diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/schema/RootSchemaProvider.java b/sql/src/main/java/org/apache/druid/sql/calcite/schema/RootSchemaProvider.java index f7f4f660e23c9..e7c3147f7387d 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/schema/RootSchemaProvider.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/schema/RootSchemaProvider.java @@ -25,7 +25,10 @@ import com.google.inject.Provider; import org.apache.calcite.jdbc.CalciteSchema; import org.apache.calcite.schema.SchemaPlus; +import org.apache.calcite.sql.type.SqlTypeName; import org.apache.druid.java.util.common.ISE; +import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.sql.calcite.table.RowSignatures; import java.util.Map; import java.util.Set; @@ -66,6 +69,10 @@ public DruidSchemaCatalog get() for (NamedSchema schema : namedSchemas) { rootSchema.add(schema.getSchemaName(), schema.getSchema()); } + rootSchema.add( + "JSON", + relDataTypeFactory -> new RowSignatures.ComplexSqlType(SqlTypeName.OTHER, ColumnType.NESTED_DATA, true) + ); return new DruidSchemaCatalog(rootSchema, ImmutableMap.copyOf(schemasByName)); } } diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteNestedDataQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteNestedDataQueryTest.java index cea1bdbef7e1a..f964a0f7f3f6b 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteNestedDataQueryTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteNestedDataQueryTest.java @@ -6523,4 +6523,218 @@ public void testJsonQueryDynamicArg() ); } + + @Test + public void testJsonQueryArrays() + { + cannotVectorize(); + testBuilder() + .sql("SELECT JSON_QUERY_ARRAY(arrayObject, '$') FROM druid.arrays") + .queryContext(QUERY_CONTEXT_NO_STRINGIFY_ARRAY) + .expectedQueries( + ImmutableList.of( + Druids.newScanQueryBuilder() + .dataSource(DATA_SOURCE_ARRAYS) + .intervals(querySegmentSpec(Filtration.eternity())) + .virtualColumns( + expressionVirtualColumn("v0", "json_query_array(\"arrayObject\",'$')", ColumnType.ofArray(ColumnType.NESTED_DATA)) + ) + .columns("v0") + .context(QUERY_CONTEXT_NO_STRINGIFY_ARRAY) + .legacy(false) + .resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST) + .build() + ) + ) + .expectedResults( + ImmutableList.of( + new Object[]{ImmutableList.of("{\"x\":1000}", "{\"y\":2000}")}, + new Object[]{ImmutableList.of("{\"x\":1}", "{\"x\":2}")}, + new Object[]{ImmutableList.of("{\"x\":null}", "{\"x\":2}")}, + new Object[]{ImmutableList.of("{\"a\":1}", "{\"b\":2}")}, + new Object[]{ImmutableList.of("{\"x\":1}", "{\"x\":2}")}, + new Object[]{Arrays.asList(null, "{\"x\":2}")}, + new Object[]{ImmutableList.of("{\"x\":3}", "{\"x\":4}")}, + new Object[]{ImmutableList.of("{\"x\":1000}", "{\"y\":2000}")}, + new Object[]{ImmutableList.of("{\"x\":1}", "{\"x\":2}")}, + new Object[]{ImmutableList.of("{\"x\":null}", "{\"x\":2}")}, + new Object[]{ImmutableList.of("{\"a\":1}", "{\"b\":2}")}, + new Object[]{ImmutableList.of("{\"x\":1}", "{\"x\":2}")}, + new Object[]{Arrays.asList(null, "{\"x\":2}")}, + new Object[]{ImmutableList.of("{\"x\":3}", "{\"x\":4}")} + ) + ) + .expectedSignature( + RowSignature.builder() + .add("EXPR$0", ColumnType.ofArray(ColumnType.NESTED_DATA)) + .build() + ) + .run(); + } + @Test + public void testUnnestJsonQueryArrays() + { + cannotVectorize(); + testBuilder() + .sql("SELECT objects FROM druid.arrays, UNNEST(JSON_QUERY_ARRAY(arrayObject, '$')) as u(objects)") + .queryContext(QUERY_CONTEXT_NO_STRINGIFY_ARRAY) + .expectedQueries( + ImmutableList.of( + Druids.newScanQueryBuilder() + .dataSource( + UnnestDataSource.create( + TableDataSource.create(DATA_SOURCE_ARRAYS), + expressionVirtualColumn("j0.unnest", "json_query_array(\"arrayObject\",'$')", ColumnType.ofArray(ColumnType.NESTED_DATA)), + null + ) + ) + .intervals(querySegmentSpec(Filtration.eternity())) + .columns("j0.unnest") + .context(QUERY_CONTEXT_NO_STRINGIFY_ARRAY) + .legacy(false) + .resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST) + .build() + ) + ) + .expectedResults( + ImmutableList.of( + new Object[]{"{\"x\":1000}"}, + new Object[]{"{\"y\":2000}"}, + new Object[]{"{\"x\":1}"}, + new Object[]{"{\"x\":2}"}, + new Object[]{"{\"x\":null}"}, + new Object[]{"{\"x\":2}"}, + new Object[]{"{\"a\":1}"}, + new Object[]{"{\"b\":2}"}, + new Object[]{"{\"x\":1}"}, + new Object[]{"{\"x\":2}"}, + new Object[]{null}, + new Object[]{"{\"x\":2}"}, + new Object[]{"{\"x\":3}"}, + new Object[]{"{\"x\":4}"}, + new Object[]{"{\"x\":1000}"}, + new Object[]{"{\"y\":2000}"}, + new Object[]{"{\"x\":1}"}, + new Object[]{"{\"x\":2}"}, + new Object[]{"{\"x\":null}"}, + new Object[]{"{\"x\":2}"}, + new Object[]{"{\"a\":1}"}, + new Object[]{"{\"b\":2}"}, + new Object[]{"{\"x\":1}"}, + new Object[]{"{\"x\":2}"}, + new Object[]{null}, + new Object[]{"{\"x\":2}"}, + new Object[]{"{\"x\":3}"}, + new Object[]{"{\"x\":4}"} + ) + ) + .expectedSignature( + RowSignature.builder() + .add("objects", ColumnType.NESTED_DATA) + .build() + ) + .run(); + } + + @Test + public void testUnnestJsonQueryArraysJsonValue() + { + cannotVectorize(); + testBuilder() + .sql( + "SELECT" + + " json_value(objects, '$.x' returning bigint) as x," + + " count(*)" + + " FROM druid.arrays, UNNEST(JSON_QUERY_ARRAY(arrayObject, '$')) as u(objects)" + + " GROUP BY 1" + ) + .queryContext(QUERY_CONTEXT_NO_STRINGIFY_ARRAY) + .expectedQueries( + ImmutableList.of( + GroupByQuery.builder() + .setDataSource( + UnnestDataSource.create( + TableDataSource.create(DATA_SOURCE_ARRAYS), + expressionVirtualColumn("j0.unnest", "json_query_array(\"arrayObject\",'$')", ColumnType.ofArray(ColumnType.NESTED_DATA)), + null + ) + ) + .setInterval(querySegmentSpec(Filtration.eternity())) + .setGranularity(Granularities.ALL) + .setVirtualColumns( + new NestedFieldVirtualColumn("j0.unnest", "$.x", "v0", ColumnType.LONG) + ) + .setDimensions( + dimensions( + new DefaultDimensionSpec("v0", "d0", ColumnType.LONG) + ) + ) + .setAggregatorSpecs(new CountAggregatorFactory("a0")) + .setContext(QUERY_CONTEXT_DEFAULT) + .build() + ) + ) + .expectedResults( + ImmutableList.of( + new Object[]{null, 10L}, + new Object[]{1L, 4L}, + new Object[]{2L, 8L}, + new Object[]{3L, 2L}, + new Object[]{4L, 2L}, + new Object[]{1000L, 2L} + ) + ) + .expectedSignature( + RowSignature.builder() + .add("x", ColumnType.LONG) + .add("EXPR$1", ColumnType.LONG) + .build() + ) + .run(); + } + + @Test + public void testUnnestJsonQueryArraysJsonValueSum() + { + cannotVectorize(); + testBuilder() + .sql( + "SELECT" + + " sum(json_value(objects, '$.x' returning bigint)) as xs" + + " FROM druid.arrays, UNNEST(JSON_QUERY_ARRAY(arrayObject, '$')) as u(objects)" + ) + .queryContext(QUERY_CONTEXT_NO_STRINGIFY_ARRAY) + .expectedQueries( + ImmutableList.of( + Druids.newTimeseriesQueryBuilder() + .intervals(querySegmentSpec(Filtration.eternity())) + .dataSource( + UnnestDataSource.create( + TableDataSource.create(DATA_SOURCE_ARRAYS), + expressionVirtualColumn("j0.unnest", "json_query_array(\"arrayObject\",'$')", ColumnType.ofArray(ColumnType.NESTED_DATA)), + null + ) + ) + .virtualColumns( + new NestedFieldVirtualColumn("j0.unnest", "$.x", "v0", ColumnType.LONG) + ) + .aggregators( + new LongSumAggregatorFactory("a0", "v0") + ) + .context(QUERY_CONTEXT_DEFAULT) + .build() + ) + ) + .expectedResults( + ImmutableList.of( + new Object[]{2034L} + ) + ) + .expectedSignature( + RowSignature.builder() + .add("xs", ColumnType.LONG) + .build() + ) + .run(); + } }