Search in sources :

Example 11 with ByteArrayInputStreamWithPos

use of org.apache.flink.core.memory.ByteArrayInputStreamWithPos in project flink by apache.

the class RocksDBFoldingState method get.

@Override
public ACC get() {
    try {
        writeCurrentKeyWithGroupAndNamespace();
        byte[] key = keySerializationStream.toByteArray();
        byte[] valueBytes = backend.db.get(columnFamily, key);
        if (valueBytes == null) {
            return null;
        }
        return valueSerializer.deserialize(new DataInputViewStreamWrapper(new ByteArrayInputStreamWithPos(valueBytes)));
    } catch (IOException | RocksDBException e) {
        throw new RuntimeException("Error while retrieving data from RocksDB", e);
    }
}
Also used : RocksDBException(org.rocksdb.RocksDBException) ByteArrayInputStreamWithPos(org.apache.flink.core.memory.ByteArrayInputStreamWithPos) IOException(java.io.IOException) DataInputViewStreamWrapper(org.apache.flink.core.memory.DataInputViewStreamWrapper)

Example 12 with ByteArrayInputStreamWithPos

use of org.apache.flink.core.memory.ByteArrayInputStreamWithPos in project flink by apache.

the class RocksDBFoldingState method add.

@Override
public void add(T value) throws IOException {
    try {
        writeCurrentKeyWithGroupAndNamespace();
        byte[] key = keySerializationStream.toByteArray();
        byte[] valueBytes = backend.db.get(columnFamily, key);
        DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(keySerializationStream);
        if (valueBytes == null) {
            keySerializationStream.reset();
            valueSerializer.serialize(foldFunction.fold(stateDesc.getDefaultValue(), value), out);
            backend.db.put(columnFamily, writeOptions, key, keySerializationStream.toByteArray());
        } else {
            ACC oldValue = valueSerializer.deserialize(new DataInputViewStreamWrapper(new ByteArrayInputStreamWithPos(valueBytes)));
            ACC newValue = foldFunction.fold(oldValue, value);
            keySerializationStream.reset();
            valueSerializer.serialize(newValue, out);
            backend.db.put(columnFamily, writeOptions, key, keySerializationStream.toByteArray());
        }
    } catch (Exception e) {
        throw new RuntimeException("Error while adding data to RocksDB", e);
    }
}
Also used : DataOutputViewStreamWrapper(org.apache.flink.core.memory.DataOutputViewStreamWrapper) ByteArrayInputStreamWithPos(org.apache.flink.core.memory.ByteArrayInputStreamWithPos) DataInputViewStreamWrapper(org.apache.flink.core.memory.DataInputViewStreamWrapper) RocksDBException(org.rocksdb.RocksDBException) IOException(java.io.IOException)

Example 13 with ByteArrayInputStreamWithPos

use of org.apache.flink.core.memory.ByteArrayInputStreamWithPos in project flink by apache.

the class RocksDBMapState method deserializeUserKey.

private UK deserializeUserKey(byte[] rawKeyBytes) throws IOException {
    ByteArrayInputStreamWithPos bais = new ByteArrayInputStreamWithPos(rawKeyBytes);
    DataInputViewStreamWrapper in = new DataInputViewStreamWrapper(bais);
    readKeyWithGroupAndNamespace(bais, in);
    return userKeySerializer.deserialize(in);
}
Also used : ByteArrayInputStreamWithPos(org.apache.flink.core.memory.ByteArrayInputStreamWithPos) DataInputViewStreamWrapper(org.apache.flink.core.memory.DataInputViewStreamWrapper)

Example 14 with ByteArrayInputStreamWithPos

use of org.apache.flink.core.memory.ByteArrayInputStreamWithPos in project flink by apache.

the class RocksDBReducingState method mergeNamespaces.

@Override
public void mergeNamespaces(N target, Collection<N> sources) throws Exception {
    if (sources == null || sources.isEmpty()) {
        return;
    }
    // cache key and namespace
    final K key = backend.getCurrentKey();
    final int keyGroup = backend.getCurrentKeyGroupIndex();
    try {
        V current = null;
        // merge the sources to the target
        for (N source : sources) {
            if (source != null) {
                writeKeyWithGroupAndNamespace(keyGroup, key, source, keySerializationStream, keySerializationDataOutputView);
                final byte[] sourceKey = keySerializationStream.toByteArray();
                final byte[] valueBytes = backend.db.get(columnFamily, sourceKey);
                if (valueBytes != null) {
                    V value = valueSerializer.deserialize(new DataInputViewStreamWrapper(new ByteArrayInputStreamWithPos(valueBytes)));
                    if (current != null) {
                        current = reduceFunction.reduce(current, value);
                    } else {
                        current = value;
                    }
                }
            }
        }
        // if something came out of merging the sources, merge it or write it to the target
        if (current != null) {
            // create the target full-binary-key 
            writeKeyWithGroupAndNamespace(keyGroup, key, target, keySerializationStream, keySerializationDataOutputView);
            final byte[] targetKey = keySerializationStream.toByteArray();
            final byte[] targetValueBytes = backend.db.get(columnFamily, targetKey);
            if (targetValueBytes != null) {
                // target also had a value, merge
                V value = valueSerializer.deserialize(new DataInputViewStreamWrapper(new ByteArrayInputStreamWithPos(targetValueBytes)));
                current = reduceFunction.reduce(current, value);
            }
            // serialize the resulting value
            keySerializationStream.reset();
            valueSerializer.serialize(current, keySerializationDataOutputView);
            // write the resulting value
            backend.db.put(columnFamily, writeOptions, targetKey, keySerializationStream.toByteArray());
        }
    } catch (Exception e) {
        throw new Exception("Error while merging state in RocksDB", e);
    }
}
Also used : ByteArrayInputStreamWithPos(org.apache.flink.core.memory.ByteArrayInputStreamWithPos) DataInputViewStreamWrapper(org.apache.flink.core.memory.DataInputViewStreamWrapper) IOException(java.io.IOException) RocksDBException(org.rocksdb.RocksDBException)

Example 15 with ByteArrayInputStreamWithPos

use of org.apache.flink.core.memory.ByteArrayInputStreamWithPos in project flink by apache.

the class VersionedIOWriteableTest method testReadCompatibleVersion.

@Test
public void testReadCompatibleVersion() throws Exception {
    String payload = "test";
    TestWriteable testWriteable = new TestWriteable(1, payload);
    byte[] serialized;
    try (ByteArrayOutputStreamWithPos out = new ByteArrayOutputStreamWithPos()) {
        testWriteable.write(new DataOutputViewStreamWrapper(out));
        serialized = out.toByteArray();
    }
    testWriteable = new TestWriteable(2) {

        @Override
        public boolean isCompatibleVersion(int version) {
            return getVersion() >= version;
        }
    };
    try (ByteArrayInputStreamWithPos in = new ByteArrayInputStreamWithPos(serialized)) {
        testWriteable.read(new DataInputViewStreamWrapper(in));
    }
    Assert.assertEquals(payload, testWriteable.getData());
}
Also used : DataOutputViewStreamWrapper(org.apache.flink.core.memory.DataOutputViewStreamWrapper) ByteArrayInputStreamWithPos(org.apache.flink.core.memory.ByteArrayInputStreamWithPos) ByteArrayOutputStreamWithPos(org.apache.flink.core.memory.ByteArrayOutputStreamWithPos) DataInputViewStreamWrapper(org.apache.flink.core.memory.DataInputViewStreamWrapper) Test(org.junit.Test)

Aggregations

ByteArrayInputStreamWithPos (org.apache.flink.core.memory.ByteArrayInputStreamWithPos)19 DataInputViewStreamWrapper (org.apache.flink.core.memory.DataInputViewStreamWrapper)19 DataOutputViewStreamWrapper (org.apache.flink.core.memory.DataOutputViewStreamWrapper)12 ByteArrayOutputStreamWithPos (org.apache.flink.core.memory.ByteArrayOutputStreamWithPos)10 Test (org.junit.Test)9 IOException (java.io.IOException)7 RocksDBException (org.rocksdb.RocksDBException)6 ArrayList (java.util.ArrayList)2 EOFException (java.io.EOFException)1 ObjectInputStream (java.io.ObjectInputStream)1 URLClassLoader (java.net.URLClassLoader)1 HashMap (java.util.HashMap)1 AggregatingStateDescriptor (org.apache.flink.api.common.state.AggregatingStateDescriptor)1 FoldingStateDescriptor (org.apache.flink.api.common.state.FoldingStateDescriptor)1 ListStateDescriptor (org.apache.flink.api.common.state.ListStateDescriptor)1 MapStateDescriptor (org.apache.flink.api.common.state.MapStateDescriptor)1 ReducingStateDescriptor (org.apache.flink.api.common.state.ReducingStateDescriptor)1 StateDescriptor (org.apache.flink.api.common.state.StateDescriptor)1 ValueStateDescriptor (org.apache.flink.api.common.state.ValueStateDescriptor)1 DataInputViewStream (org.apache.flink.api.java.typeutils.runtime.DataInputViewStream)1