private void rewriteToBlockSize(int targetBlockBits) { assert targetBlockBits <= maxBitsPerBlock; // We copy over data blocks to an output with one-larger block bit size. // We also discard references to blocks as we're copying to allow GC to // clean up partial results in case of memory pressure. ByteBuffersDataOutput cloned = new ByteBuffersDataOutput(targetBlockBits, targetBlockBits, blockAllocate, NO_REUSE); ByteBuffer block; while ((block = blocks.pollFirst()) != null) { block.flip(); cloned.writeBytes(block); if (blockReuse != NO_REUSE) { blockReuse.accept(block); } } assert blocks.isEmpty(); this.blockBits = targetBlockBits; blocks.addAll(cloned.blocks); }
@Override public void writeString(String v) { try { final int MAX_CHARS_PER_WINDOW = 1024; if (v.length() <= MAX_CHARS_PER_WINDOW) { final BytesRef utf8 = new BytesRef(v); writeVInt(utf8.length); writeBytes(utf8.bytes, utf8.offset, utf8.length); } else { writeVInt(UnicodeUtil.calcUTF16toUTF8Length(v, 0, v.length())); final byte [] buf = new byte [UnicodeUtil.MAX_UTF8_BYTES_PER_CHAR * MAX_CHARS_PER_WINDOW]; UTF16toUTF8(v, 0, v.length(), buf, (len) -> { writeBytes(buf, 0, len); }); } } catch (IOException e) { throw new UncheckedIOException(e); } }
@Override public void writeByte(byte b) { if (!currentBlock.hasRemaining()) { appendBlock(); } currentBlock.put(b); }
/** * Return a contiguous array with the current content written to the output. The returned * array is always a copy (can be mutated). */ public byte[] toArrayCopy() { if (blocks.size() == 0) { return EMPTY_BYTE_ARRAY; } // We could try to detect single-block, array-based ByteBuffer here // and use Arrays.copyOfRange, but I don't think it's worth the extra // instance checks. byte [] arr = new byte[Math.toIntExact(size())]; int offset = 0; for (ByteBuffer bb : toBufferList()) { int len = bb.remaining(); bb.get(arr, offset, len); offset += len; } return arr; }
/** * Return a {@link ByteBuffersDataInput} for the set of current buffers ({@link #toBufferList()}). */ public ByteBuffersDataInput toDataInput() { return new ByteBuffersDataInput(toBufferList()); }
@Override public long getFilePointer() { ensureOpen(); return delegate.size(); }
@Override public void writeByte(byte b) throws IOException { ensureOpen(); delegate.writeByte(b); }
private void appendBlock() { if (blocks.size() >= MAX_BLOCKS_BEFORE_BLOCK_EXPANSION && blockBits < maxBitsPerBlock) { rewriteToBlockSize(blockBits + 1); if (blocks.getLast().hasRemaining()) { return; } } final int requiredBlockSize = 1 << blockBits; currentBlock = blockAllocate.apply(requiredBlockSize); assert currentBlock.capacity() == requiredBlockSize; blocks.add(currentBlock); }
@Override public void copyBytes(DataInput input, long numBytes) throws IOException { ensureOpen(); delegate.copyBytes(input, numBytes); }
/** * @return The number of bytes written to this output so far. */ public long size() { long size = 0; int blockCount = blocks.size(); if (blockCount >= 1) { int fullBlockSize = (blockCount - 1) * blockSize(); int lastBlockSize = blocks.getLast().position(); size = fullBlockSize + lastBlockSize; } return size; }
/** * @return Returns a new {@link ByteBuffersDataOutput} with the {@link #reset()} capability. */ // TODO: perhaps we can move it out to an utility class (as a supplier of preconfigured instances?) public static ByteBuffersDataOutput newResettableInstance() { ByteBuffersDataOutput.ByteBufferRecycler reuser = new ByteBuffersDataOutput.ByteBufferRecycler( ByteBuffersDataOutput.ALLOCATE_BB_ON_HEAP); return new ByteBuffersDataOutput( ByteBuffersDataOutput.DEFAULT_MIN_BITS_PER_BLOCK, ByteBuffersDataOutput.DEFAULT_MAX_BITS_PER_BLOCK, reuser::allocate, reuser::reuse); }
public ByteBuffersDataOutput(long expectedSize) { this(computeBlockSizeBitsFor(expectedSize), DEFAULT_MAX_BITS_PER_BLOCK, ALLOCATE_BB_ON_HEAP, NO_REUSE); }
if (lastChecksumPosition != delegate.size()) { lastChecksumPosition = delegate.size(); checksum.reset(); byte [] buffer = null; for (ByteBuffer bb : delegate.toBufferList()) { if (bb.hasArray()) { checksum.update(bb.array(), bb.arrayOffset() + bb.position(), bb.remaining());
/** * Copy the current content of this object into another {@link DataOutput}. */ public void copyTo(DataOutput output) throws IOException { for (ByteBuffer bb : toBufferList()) { if (bb.hasArray()) { output.writeBytes(bb.array(), bb.arrayOffset() + bb.position(), bb.remaining()); } else { output.copyBytes(new ByteBuffersDataInput(Arrays.asList(bb)), bb.remaining()); } } }
final IndexOutput createOutput(BiFunction<String, ByteBuffersDataOutput, IndexInput> outputToInput) throws IOException { if (content != null) { throw new IOException("Can only write to a file once: " + fileName); } String clazzName = ByteBuffersDirectory.class.getSimpleName(); String outputName = String.format(Locale.ROOT, "%s output (file=%s)", clazzName, fileName); return new ByteBuffersIndexOutput( bbOutputSupplier.get(), outputName, fileName, new CRC32(), (output) -> { content = outputToInput.apply(fileName, output); cachedLength = output.size(); }); } }
@Override public void writeByte(byte b) throws IOException { ensureOpen(); delegate.writeByte(b); }
private void appendBlock() { if (blocks.size() >= MAX_BLOCKS_BEFORE_BLOCK_EXPANSION && blockBits < maxBitsPerBlock) { rewriteToBlockSize(blockBits + 1); if (blocks.getLast().hasRemaining()) { return; } } final int requiredBlockSize = 1 << blockBits; currentBlock = blockAllocate.apply(requiredBlockSize); assert currentBlock.capacity() == requiredBlockSize; blocks.add(currentBlock); }
@Override public void copyBytes(DataInput input, long numBytes) throws IOException { ensureOpen(); delegate.copyBytes(input, numBytes); }