Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance. Project price only 1 $
You can buy this project and download/modify it how often you want.
/*
* 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.operator.scalar;
import com.fasterxml.jackson.core.JsonFactory;
import com.fasterxml.jackson.core.JsonGenerator;
import com.fasterxml.jackson.core.JsonParseException;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.core.JsonToken;
import com.fasterxml.jackson.core.io.SerializedString;
import com.google.common.collect.ImmutableList;
import io.airlift.slice.DynamicSliceOutput;
import io.airlift.slice.Slice;
import io.trino.spi.TrinoException;
import java.io.IOException;
import java.io.UncheckedIOException;
import static com.fasterxml.jackson.core.JsonFactory.Feature.CANONICALIZE_FIELD_NAMES;
import static com.fasterxml.jackson.core.JsonToken.END_ARRAY;
import static com.fasterxml.jackson.core.JsonToken.END_OBJECT;
import static com.fasterxml.jackson.core.JsonToken.FIELD_NAME;
import static com.fasterxml.jackson.core.JsonToken.START_ARRAY;
import static com.fasterxml.jackson.core.JsonToken.START_OBJECT;
import static com.fasterxml.jackson.core.JsonToken.VALUE_NULL;
import static io.airlift.slice.Slices.utf8Slice;
import static io.trino.plugin.base.util.JsonUtils.jsonFactoryBuilder;
import static io.trino.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT;
import static io.trino.util.JsonUtil.createJsonGenerator;
import static io.trino.util.JsonUtil.createJsonParser;
import static java.util.Objects.requireNonNull;
/**
* Extracts values from JSON
*
* Supports the following JSON path primitives:
*
* $ : Root object
* . or [] : Child operator
* [] : Subscript operator for array
*
*/
public final class JsonExtract
{
private static final int ESTIMATED_JSON_OUTPUT_SIZE = 512;
private static final JsonFactory JSON_FACTORY = jsonFactoryBuilder()
.disable(CANONICALIZE_FIELD_NAMES)
.build();
private JsonExtract() {}
public static T extract(Slice jsonInput, JsonExtractor jsonExtractor)
{
requireNonNull(jsonInput, "jsonInput is null");
try {
try (JsonParser jsonParser = createJsonParser(JSON_FACTORY, jsonInput)) {
// Initialize by advancing to first token and make sure it exists
if (jsonParser.nextToken() == null) {
return null;
}
return jsonExtractor.extract(jsonParser);
}
}
catch (JsonParseException e) {
// Return null if we failed to parse something
return null;
}
catch (IOException e) {
throw new UncheckedIOException(e);
}
}
public static JsonExtractor generateExtractor(String path, JsonExtractor rootExtractor)
{
return generateExtractor(path, rootExtractor, false);
}
public static JsonExtractor generateExtractor(String path, JsonExtractor rootExtractor, boolean exceptionOnOutOfBounds)
{
ImmutableList tokens = ImmutableList.copyOf(new JsonPathTokenizer(path));
JsonExtractor jsonExtractor = rootExtractor;
for (String token : tokens.reverse()) {
jsonExtractor = new ObjectFieldJsonExtractor<>(token, jsonExtractor, exceptionOnOutOfBounds);
}
return jsonExtractor;
}
public interface JsonExtractor
{
/**
* Executes the extraction on the existing content of the JsonParser and outputs the match.
*
* Notes:
*
*
JsonParser must be on the FIRST token of the value to be processed when extract is called
*
INVARIANT: when extract() returns, the current token of the parser will be the LAST token of the value
*
*
* @return the value, or null if not applicable
*/
T extract(JsonParser jsonParser)
throws IOException;
}
public static class ObjectFieldJsonExtractor
implements JsonExtractor
{
private final SerializedString fieldName;
private final JsonExtractor delegate;
private final int index;
private final boolean exceptionOnOutOfBounds;
public ObjectFieldJsonExtractor(String fieldName, JsonExtractor delegate)
{
this(fieldName, delegate, false);
}
public ObjectFieldJsonExtractor(String fieldName, JsonExtractor delegate, boolean exceptionOnOutOfBounds)
{
this.fieldName = new SerializedString(requireNonNull(fieldName, "fieldName is null"));
this.delegate = requireNonNull(delegate, "delegate is null");
this.exceptionOnOutOfBounds = exceptionOnOutOfBounds;
this.index = tryParseInt(fieldName, -1);
}
@Override
public T extract(JsonParser jsonParser)
throws IOException
{
if (jsonParser.getCurrentToken() == START_OBJECT) {
return processJsonObject(jsonParser);
}
if (jsonParser.getCurrentToken() == START_ARRAY) {
return processJsonArray(jsonParser);
}
throw new JsonParseException(jsonParser, "Expected a JSON object or array");
}
public T processJsonObject(JsonParser jsonParser)
throws IOException
{
while (!jsonParser.nextFieldName(fieldName)) {
if (!jsonParser.hasCurrentToken()) {
throw new JsonParseException(jsonParser, "Unexpected end of object");
}
if (jsonParser.getCurrentToken() == END_OBJECT) {
// Unable to find matching field
return null;
}
jsonParser.skipChildren(); // Skip nested structure if currently at the start of one
}
jsonParser.nextToken(); // Shift to first token of the value
return delegate.extract(jsonParser);
}
public T processJsonArray(JsonParser jsonParser)
throws IOException
{
int currentIndex = 0;
while (true) {
JsonToken token = jsonParser.nextToken();
if (token == null) {
throw new JsonParseException(jsonParser, "Unexpected end of array");
}
if (token == END_ARRAY) {
// Index out of bounds
if (exceptionOnOutOfBounds) {
throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "Index out of bounds");
}
return null;
}
if (currentIndex == index) {
break;
}
currentIndex++;
jsonParser.skipChildren(); // Skip nested structure if currently at the start of one
}
return delegate.extract(jsonParser);
}
}
public static class ScalarValueJsonExtractor
implements JsonExtractor
{
@Override
public Slice extract(JsonParser jsonParser)
throws IOException
{
JsonToken token = jsonParser.getCurrentToken();
if (token == null) {
throw new JsonParseException(jsonParser, "Unexpected end of value");
}
if (!token.isScalarValue() || token == VALUE_NULL) {
return null;
}
return utf8Slice(jsonParser.getText());
}
}
public static class JsonValueJsonExtractor
implements JsonExtractor
{
@Override
public Slice extract(JsonParser jsonParser)
throws IOException
{
if (!jsonParser.hasCurrentToken()) {
throw new JsonParseException(jsonParser, "Unexpected end of value");
}
DynamicSliceOutput dynamicSliceOutput = new DynamicSliceOutput(ESTIMATED_JSON_OUTPUT_SIZE);
try (JsonGenerator jsonGenerator = createJsonGenerator(JSON_FACTORY, dynamicSliceOutput)) {
jsonGenerator.copyCurrentStructure(jsonParser);
}
return dynamicSliceOutput.slice();
}
}
public static class JsonSizeExtractor
implements JsonExtractor
{
@Override
public Long extract(JsonParser jsonParser)
throws IOException
{
if (!jsonParser.hasCurrentToken()) {
throw new JsonParseException(jsonParser, "Unexpected end of value");
}
if (jsonParser.getCurrentToken() == START_ARRAY) {
long length = 0;
while (true) {
JsonToken token = jsonParser.nextToken();
if (token == null) {
return null;
}
if (token == END_ARRAY) {
return length;
}
jsonParser.skipChildren();
length++;
}
}
if (jsonParser.getCurrentToken() == START_OBJECT) {
long length = 0;
while (true) {
JsonToken token = jsonParser.nextToken();
if (token == null) {
return null;
}
if (token == END_OBJECT) {
return length;
}
if (token == FIELD_NAME) {
length++;
}
else {
jsonParser.skipChildren();
}
}
}
return 0L;
}
}
private static int tryParseInt(String fieldName, int defaultValue)
{
int index = defaultValue;
try {
index = Integer.parseInt(fieldName);
}
catch (NumberFormatException ignored) {
}
return index;
}
}
