Skip to content

Commit

Permalink
support non-constant expressions for path arguments for json_value an…
Browse files Browse the repository at this point in the history 
…d json_query (apache#15320)

* support dynamic expressions for path arguments for json_value and json_query
  • Loading branch information
@clintropolis
clintropolis authored Nov 17, 2023
1 parent a291478 commit a95c22c
Show file tree
Hide file tree
Showing 5 changed files with 329 additions and 98 deletions.
2 changes: 1 addition & 1 deletion docs/querying/math-expr.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -235,7 +235,7 @@ JSON functions provide facilities to extract, transform, and create `COMPLEX<jso

| function | description |
|---|---|
| json_value(expr, path[, type]) | Extract a Druid literal (`STRING`, `LONG`, `DOUBLE`) value from `expr` using JSONPath syntax of `path`. The optional `type` argument can be set to `'LONG'`,`'DOUBLE'` or `'STRING'` to cast values to that type. |
| json_value(expr, path[, type]) | Extract a Druid literal (`STRING`, `LONG`, `DOUBLE`, `ARRAY<STRING>`, `ARRAY<LONG>`, or `ARRAY<DOUBLE>`) value from `expr` using JSONPath syntax of `path`. The optional `type` argument can be set to `'LONG'`,`'DOUBLE'`, `'STRING'`, `'ARRAY<LONG>'`, `'ARRAY<DOUBLE>'`, or `'ARRAY<STRING>'` to cast values to that type. |
| json_query(expr, path) | Extract a `COMPLEX<json>` value from `expr` using JSONPath syntax of `path` |
| json_object(expr1, expr2[, expr3, expr4 ...]) | Construct a `COMPLEX<json>` with alternating 'key' and 'value' arguments|
| parse_json(expr) | Deserialize a JSON `STRING` into a `COMPLEX<json>`. If the input is not a `STRING` or it is invalid JSON, this function will result in an error.|
Expand Down
307 changes: 216 additions & 91 deletions processing/src/main/java/org/apache/druid/query/expression/NestedDataExpressions.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -330,88 +330,163 @@ public String name()
@Override
public Expr apply(List<Expr> args)
{
final List<NestedPathPart> parts = getJsonPathPartsFromLiteral(this, args.get(1));
if (args.size() == 3 && args.get(2).isLiteral()) {
final ExpressionType castTo = ExpressionType.fromString((String) args.get(2).getLiteralValue());
if (args.get(1).isLiteral()) {
if (args.size() == 3 && args.get(2).isLiteral()) {
return new JsonValueCastExpr(args);
} else {
return new JsonValueExpr(args);
}
} else {
return new JsonValueDynamicExpr(args);
}
}

final class JsonValueExpr extends ExprMacroTable.BaseScalarMacroFunctionExpr
{
private final List<NestedPathPart> parts;

public JsonValueExpr(List<Expr> args)
{
super(name(), args);
this.parts = getJsonPathPartsFromLiteral(JsonValueExprMacro.this, args.get(1));
}

@Override
public ExprEval eval(ObjectBinding bindings)
{
final ExprEval input = args.get(0).eval(bindings);
final ExprEval valAtPath = ExprEval.bestEffortOf(
NestedPathFinder.find(unwrap(input), parts)
);
if (valAtPath.type().isPrimitive() || valAtPath.type().isPrimitiveArray()) {
return valAtPath;
}
return ExprEval.of(null);
}

@Override
public Expr visit(Shuttle shuttle)
{
List<Expr> newArgs = args.stream().map(x -> x.visit(shuttle)).collect(Collectors.toList());
if (newArgs.get(1).isLiteral()) {
return shuttle.visit(new JsonValueExpr(newArgs));
} else {
return shuttle.visit(new JsonValueDynamicExpr(newArgs));
}
}

@Nullable
@Override
public ExpressionType getOutputType(InputBindingInspector inspector)
{
// we cannot infer output type because there could be anything at the path, and, we lack a proper VARIANT type
return null;
}
}

final class JsonValueCastExpr extends ExprMacroTable.BaseScalarMacroFunctionExpr
{
private final List<NestedPathPart> parts;
private final ExpressionType castTo;

public JsonValueCastExpr(List<Expr> args)
{
super(name(), args);
this.parts = getJsonPathPartsFromLiteral(JsonValueExprMacro.this, args.get(1));
this.castTo = ExpressionType.fromString((String) args.get(2).getLiteralValue());
if (castTo == null) {
throw JsonValueExprMacro.this.validationFailed(
"invalid output type: [%s]",
args.get(2).getLiteralValue()
);
}
final class JsonValueCastExpr extends ExprMacroTable.BaseScalarMacroFunctionExpr
{
public JsonValueCastExpr(List<Expr> args)
{
super(name(), args);
}

@Override
public ExprEval eval(ObjectBinding bindings)
{
final ExprEval input = args.get(0).eval(bindings);
final ExprEval valAtPath = ExprEval.bestEffortOf(
NestedPathFinder.find(unwrap(input), parts)
);
if (valAtPath.type().isPrimitive() || valAtPath.type().isPrimitiveArray()) {
return valAtPath.castTo(castTo);
}
return ExprEval.ofType(castTo, null);
}
}

@Override
public Expr visit(Shuttle shuttle)
{
List<Expr> newArgs = args.stream().map(x -> x.visit(shuttle)).collect(Collectors.toList());
return shuttle.visit(new JsonValueCastExpr(newArgs));
}
@Override
public ExprEval eval(ObjectBinding bindings)
{
final ExprEval input = args.get(0).eval(bindings);
final ExprEval valAtPath = ExprEval.bestEffortOf(
NestedPathFinder.find(unwrap(input), parts)
);
if (valAtPath.type().isPrimitive() || valAtPath.type().isPrimitiveArray()) {
return valAtPath.castTo(castTo);
}
return ExprEval.ofType(castTo, null);
}

@Nullable
@Override
public ExpressionType getOutputType(InputBindingInspector inspector)
{
return castTo;
}
@Override
public Expr visit(Shuttle shuttle)
{
List<Expr> newArgs = args.stream().map(x -> x.visit(shuttle)).collect(Collectors.toList());
if (newArgs.get(1).isLiteral()) {
return shuttle.visit(new JsonValueCastExpr(newArgs));
} else {
return shuttle.visit(new JsonValueDynamicExpr(newArgs));
}
return new JsonValueCastExpr(args);
} else {
final class JsonValueExpr extends ExprMacroTable.BaseScalarMacroFunctionExpr
{
}

public JsonValueExpr(List<Expr> args)
{
super(name(), args);
}
@Nullable
@Override
public ExpressionType getOutputType(InputBindingInspector inspector)
{
return castTo;
}
}

final class JsonValueDynamicExpr extends ExprMacroTable.BaseScalarMacroFunctionExpr
{
public JsonValueDynamicExpr(List<Expr> args)
{
super(name(), args);
}

@Override
public ExprEval eval(ObjectBinding bindings)
{
final ExprEval input = args.get(0).eval(bindings);
final ExprEval valAtPath = ExprEval.bestEffortOf(
NestedPathFinder.find(unwrap(input), parts)
@Override
public ExprEval eval(ObjectBinding bindings)
{
final ExprEval input = args.get(0).eval(bindings);
final ExprEval path = args.get(1).eval(bindings);
final ExpressionType castTo;
if (args.size() == 3) {
castTo = ExpressionType.fromString(args.get(2).eval(bindings).asString());
if (castTo == null) {
throw JsonValueExprMacro.this.validationFailed(
"invalid output type: [%s]",
args.get(2).getLiteralValue()
);
if (valAtPath.type().isPrimitive() || valAtPath.type().isPrimitiveArray()) {
return valAtPath;
}
return ExprEval.of(null);
}
} else {
castTo = null;
}
final List<NestedPathPart> parts = NestedPathFinder.parseJsonPath(path.asString());
final ExprEval<?> valAtPath = ExprEval.bestEffortOf(NestedPathFinder.find(unwrap(input), parts));
if (valAtPath.type().isPrimitive() || valAtPath.type().isPrimitiveArray()) {
return castTo == null ? valAtPath : valAtPath.castTo(castTo);
}
return castTo == null ? ExprEval.of(null) : ExprEval.ofType(castTo, null);
}

@Override
public Expr visit(Shuttle shuttle)
{
List<Expr> newArgs = args.stream().map(x -> x.visit(shuttle)).collect(Collectors.toList());
@Override
public Expr visit(Shuttle shuttle)
{
List<Expr> newArgs = args.stream().map(x -> x.visit(shuttle)).collect(Collectors.toList());
if (newArgs.get(1).isLiteral()) {
if (newArgs.size() == 3 && newArgs.get(2).isLiteral()) {
return shuttle.visit(new JsonValueCastExpr(newArgs));
} else {
return shuttle.visit(new JsonValueExpr(newArgs));
}

@Nullable
@Override
public ExpressionType getOutputType(InputBindingInspector inspector)
{
// we cannot infer output type because there could be anything at the path, and, we lack a proper VARIANT type
return null;
}
} else {
return shuttle.visit(new JsonValueDynamicExpr(newArgs));
}
return new JsonValueExpr(args);
}

@Nullable
@Override
public ExpressionType getOutputType(InputBindingInspector inspector)
{
// we cannot infer output type because there could be anything at the path, and, we lack a proper VARIANT type
return null;
}
}
}
Expand All @@ -429,40 +504,90 @@ public String name()
@Override
public Expr apply(List<Expr> args)
{
final List<NestedPathPart> parts = getJsonPathPartsFromLiteral(this, args.get(1));
final class JsonQueryExpr extends ExprMacroTable.BaseScalarMacroFunctionExpr
if (args.get(1).isLiteral()) {
return new JsonQueryExpr(args);
} else {
return new JsonQueryDynamicExpr(args);
}
}

final class JsonQueryExpr extends ExprMacroTable.BaseScalarMacroFunctionExpr
{
private final List<NestedPathPart> parts;

public JsonQueryExpr(List<Expr> args)
{
public JsonQueryExpr(List<Expr> args)
{
super(name(), args);
}
super(name(), args);
this.parts = getJsonPathPartsFromLiteral(JsonQueryExprMacro.this, args.get(1));
}

@Override
public ExprEval eval(ObjectBinding bindings)
{
ExprEval input = args.get(0).eval(bindings);
return ExprEval.ofComplex(
ExpressionType.NESTED_DATA,
NestedPathFinder.find(unwrap(input), parts)
);
}
@Override
public ExprEval eval(ObjectBinding bindings)
{
ExprEval input = args.get(0).eval(bindings);
return ExprEval.ofComplex(
ExpressionType.NESTED_DATA,
NestedPathFinder.find(unwrap(input), parts)
);
}

@Override
public Expr visit(Shuttle shuttle)
{
List<Expr> newArgs = args.stream().map(x -> x.visit(shuttle)).collect(Collectors.toList());
@Override
public Expr visit(Shuttle shuttle)
{
List<Expr> newArgs = args.stream().map(x -> x.visit(shuttle)).collect(Collectors.toList());
if (newArgs.get(1).isLiteral()) {
return shuttle.visit(new JsonQueryExpr(newArgs));
} else {
return shuttle.visit(new JsonQueryDynamicExpr(newArgs));
}
}

@Nullable
@Override
public ExpressionType getOutputType(InputBindingInspector inspector)
{
// call all the output JSON typed
return ExpressionType.NESTED_DATA;
@Nullable
@Override
public ExpressionType getOutputType(InputBindingInspector inspector)
{
// call all the output JSON typed
return ExpressionType.NESTED_DATA;
}
}

final class JsonQueryDynamicExpr extends ExprMacroTable.BaseScalarMacroFunctionExpr
{
public JsonQueryDynamicExpr(List<Expr> 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<NestedPathPart> parts = NestedPathFinder.parseJsonPath(path.asString());
return ExprEval.ofComplex(
ExpressionType.NESTED_DATA,
NestedPathFinder.find(unwrap(input), parts)
);
}

@Override
public Expr visit(Shuttle shuttle)
{
List<Expr> newArgs = args.stream().map(x -> x.visit(shuttle)).collect(Collectors.toList());
if (newArgs.get(1).isLiteral()) {
return shuttle.visit(new JsonQueryExpr(newArgs));
} else {
return shuttle.visit(new JsonQueryDynamicExpr(newArgs));
}
}
return new JsonQueryExpr(args);

@Nullable
@Override
public ExpressionType getOutputType(InputBindingInspector inspector)
{
// call all the output JSON typed
return ExpressionType.NESTED_DATA;
}
}
}

Expand Down
10 changes: 10 additions & 0 deletions processing/src/test/java/org/apache/druid/query/expression/NestedDataExpressionsTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -272,6 +272,11 @@ public void testJsonValueExpression()
eval = expr.eval(inputBindings);
Assert.assertArrayEquals(new Object[]{"1", "2", "3"}, (Object[]) eval.value());
Assert.assertEquals(ExpressionType.STRING_ARRAY, eval.type());

expr = Parser.parse("json_value(nester, array_offset(json_paths(nester), 0))", MACRO_TABLE);
eval = expr.eval(inputBindings);
Assert.assertArrayEquals(new Object[]{"a", "b", "c"}, (Object[]) eval.value());
Assert.assertEquals(ExpressionType.STRING_ARRAY, eval.type());
}

@Test
Expand Down Expand Up @@ -317,6 +322,11 @@ public void testJsonQueryExpression()
eval = expr.eval(inputBindings);
Assert.assertEquals(1234L, eval.value());
Assert.assertEquals(ExpressionType.NESTED_DATA, eval.type());

expr = Parser.parse("json_query(nester, array_offset(json_paths(nester), 0))", MACRO_TABLE);
eval = expr.eval(inputBindings);
Assert.assertEquals(NESTER.get("x"), eval.value());
Assert.assertEquals(ExpressionType.NESTED_DATA, eval.type());
}

@Test
Expand Down
Loading

0 comments on commit a95c22c

Please sign in to comment.