io.airlift.slice.DynamicSliceOutput Java Examples

@Test
public void testStackRepresentation()
{
    Block actualBlock = arrayBlockOf(new ArrayType(BIGINT), arrayBlockOf(BIGINT, 1L, 2L), arrayBlockOf(BIGINT, 3L));
    DynamicSliceOutput actualSliceOutput = new DynamicSliceOutput(100);
    writeBlock(functionAssertions.getMetadata().getBlockEncodingSerde(), actualSliceOutput, actualBlock);

    Block expectedBlock = new ArrayType(BIGINT)
            .createBlockBuilder(null, 3)
            .appendStructure(BIGINT.createBlockBuilder(null, 2).writeLong(1).closeEntry().writeLong(2).closeEntry().build())
            .appendStructure(BIGINT.createBlockBuilder(null, 1).writeLong(3).closeEntry().build())
            .build();
    DynamicSliceOutput expectedSliceOutput = new DynamicSliceOutput(100);
    writeBlock(functionAssertions.getMetadata().getBlockEncodingSerde(), expectedSliceOutput, expectedBlock);

    assertEquals(actualSliceOutput.slice(), expectedSliceOutput.slice());
}
 

public static Slice toJson(JsonGeneratorWriter writer, ConnectorSession session, Block block)
{
    try {
        SliceOutput output = new DynamicSliceOutput(40);
        try (JsonGenerator jsonGenerator = createJsonGenerator(JSON_FACTORY, output)) {
            jsonGenerator.writeStartArray();
            for (int i = 0; i < block.getPositionCount(); i++) {
                writer.writeJsonValue(jsonGenerator, block, i, session);
            }
            jsonGenerator.writeEndArray();
        }
        return output.slice();
    }
    catch (IOException e) {
        throwIfUnchecked(e);
        throw new RuntimeException(e);
    }
}
 

@Test
public void testWriteHugeByteChucks()
{
    int size = 1024 * 1024;
    byte[] largeByteArray = new byte[size];
    Arrays.fill(largeByteArray, (byte) 0xA);
    OrcOutputBuffer sliceOutput = new OrcOutputBuffer(CompressionKind.NONE, 256 * 1024);

    DynamicSliceOutput output = new DynamicSliceOutput(size);
    sliceOutput.writeBytes(largeByteArray, 10, size - 10);
    assertEquals(sliceOutput.writeDataTo(output), size - 10);
    assertEquals(output.slice(), wrappedBuffer(largeByteArray, 10, size - 10));

    sliceOutput.reset();
    output.reset();
    sliceOutput.writeBytes(wrappedBuffer(largeByteArray), 100, size - 100);
    assertEquals(sliceOutput.writeDataTo(output), size - 100);
    assertEquals(output.slice(), wrappedBuffer(largeByteArray, 100, size - 100));
}
 

@ScalarOperator(CAST)
@LiteralParameters("x")
@SqlType(JSON)
public static Slice castFromVarchar(@SqlType("varchar(x)") Slice value)
{
    try {
        SliceOutput output = new DynamicSliceOutput(value.length() + 2);
        try (JsonGenerator jsonGenerator = createJsonGenerator(JSON_FACTORY, output)) {
            jsonGenerator.writeString(value.toStringUtf8());
        }
        return output.slice();
    }
    catch (IOException e) {
        throw new PrestoException(INVALID_CAST_ARGUMENT, format("Cannot cast '%s' to %s", value.toStringUtf8(), JSON));
    }
}
 

private static void writePoint(DynamicSliceOutput output, OGCGeometry geometry)
{
    Geometry esriGeometry = geometry.getEsriGeometry();
    verify(esriGeometry instanceof Point, "geometry is expected to be an instance of Point");
    Point point = (Point) esriGeometry;
    verify(!point.hasAttribute(VertexDescription.Semantics.Z) &&
                    !point.hasAttribute(VertexDescription.Semantics.M) &&
                    !point.hasAttribute(VertexDescription.Semantics.ID),
            "Only 2D points with no ID nor M attribute are supported");
    output.appendByte(GeometrySerializationType.POINT.code());
    if (!point.isEmpty()) {
        output.appendDouble(point.getX());
        output.appendDouble(point.getY());
    }
    else {
        output.appendDouble(NaN);
        output.appendDouble(NaN);
    }
}
 

private static void writeGeometryCollection(DynamicSliceOutput output, OGCGeometryCollection collection)
{
    output.appendByte(GeometrySerializationType.GEOMETRY_COLLECTION.code());
    for (int geometryIndex = 0; geometryIndex < collection.numGeometries(); geometryIndex++) {
        OGCGeometry geometry = collection.geometryN(geometryIndex);
        int startPosition = output.size();

        // leave 4 bytes for the shape length
        output.appendInt(0);
        writeGeometry(output, geometry);

        int endPosition = output.size();
        int length = endPosition - startPosition - Integer.BYTES;

        output.getUnderlyingSlice().setInt(startPosition, length);
    }
}
 

@UsedByGeneratedCode
public static Slice toJson(ObjectKeyProvider provider, JsonGeneratorWriter writer, ConnectorSession session, Block block)
{
    try {
        Map<String, Integer> orderedKeyToValuePosition = new TreeMap<>();
        for (int i = 0; i < block.getPositionCount(); i += 2) {
            String objectKey = provider.getObjectKey(block, i);
            orderedKeyToValuePosition.put(objectKey, i + 1);
        }

        SliceOutput output = new DynamicSliceOutput(40);
        try (JsonGenerator jsonGenerator = createJsonGenerator(JSON_FACTORY, output)) {
            jsonGenerator.writeStartObject();
            for (Map.Entry<String, Integer> entry : orderedKeyToValuePosition.entrySet()) {
                jsonGenerator.writeFieldName(entry.getKey());
                writer.writeJsonValue(jsonGenerator, block, entry.getValue(), session);
            }
            jsonGenerator.writeEndObject();
        }
        return output.slice();
    }
    catch (IOException e) {
        throwIfUnchecked(e);
        throw new RuntimeException(e);
    }
}
 

@ScalarFunction
@LiteralParameters("x")
@SqlType(StandardTypes.JSON)
public static Slice jsonParse(@SqlType("varchar(x)") Slice slice)
{
    // cast(json_parse(x) AS t)` will be optimized into `$internal$json_string_to_array/map/row_cast` in ExpressionOptimizer
    // If you make changes to this function (e.g. use parse JSON string into some internal representation),
    // make sure `$internal$json_string_to_array/map/row_cast` is changed accordingly.
    try (JsonParser parser = createJsonParser(JSON_FACTORY, slice)) {
        SliceOutput dynamicSliceOutput = new DynamicSliceOutput(slice.length());
        SORTED_MAPPER.writeValue((OutputStream) dynamicSliceOutput, SORTED_MAPPER.readValue(parser, Object.class));
        // nextToken() returns null if the input is parsed correctly,
        // but will throw an exception if there are trailing characters.
        parser.nextToken();
        return dynamicSliceOutput.slice();
    }
    catch (Exception e) {
        throw new PrestoException(INVALID_FUNCTION_ARGUMENT, format("Cannot convert '%s' to JSON", slice.toStringUtf8()));
    }
}
 

private static void writeGeometryCollection(Geometry collection, DynamicSliceOutput output)
{
    output.appendByte(GeometrySerializationType.GEOMETRY_COLLECTION.code());
    for (int geometryIndex = 0; geometryIndex < collection.getNumGeometries(); geometryIndex++) {
        Geometry geometry = collection.getGeometryN(geometryIndex);
        int startPosition = output.size();

        // leave 4 bytes for the shape length
        output.appendInt(0);
        writeGeometry(geometry, output);

        int endPosition = output.size();
        int length = endPosition - startPosition - Integer.BYTES;

        output.getUnderlyingSlice().setInt(startPosition, length);
    }
}
 

public Slice serialize()
{
    try (SliceOutput output = new DynamicSliceOutput(estimatedSerializedSize())) {
        output.appendByte(UNCOMPRESSED_FORMAT);
        Slice serializedHll = hll.serialize();
        output.appendInt(serializedHll.length());
        output.appendBytes(serializedHll);
        output.appendInt(maxHashes);
        output.appendInt(minhash.size());
        for (long key : minhash.keySet()) {
            output.appendLong(key);
        }
        for (short value : minhash.values()) {
            output.appendShort(value);
        }
        return output.slice();
    }
    catch (IOException e) {
        throw new UncheckedIOException(e);
    }
}
 

private static Block createChannel(int positionCount, int mapSize, Type valueType)
{
    BlockBuilder blockBuilder = JSON.createBlockBuilder(null, positionCount);
    for (int position = 0; position < positionCount; position++) {
        SliceOutput jsonSlice = new DynamicSliceOutput(20 * mapSize);
        jsonSlice.appendByte('{');
        for (int i = 0; i < mapSize; i++) {
            if (i != 0) {
                jsonSlice.appendByte(',');
            }
            String key = "key" + i;
            String value = generateRandomJsonValue(valueType);
            jsonSlice.appendByte('"');
            jsonSlice.appendBytes(key.getBytes(UTF_8));
            jsonSlice.appendBytes("\":".getBytes(UTF_8));
            jsonSlice.appendBytes(value.getBytes(UTF_8));
        }
        jsonSlice.appendByte('}');

        JSON.writeSlice(blockBuilder, jsonSlice.slice());
    }
    return blockBuilder.build();
}
 

private static Slice encodePolyline(MultiVertexGeometry multiVertexGeometry)
{
    long lastLatitude = 0;
    long lastLongitude = 0;

    DynamicSliceOutput output = new DynamicSliceOutput(0);

    for (int i = 0; i < multiVertexGeometry.getPointCount(); i++) {
        Point point = multiVertexGeometry.getPoint(i);

        long latitude = Math.round(point.getY() * 1e5);
        long longitude = Math.round(point.getX() * 1e5);

        long latitudeDelta = latitude - lastLatitude;
        long longitudeDelta = longitude - lastLongitude;

        encode(latitudeDelta, output);
        encode(longitudeDelta, output);

        lastLatitude = latitude;
        lastLongitude = longitude;
    }
    return output.slice();
}
 

private static Block createChannel(int positionCount, int mapSize, Type elementType)
{
    BlockBuilder blockBuilder = JSON.createBlockBuilder(null, positionCount);
    for (int position = 0; position < positionCount; position++) {
        SliceOutput jsonSlice = new DynamicSliceOutput(20 * mapSize);
        jsonSlice.appendByte('[');
        for (int i = 0; i < mapSize; i++) {
            if (i != 0) {
                jsonSlice.appendByte(',');
            }
            String value = generateRandomJsonValue(elementType);
            jsonSlice.appendBytes(value.getBytes(UTF_8));
        }
        jsonSlice.appendByte(']');

        JSON.writeSlice(blockBuilder, jsonSlice.slice());
    }
    return blockBuilder.build();
}
 

@Override
public byte[] getSerializedData()
{
    byte[] classifierBytes = serialize(classifier).getBytes();
    DynamicSliceOutput output = new DynamicSliceOutput(classifierBytes.length + 64 * labelEnumeration.size());
    output.appendInt(classifierBytes.length);
    output.appendBytes(classifierBytes);
    output.appendInt(labelEnumeration.size());
    // Write the enumeration keys
    for (Map.Entry<Integer, String> entry : labelEnumeration.entrySet()) {
        output.appendInt(entry.getKey());
        byte[] bytes = entry.getValue().getBytes(UTF_8);
        output.appendInt(bytes.length);
        output.appendBytes(bytes);
    }

    return output.slice().getBytes();
}
 

@UsedByGeneratedCode
public static Slice toJson(List<JsonGeneratorWriter> fieldWriters, ConnectorSession session, Block block)
{
    try {
        SliceOutput output = new DynamicSliceOutput(40);
        try (JsonGenerator jsonGenerator = createJsonGenerator(JSON_FACTORY, output)) {
            jsonGenerator.writeStartArray();
            for (int i = 0; i < block.getPositionCount(); i++) {
                fieldWriters.get(i).writeJsonValue(jsonGenerator, block, i, session);
            }
            jsonGenerator.writeEndArray();
        }
        return output.slice();
    }
    catch (IOException e) {
        throwIfUnchecked(e);
        throw new RuntimeException(e);
    }
}
 

private static int serializedSize(List<? extends Type> types, Page expectedPage)
{
    PagesSerde serde = new TestingPagesSerdeFactory().createPagesSerde();
    DynamicSliceOutput sliceOutput = new DynamicSliceOutput(1024);
    writePages(serde, sliceOutput, expectedPage);
    Slice slice = sliceOutput.slice();

    Iterator<Page> pageIterator = readPages(serde, slice.getInput());
    if (pageIterator.hasNext()) {
        assertPageEquals(types, pageIterator.next(), expectedPage);
    }
    else {
        assertEquals(expectedPage.getPositionCount(), 0);
    }
    assertFalse(pageIterator.hasNext());

    return slice.length();
}
 

@Test
public void testRoundTrip()
{
    PagesSerde serde = new TestingPagesSerdeFactory().createPagesSerde();
    BlockBuilder expectedBlockBuilder = VARCHAR.createBlockBuilder(null, 5);
    VARCHAR.writeString(expectedBlockBuilder, "alice");
    VARCHAR.writeString(expectedBlockBuilder, "bob");
    VARCHAR.writeString(expectedBlockBuilder, "charlie");
    VARCHAR.writeString(expectedBlockBuilder, "dave");
    Block expectedBlock = expectedBlockBuilder.build();

    Page expectedPage = new Page(expectedBlock, expectedBlock, expectedBlock);

    DynamicSliceOutput sliceOutput = new DynamicSliceOutput(1024);
    writePages(serde, sliceOutput, expectedPage, expectedPage, expectedPage);

    List<Type> types = ImmutableList.of(VARCHAR, VARCHAR, VARCHAR);
    Iterator<Page> pageIterator = readPages(serde, sliceOutput.slice().getInput());
    assertPageEquals(types, pageIterator.next(), expectedPage);
    assertPageEquals(types, pageIterator.next(), expectedPage);
    assertPageEquals(types, pageIterator.next(), expectedPage);
    assertFalse(pageIterator.hasNext());
}
 

/**
 * Serialize JTS {@link Geometry} shape into an ESRI shape
 */
public static Slice serialize(Geometry geometry)
{
    requireNonNull(geometry, "input is null");
    DynamicSliceOutput output = new DynamicSliceOutput(100);
    writeGeometry(geometry, output);
    return output.slice();
}
 

private static void writeSimpleGeometry(DynamicSliceOutput output, GeometrySerializationType type, OGCGeometry geometry)
{
    output.appendByte(type.code());
    Geometry esriGeometry = requireNonNull(geometry.getEsriGeometry(), "esriGeometry is null");
    byte[] shape = geometryToEsriShape(esriGeometry);
    output.appendBytes(shape);
}
 

@Override
public void serialize(Block block, JsonGenerator jsonGenerator, SerializerProvider serializerProvider)
        throws IOException
{
    SliceOutput output = new DynamicSliceOutput(64);
    blockEncodingSerde.writeBlock(output, block);
    String encoded = Base64.getEncoder().encodeToString(output.slice().getBytes());
    jsonGenerator.writeString(encoded);
}
 

@Test
public void testRoundTrip()
{
    BlockBuilder expectedBlockBuilder = VARCHAR.createBlockBuilder(null, 4);
    VARCHAR.writeString(expectedBlockBuilder, "alice");
    VARCHAR.writeString(expectedBlockBuilder, "bob");
    VARCHAR.writeString(expectedBlockBuilder, "charlie");
    VARCHAR.writeString(expectedBlockBuilder, "dave");
    Block expectedBlock = expectedBlockBuilder.build();

    DynamicSliceOutput sliceOutput = new DynamicSliceOutput(1024);
    blockEncodingSerde.writeBlock(sliceOutput, expectedBlock);
    Block actualBlock = blockEncodingSerde.readBlock(sliceOutput.slice().getInput());
    assertBlockEquals(VARCHAR, actualBlock, expectedBlock);
}
 

private static void writeEnvelopeCoordinates(DynamicSliceOutput output, Envelope envelope)
{
    if (envelope.isEmpty()) {
        output.appendDouble(NaN);
        output.appendDouble(NaN);
        output.appendDouble(NaN);
        output.appendDouble(NaN);
    }
    else {
        output.appendDouble(envelope.getXMin());
        output.appendDouble(envelope.getYMin());
        output.appendDouble(envelope.getXMax());
        output.appendDouble(envelope.getYMax());
    }
}
 

private static Block createVariableWidthBlock(int entries)
{
    int[] offsets = new int[entries + 1];
    DynamicSliceOutput dynamicSliceOutput = new DynamicSliceOutput(entries);
    for (int i = 0; i < entries; i++) {
        dynamicSliceOutput.writeByte(i);
        offsets[i + 1] = dynamicSliceOutput.size();
    }
    return new VariableWidthBlock(entries, dynamicSliceOutput.slice(), offsets, Optional.empty());
}
 

@Test
public void testRoundTrip()
{
    int positionCount = 40;

    // build dictionary
    BlockBuilder dictionaryBuilder = VARCHAR.createBlockBuilder(null, 4);
    VARCHAR.writeString(dictionaryBuilder, "alice");
    VARCHAR.writeString(dictionaryBuilder, "bob");
    VARCHAR.writeString(dictionaryBuilder, "charlie");
    VARCHAR.writeString(dictionaryBuilder, "dave");
    Block dictionary = dictionaryBuilder.build();

    // build ids
    int[] ids = new int[positionCount];
    for (int i = 0; i < 40; i++) {
        ids[i] = i % 4;
    }

    DictionaryBlock dictionaryBlock = new DictionaryBlock(dictionary, ids);

    DynamicSliceOutput sliceOutput = new DynamicSliceOutput(1024);
    blockEncodingSerde.writeBlock(sliceOutput, dictionaryBlock);
    Block actualBlock = blockEncodingSerde.readBlock(sliceOutput.slice().getInput());

    assertTrue(actualBlock instanceof DictionaryBlock);
    DictionaryBlock actualDictionaryBlock = (DictionaryBlock) actualBlock;
    assertBlockEquals(VARCHAR, actualDictionaryBlock.getDictionary(), dictionary);
    for (int position = 0; position < actualDictionaryBlock.getPositionCount(); position++) {
        assertEquals(actualDictionaryBlock.getId(position), ids[position]);
    }
    assertEquals(actualDictionaryBlock.getDictionarySourceId(), dictionaryBlock.getDictionarySourceId());
}
 

private void initializeCapacity()
{
    if (positions != 0 || currentEntrySize != 0) {
        throw new IllegalStateException(getClass().getSimpleName() + " was used before initialization");
    }
    initialized = true;
    valueIsNull = new boolean[initialEntryCount];
    offsets = new int[initialEntryCount + 1];
    sliceOutput = new DynamicSliceOutput(initialSliceOutputSize);
    updateArraysDataSize();
}
 

private Slice blockToSlice(Block block)
{
    // This function is strictly for testing use only
    SliceOutput sliceOutput = new DynamicSliceOutput(1000);
    BlockSerdeUtil.writeBlock(blockEncodingSerde, sliceOutput, block);
    return sliceOutput.slice();
}
 

static Slice getFooter(List<RowGroup> rowGroups, MessageType messageType)
        throws IOException
{
    FileMetaData fileMetaData = new FileMetaData();
    fileMetaData.setVersion(1);
    fileMetaData.setSchema(MessageTypeConverter.toParquetSchema(messageType));
    long totalRows = rowGroups.stream().mapToLong(RowGroup::getNum_rows).sum();
    fileMetaData.setNum_rows(totalRows);
    fileMetaData.setRow_groups(ImmutableList.copyOf(rowGroups));

    DynamicSliceOutput dynamicSliceOutput = new DynamicSliceOutput(40);
    Util.writeFileMetaData(fileMetaData, dynamicSliceOutput);
    return dynamicSliceOutput.slice();
}
 

static RE2 compileImpl(String expr, int mode, MatchKind matchKind, Options options)
        throws com.github.aaronshan.functions.regexp.re2j.PatternSyntaxException {
    Regexp re = com.github.aaronshan.functions.regexp.re2j.Parser.parse(expr, mode);
    int maxCap = re.maxCap();  // (may shrink during simplify)
    re = Simplify.simplify(re);
    com.github.aaronshan.functions.regexp.re2j.Prog prog = com.github.aaronshan.functions.regexp.re2j.Compiler.compileRegexp(re, false);
    com.github.aaronshan.functions.regexp.re2j.Prog reverseProg = com.github.aaronshan.functions.regexp.re2j.Compiler.compileRegexp(re, true);
    SliceOutput prefixBuilder = new DynamicSliceOutput(prog.numInst());
    boolean prefixComplete = prog.prefix(prefixBuilder);
    Slice prefixUTF8 = prefixBuilder.slice();
    return new RE2(expr, prog, reverseProg, maxCap, re.namedGroupIndexes(), matchKind, options, prefixComplete, prefixUTF8);
}
 

/**
 * Helper: replaceAll/replaceFirst hybrid.
 */
private Slice replace(Slice replacement, boolean all) {
    reset();
    SliceOutput so = new DynamicSliceOutput(input.length());
    while (find()) {
        appendReplacement(so, replacement);
        if (!all) {
            break;
        }
    }
    appendTail(so);
    return so.slice();
}
 

@Test
public void testNoStartSync()
        throws Exception
{
    SliceOutput output = new DynamicSliceOutput(10 * 1024);

    List<Segment> segments = ImmutableList.of(
            writeSegment(output, ImmutableList.of(ImmutableList.of(0, 2, 3, 4), ImmutableList.of(10, 12, 13))),
            writeSegment(output, ImmutableList.of(ImmutableList.of(20, 22), ImmutableList.of(30, 33), ImmutableList.of(40, 44))),
            writeSegment(output, ImmutableList.of(ImmutableList.of(100, 101, 102))));

    assertFileSegments(output.slice(), segments);
}