Skip to content

Commit

Permalink
Cache resolved operators for JSON path binary expressions
Browse files Browse the repository at this point in the history
kasiafi committed May 27, 2022
1 parent df4eddd commit 9376c36
Showing 18 changed files with 2,155 additions and 1,776 deletions.
193 changes: 193 additions & 0 deletions core/trino-main/src/main/java/io/trino/json/CachingResolver.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,193 @@
/*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.trino.json;

import com.google.common.cache.CacheBuilder;
import com.google.common.collect.ImmutableList;
import io.trino.FullConnectorSession;
import io.trino.Session;
import io.trino.collect.cache.NonEvictableCache;
import io.trino.json.ir.IrPathNode;
import io.trino.metadata.BoundSignature;
import io.trino.metadata.Metadata;
import io.trino.metadata.OperatorNotFoundException;
import io.trino.metadata.ResolvedFunction;
import io.trino.spi.connector.ConnectorSession;
import io.trino.spi.function.OperatorType;
import io.trino.spi.type.Type;
import io.trino.spi.type.TypeManager;
import io.trino.type.TypeCoercion;

import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.ExecutionException;

import static com.google.common.base.Preconditions.checkState;
import static io.trino.collect.cache.SafeCaches.buildNonEvictableCache;
import static io.trino.json.CachingResolver.ResolvedOperatorAndCoercions.RESOLUTION_ERROR;
import static io.trino.json.CachingResolver.ResolvedOperatorAndCoercions.operators;
import static java.util.Objects.requireNonNull;

/**
* A resolver providing coercions and binary operators used for JSON path evaluation,
* based on the operation type and the input types.
* <p>
* It is instantiated per-driver, and caches the resolved operators and coercions.
* Caching is applied to IrArithmeticBinary and IrComparisonPredicate path nodes.
* Its purpose is to avoid resolving operators and coercions on a per-row basis, assuming
* that the input types to the JSON path operations repeat across rows.
* <p>
* If an operator or a component coercion cannot be resolved for certain input types,
* it is cached as RESOLUTION_ERROR. It depends on the caller to handle this condition.
*/
public class CachingResolver
{
private static final int MAX_CACHE_SIZE = 1000;

private final Metadata metadata;
private final Session session;
private final TypeCoercion typeCoercion;
private final NonEvictableCache<NodeAndTypes, ResolvedOperatorAndCoercions> operators = buildNonEvictableCache(CacheBuilder.newBuilder().maximumSize(MAX_CACHE_SIZE));

public CachingResolver(Metadata metadata, ConnectorSession connectorSession, TypeManager typeManager)
{
requireNonNull(metadata, "metadata is null");
requireNonNull(connectorSession, "connectorSession is null");
requireNonNull(typeManager, "typeManager is null");

this.metadata = metadata;
this.session = ((FullConnectorSession) connectorSession).getSession();
this.typeCoercion = new TypeCoercion(typeManager::getType);
}

public ResolvedOperatorAndCoercions getOperators(IrPathNode node, OperatorType operatorType, Type leftType, Type rightType)
{
try {
return operators
.get(new NodeAndTypes(IrPathNodeRef.of(node), leftType, rightType), () -> resolveOperators(operatorType, leftType, rightType));
}
catch (ExecutionException e) {
throw new RuntimeException(e);
}
}

private ResolvedOperatorAndCoercions resolveOperators(OperatorType operatorType, Type leftType, Type rightType)
{
ResolvedFunction operator;
try {
operator = metadata.resolveOperator(session, operatorType, ImmutableList.of(leftType, rightType));
}
catch (OperatorNotFoundException e) {
return RESOLUTION_ERROR;
}

BoundSignature signature = operator.getSignature();

Optional<ResolvedFunction> leftCast = Optional.empty();
if (!signature.getArgumentTypes().get(0).equals(leftType) && !typeCoercion.isTypeOnlyCoercion(leftType, signature.getArgumentTypes().get(0))) {
try {
leftCast = Optional.of(metadata.getCoercion(session, leftType, signature.getArgumentTypes().get(0)));
}
catch (OperatorNotFoundException e) {
return RESOLUTION_ERROR;
}
}

Optional<ResolvedFunction> rightCast = Optional.empty();
if (!signature.getArgumentTypes().get(1).equals(rightType) && !typeCoercion.isTypeOnlyCoercion(rightType, signature.getArgumentTypes().get(1))) {
try {
rightCast = Optional.of(metadata.getCoercion(session, rightType, signature.getArgumentTypes().get(1)));
}
catch (OperatorNotFoundException e) {
return RESOLUTION_ERROR;
}
}

return operators(operator, leftCast, rightCast);
}

private static class NodeAndTypes
{
private final IrPathNodeRef<IrPathNode> node;
private final Type leftType;
private final Type rightType;

public NodeAndTypes(IrPathNodeRef<IrPathNode> node, Type leftType, Type rightType)
{
this.node = node;
this.leftType = leftType;
this.rightType = rightType;
}

@Override
public boolean equals(Object o)
{
if (this == o) {
return true;
}
if (o == null || getClass() != o.getClass()) {
return false;
}
NodeAndTypes that = (NodeAndTypes) o;
return Objects.equals(node, that.node) &&
Objects.equals(leftType, that.leftType) &&
Objects.equals(rightType, that.rightType);
}

@Override
public int hashCode()
{
return Objects.hash(node, leftType, rightType);
}
}

public static class ResolvedOperatorAndCoercions
{
public static final ResolvedOperatorAndCoercions RESOLUTION_ERROR = new ResolvedOperatorAndCoercions(null, Optional.empty(), Optional.empty());

private final ResolvedFunction operator;
private final Optional<ResolvedFunction> leftCoercion;
private final Optional<ResolvedFunction> rightCoercion;

public static ResolvedOperatorAndCoercions operators(ResolvedFunction operator, Optional<ResolvedFunction> leftCoercion, Optional<ResolvedFunction> rightCoercion)
{
return new ResolvedOperatorAndCoercions(requireNonNull(operator, "operator is null"), leftCoercion, rightCoercion);
}

private ResolvedOperatorAndCoercions(ResolvedFunction operator, Optional<ResolvedFunction> leftCoercion, Optional<ResolvedFunction> rightCoercion)
{
this.operator = operator;
this.leftCoercion = requireNonNull(leftCoercion, "leftCoercion is null");
this.rightCoercion = requireNonNull(rightCoercion, "rightCoercion is null");
}

public ResolvedFunction getOperator()
{
checkState(this != RESOLUTION_ERROR, "accessing operator on RESOLUTION_ERROR");
return operator;
}

public Optional<ResolvedFunction> getLeftCoercion()
{
checkState(this != RESOLUTION_ERROR, "accessing coercion on RESOLUTION_ERROR");
return leftCoercion;
}

public Optional<ResolvedFunction> getRightCoercion()
{
checkState(this != RESOLUTION_ERROR, "accessing coercion on RESOLUTION_ERROR");
return rightCoercion;
}
}
}
68 changes: 68 additions & 0 deletions core/trino-main/src/main/java/io/trino/json/IrPathNodeRef.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,68 @@
/*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.trino.json;

import io.trino.json.ir.IrPathNode;

import static java.lang.String.format;
import static java.lang.System.identityHashCode;
import static java.util.Objects.requireNonNull;

public final class IrPathNodeRef<T extends IrPathNode>
{
public static <T extends IrPathNode> IrPathNodeRef<T> of(T pathNode)
{
return new IrPathNodeRef<>(pathNode);
}

private final T pathNode;

private IrPathNodeRef(T pathNode)
{
this.pathNode = requireNonNull(pathNode, "pathNode is null");
}

public T getNode()
{
return pathNode;
}

@Override
public boolean equals(Object o)
{
if (this == o) {
return true;
}
if (o == null || getClass() != o.getClass()) {
return false;
}
IrPathNodeRef<?> other = (IrPathNodeRef<?>) o;
return pathNode == other.pathNode;
}

@Override
public int hashCode()
{
return identityHashCode(pathNode);
}

@Override
public String toString()
{
return format(
"@%s: %s",
Integer.toHexString(identityHashCode(pathNode)),
pathNode);
}
}
78 changes: 78 additions & 0 deletions core/trino-main/src/main/java/io/trino/json/JsonPathEvaluator.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
/*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.trino.json;

import com.fasterxml.jackson.databind.JsonNode;
import io.trino.json.ir.IrJsonPath;
import io.trino.metadata.FunctionManager;
import io.trino.metadata.Metadata;
import io.trino.metadata.ResolvedFunction;
import io.trino.spi.connector.ConnectorSession;
import io.trino.spi.type.TypeManager;
import io.trino.sql.InterpretedFunctionInvoker;

import java.util.List;

import static java.util.Objects.requireNonNull;

/**
* Evaluates the JSON path expression using given JSON input and parameters,
* respecting the path mode `strict` or `lax`.
* Successful evaluation results in a sequence of objects.
* Each object in the sequence is either a `JsonNode` or a `TypedValue`
* Certain error conditions might be suppressed in `lax` mode.
* Any unsuppressed error condition causes evaluation failure.
* In such case, `PathEvaluationError` is thrown.
*/
public class JsonPathEvaluator
{
private final IrJsonPath path;
private final Invoker invoker;
private final CachingResolver resolver;

public JsonPathEvaluator(IrJsonPath path, ConnectorSession session, Metadata metadata, TypeManager typeManager, FunctionManager functionManager)
{
requireNonNull(path, "path is null");
requireNonNull(session, "session is null");
requireNonNull(metadata, "metadata is null");
requireNonNull(typeManager, "typeManager is null");
requireNonNull(functionManager, "functionManager is null");

this.path = path;
this.invoker = new Invoker(session, functionManager);
this.resolver = new CachingResolver(metadata, session, typeManager);
}

public List<Object> evaluate(JsonNode input, Object[] parameters)
{
return new PathEvaluationVisitor(path.isLax(), input, parameters, invoker, resolver).process(path.getRoot(), new PathEvaluationContext());
}

public static class Invoker
{
private final ConnectorSession connectorSession;
private final InterpretedFunctionInvoker functionInvoker;

public Invoker(ConnectorSession connectorSession, FunctionManager functionManager)
{
this.connectorSession = connectorSession;
this.functionInvoker = new InterpretedFunctionInvoker(functionManager);
}

public Object invoke(ResolvedFunction function, List<Object> arguments)
{
return functionInvoker.invoke(function, connectorSession, arguments);
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
/*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.trino.json;

/**
* A class representing state, used by JSON-processing functions: JSON_EXISTS, JSON_VALUE, and JSON_QUERY.
* It is instantiated per driver, and allows gathering and reusing information across the processed rows.
* It contains a JsonPathEvaluator object, which caches ResolvedFunctions used by certain path nodes.
* Caching the ResolvedFunctions addresses the assumption that all or most rows shall provide values
* of the same types to certain JSON path operators.
*/
public class JsonPathInvocationContext
{
private JsonPathEvaluator evaluator;

public JsonPathEvaluator getEvaluator()
{
return evaluator;
}

public void setEvaluator(JsonPathEvaluator evaluator)
{
this.evaluator = evaluator;
}
}
Loading

0 comments on commit 9376c36

Please sign in to comment.