fix bug with json_object expression not fully unwrapping inputs (#12893)

This commit is contained in:
Clint Wylie 2022-08-13 21:15:19 -07:00 committed by GitHub
parent 5394838030
commit f4e0909e92
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 24 additions and 4 deletions

View File

@ -38,6 +38,7 @@ import org.apache.druid.segment.nested.StructuredData;
import org.apache.druid.segment.nested.StructuredDataProcessor;
import javax.annotation.Nullable;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.stream.Collectors;
@ -80,7 +81,7 @@ public class NestedDataExpressions
ExprEval value = args.get(i + 1).eval(bindings);
Preconditions.checkArgument(field.type().is(ExprType.STRING), "field name must be a STRING");
theMap.put(field.asString(), value.value());
theMap.put(field.asString(), maybeUnwrapStructuredData(value.value()));
}
return ExprEval.ofComplex(TYPE, theMap);
@ -612,13 +613,22 @@ public class NestedDataExpressions
@Nullable
static Object maybeUnwrapStructuredData(ExprEval input)
{
if (input.value() instanceof StructuredData) {
StructuredData data = (StructuredData) input.value();
return maybeUnwrapStructuredData(input.value());
}
static Object maybeUnwrapStructuredData(Object input)
{
if (input instanceof StructuredData) {
StructuredData data = (StructuredData) input;
return data.getValue();
}
return input.value();
if (input instanceof Object[]) {
return Arrays.stream((Object[]) input).map(x -> maybeUnwrapStructuredData(x)).toArray();
}
return input;
}
static List<NestedPathPart> getArg1PathPartsFromLiteral(String fnName, List<Expr> args)
{
if (!(args.get(1).isLiteral() && args.get(1).getLiteralValue() instanceof String)) {

View File

@ -29,6 +29,7 @@ import org.apache.druid.math.expr.ExprMacroTable;
import org.apache.druid.math.expr.ExpressionType;
import org.apache.druid.math.expr.InputBindings;
import org.apache.druid.math.expr.Parser;
import org.apache.druid.segment.nested.StructuredData;
import org.apache.druid.testing.InitializedNullHandlingTest;
import org.junit.Assert;
import org.junit.Test;
@ -67,6 +68,7 @@ public class NestedDataExpressionsTest extends InitializedNullHandlingTest
Expr.ObjectBinding inputBindings = InputBindings.withTypedSuppliers(
new ImmutableMap.Builder<String, Pair<ExpressionType, Supplier<Object>>>()
.put("nest", new Pair<>(NestedDataExpressions.TYPE, () -> NEST))
.put("nestWrapped", new Pair<>(NestedDataExpressions.TYPE, () -> new StructuredData(NEST)))
.put("nester", new Pair<>(NestedDataExpressions.TYPE, () -> NESTER))
.put("string", new Pair<>(ExpressionType.STRING, () -> "abcdef"))
.put("long", new Pair<>(ExpressionType.LONG, () -> 1234L))
@ -345,5 +347,13 @@ public class NestedDataExpressionsTest extends InitializedNullHandlingTest
eval = expr.eval(inputBindings);
Assert.assertEquals(100L, eval.value());
Assert.assertEquals(ExpressionType.LONG, eval.type());
expr = Parser.parse("to_json_string(json_object('x', nestWrapped))", MACRO_TABLE);
eval = expr.eval(inputBindings);
Assert.assertEquals("{\"x\":{\"x\":100,\"y\":200,\"z\":300}}", eval.value());
expr = Parser.parse("to_json_string(json_object('xs', array(nest, nestWrapped)))", MACRO_TABLE);
eval = expr.eval(inputBindings);
Assert.assertEquals("{\"xs\":[{\"x\":100,\"y\":200,\"z\":300},{\"x\":100,\"y\":200,\"z\":300}]}", eval.value());
}
}