Examples with SqlNullable io.trino.spi.function.SqlNullable used on opensource projects

Example 31 with SqlNullable

use of io.trino.spi.function.SqlNullable in project trino by trinodb.

the class JoniRegexpReplaceLambdaFunction method regexpReplace.

@LiteralParameters("x")
@SqlType("varchar")
@SqlNullable
public Slice regexpReplace(@SqlType("varchar") Slice source, @SqlType(JoniRegexpType.NAME) JoniRegexp pattern, @SqlType("function(array(varchar), varchar(x))") UnaryFunctionInterface replaceFunction) {
    // If there is no match we can simply return the original source without doing copy.
    Matcher matcher = pattern.matcher(source.getBytes());
    if (matcher.search(0, source.length(), Option.DEFAULT) == -1) {
        return source;
    }
    SliceOutput output = new DynamicSliceOutput(source.length());
    // that will be passed to the lambda function.
    if (pageBuilder.isFull()) {
        pageBuilder.reset();
    }
    BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0);
    int groupCount = pattern.regex().numberOfCaptures();
    int appendPosition = 0;
    int nextStart;
    do {
        // nextStart is the same as the last appendPosition, unless the last match was zero-width.
        if (matcher.getEnd() == matcher.getBegin()) {
            if (matcher.getBegin() < source.length()) {
                nextStart = matcher.getEnd() + lengthOfCodePointFromStartByte(source.getByte(matcher.getBegin()));
            } else {
                // last match is empty and we matched end of source, move past the source length to terminate the loop
                nextStart = matcher.getEnd() + 1;
            }
        } else {
            nextStart = matcher.getEnd();
        }
        // Append the un-matched part
        Slice unmatched = source.slice(appendPosition, matcher.getBegin() - appendPosition);
        appendPosition = matcher.getEnd();
        output.appendBytes(unmatched);
        // Append the capturing groups to the target block that will be passed to lambda
        Region matchedRegion = matcher.getEagerRegion();
        for (int i = 1; i <= groupCount; i++) {
            // Add to the block builder if the matched region is not null. In Joni null is represented as [-1, -1]
            if (matchedRegion.beg[i] >= 0 && matchedRegion.end[i] >= 0) {
                VARCHAR.writeSlice(blockBuilder, source, matchedRegion.beg[i], matchedRegion.end[i] - matchedRegion.beg[i]);
            } else {
                blockBuilder.appendNull();
            }
        }
        pageBuilder.declarePositions(groupCount);
        Block target = blockBuilder.getRegion(blockBuilder.getPositionCount() - groupCount, groupCount);
        // Call the lambda function to replace the block, and append the result to output
        Slice replaced = (Slice) replaceFunction.apply(target);
        if (replaced == null) {
            // replacing a substring with null (unknown) makes the entire string null
            return null;
        }
        output.appendBytes(replaced);
    } while (matcher.search(nextStart, source.length(), Option.DEFAULT) != -1);
    // Append the last un-matched part
    output.writeBytes(source, appendPosition, source.length() - appendPosition);
    return output.slice();
}

Example 32 with SqlNullable

use of io.trino.spi.function.SqlNullable in project trino by trinodb.

the class JsonFunctions method jsonArrayGet.

@SqlNullable
@ScalarFunction
@SqlType(StandardTypes.JSON)
public static Slice jsonArrayGet(@SqlType(StandardTypes.JSON) Slice json, @SqlType(StandardTypes.BIGINT) long index) {
    // this value cannot be converted to positive number
    if (index == Long.MIN_VALUE) {
        return null;
    }
    try (JsonParser parser = createJsonParser(MAPPING_JSON_FACTORY, json)) {
        if (parser.nextToken() != START_ARRAY) {
            return null;
        }
        List<String> tokens = null;
        if (index < 0) {
            tokens = new LinkedList<>();
        }
        long count = 0;
        while (true) {
            JsonToken token = parser.nextToken();
            if (token == null) {
                return null;
            }
            if (token == END_ARRAY) {
                if (tokens != null && count >= index * -1) {
                    return utf8Slice(tokens.get(0));
                }
                return null;
            }
            String arrayElement;
            if (token == START_OBJECT || token == START_ARRAY) {
                arrayElement = parser.readValueAsTree().toString();
            } else {
                arrayElement = parser.getValueAsString();
            }
            if (count == index) {
                return arrayElement == null ? null : utf8Slice(arrayElement);
            }
            if (tokens != null) {
                tokens.add(arrayElement);
                if (count >= index * -1) {
                    tokens.remove(0);
                }
            }
            count++;
        }
    } catch (IOException e) {
        return null;
    }
}

Example 33 with SqlNullable

use of io.trino.spi.function.SqlNullable in project trino by trinodb.

the class JsonOperators method castToVarchar.

@ScalarOperator(CAST)
@SqlNullable
@LiteralParameters("x")
@SqlType("varchar(x)")
public static Slice castToVarchar(@SqlType(JSON) Slice json) {
    try (JsonParser parser = createJsonParser(JSON_FACTORY, json)) {
        parser.nextToken();
        Slice result = currentTokenAsVarchar(parser);
        // check no trailing token
        checkCondition(parser.nextToken() == null, INVALID_CAST_ARGUMENT, "Cannot cast input json to VARCHAR");
        return result;
    } catch (IOException | JsonCastException e) {
        throw new TrinoException(INVALID_CAST_ARGUMENT, format("Cannot cast '%s' to %s", json.toStringUtf8(), VARCHAR), e);
    }
}

Example 34 with SqlNullable

use of io.trino.spi.function.SqlNullable in project trino by trinodb.

the class JsonOperators method castToDouble.

@ScalarOperator(CAST)
@SqlNullable
@SqlType(DOUBLE)
public static Double castToDouble(@SqlType(JSON) Slice json) {
    try (JsonParser parser = createJsonParser(JSON_FACTORY, json)) {
        parser.nextToken();
        Double result = currentTokenAsDouble(parser);
        // check no trailing token
        checkCondition(parser.nextToken() == null, INVALID_CAST_ARGUMENT, "Cannot cast input json to DOUBLE");
        return result;
    } catch (IOException | JsonCastException e) {
        throw new TrinoException(INVALID_CAST_ARGUMENT, format("Cannot cast '%s' to %s", json.toStringUtf8(), DOUBLE), e);
    }
}

Example 35 with SqlNullable

use of io.trino.spi.function.SqlNullable in project trino by trinodb.

the class JsonOperators method castToReal.

@ScalarOperator(CAST)
@SqlNullable
@SqlType(REAL)
public static Long castToReal(@SqlType(JSON) Slice json) {
    try (JsonParser parser = createJsonParser(JSON_FACTORY, json)) {
        parser.nextToken();
        Long result = currentTokenAsReal(parser);
        // check no trailing token
        checkCondition(parser.nextToken() == null, INVALID_CAST_ARGUMENT, "Cannot cast input json to REAL");
        return result;
    } catch (IOException | JsonCastException e) {
        throw new TrinoException(INVALID_CAST_ARGUMENT, format("Cannot cast '%s' to %s", json.toStringUtf8(), REAL), e);
    }
}