Java Code Examples for org.apache.spark.unsafe.Platform#copyMemory()

/**
 * UnsafeArrayWriter doesn't have a binary form that lets the user pass an
 * offset and length, so I've added one here. It is the minor tweak of the
 * UnsafeArrayWriter.write(int, byte[]) method.
 * @param holder the BufferHolder where the bytes are being written
 * @param writer the UnsafeArrayWriter
 * @param ordinal the element that we are writing into
 * @param input the input bytes
 * @param offset the first byte from input
 * @param length the number of bytes to write
 */
static void write(BufferHolder holder, UnsafeArrayWriter writer, int ordinal,
                  byte[] input, int offset, int length) {
  final int roundedSize = ByteArrayMethods.roundNumberOfBytesToNearestWord(length);

  // grow the global buffer before writing data.
  holder.grow(roundedSize);

  if ((length & 0x07) > 0) {
    Platform.putLong(holder.buffer, holder.cursor + ((length >> 3) << 3), 0L);
  }

  // Write the bytes to the variable length portion.
  Platform.copyMemory(input, Platform.BYTE_ARRAY_OFFSET + offset,
      holder.buffer, holder.cursor, length);

  writer.setOffsetAndSize(ordinal, holder.cursor, length);

  // move the cursor forward.
  holder.cursor += roundedSize;
}
 

/**
 * Write a record to the sorter.
 */
public void insertRecord(Object recordBase, long recordOffset, int length, long prefix)
        throws IOException {

    growPointerArrayIfNecessary();
    // Need 4 bytes to store the record length.
    final int required = length + 4;
    acquireNewPageIfNecessary(required);

    final Object base = currentPage.getBaseObject();
    final long recordAddress = taskMemoryManager.encodePageNumberAndOffset(currentPage, pageCursor);
    Platform.putInt(base, pageCursor, length);
    pageCursor += 4;
    Platform.copyMemory(recordBase, recordOffset, base, pageCursor, length);
    pageCursor += length;
    assert (inMemSorter != null);
    inMemSorter.insertRecord(recordAddress, prefix);
}
 

/**
 * Write a key-value record to the sorter. The key and value will be put together in-memory,
 * using the following format:
 *
 * record length (4 bytes), key length (4 bytes), key data, value data
 *
 * record length = key length + value length + 4
 */
public void insertKVRecord(Object keyBase, long keyOffset, int keyLen,
                           Object valueBase, long valueOffset, int valueLen, long prefix)
        throws IOException {

    growPointerArrayIfNecessary();
    final int required = keyLen + valueLen + 4 + 4;
    acquireNewPageIfNecessary(required);

    final Object base = currentPage.getBaseObject();
    final long recordAddress = taskMemoryManager.encodePageNumberAndOffset(currentPage, pageCursor);
    Platform.putInt(base, pageCursor, keyLen + valueLen + 4);
    pageCursor += 4;
    Platform.putInt(base, pageCursor, keyLen);
    pageCursor += 4;
    Platform.copyMemory(keyBase, keyOffset, base, pageCursor, keyLen);
    pageCursor += keyLen;
    Platform.copyMemory(valueBase, valueOffset, base, pageCursor, valueLen);
    pageCursor += valueLen;

    assert (inMemSorter != null);
    inMemSorter.insertRecord(recordAddress, prefix);
}
 

/**
 * Copies this row, returning a self-contained UnsafeRow that stores its data in an internal
 * byte array rather than referencing data stored in a data page.
 */
@Override
public UnsafeRow copy() {
    UnsafeRow rowCopy = new UnsafeRow(numFields);
    final byte[] rowDataCopy = new byte[sizeInBytes];
    Platform.copyMemory(
            baseObject,
            baseOffset,
            rowDataCopy,
            Platform.BYTE_ARRAY_OFFSET,
            sizeInBytes
    );
    rowCopy.pointTo(rowDataCopy, Platform.BYTE_ARRAY_OFFSET, sizeInBytes);
    if (stringValues != null) {
        rowCopy.stringValues = new ArrayList<>(stringValues);
    }
    return rowCopy;
}
 

/**
 * Write this UnsafeRow's underlying bytes to the given OutputStream.
 *
 * @param out         the stream to write to.
 * @param writeBuffer a byte array for buffering chunks of off-heap data while writing to the
 *                    output stream. If this row is backed by an on-heap byte array, then this
 *                    buffer will not be used and may be null.
 */
public void writeToStream(OutputStream out, byte[] writeBuffer) throws IOException {
    if (baseObject instanceof byte[]) {
        int offsetInByteArray = (int) (Platform.BYTE_ARRAY_OFFSET - baseOffset);
        out.write((byte[]) baseObject, offsetInByteArray, sizeInBytes);
    } else {
        int dataRemaining = sizeInBytes;
        long rowReadPosition = baseOffset;
        while (dataRemaining > 0) {
            int toTransfer = Math.min(writeBuffer.length, dataRemaining);
            Platform.copyMemory(
                    baseObject, rowReadPosition, writeBuffer, Platform.BYTE_ARRAY_OFFSET, toTransfer);
            out.write(writeBuffer, 0, toTransfer);
            rowReadPosition += toTransfer;
            dataRemaining -= toTransfer;
        }
    }
}
 

/**
 * Write the bytes of var-length field into ByteBuffer
 *
 * Note: only work with HeapByteBuffer
 */
public void writeFieldTo(int ordinal, ByteBuffer buffer) {
    final long offsetAndSize = getLong(ordinal);
    final int offset = (int) (offsetAndSize >> 32);
    final int size = (int) offsetAndSize;

    buffer.putInt(size);
    int pos = buffer.position();
    buffer.position(pos + size);
    Platform.copyMemory(
            baseObject,
            baseOffset + offset,
            buffer.array(),
            Platform.BYTE_ARRAY_OFFSET + buffer.arrayOffset() + pos,
            size);
}
 

public boolean onStringValue(Object base, long offset, int size) {
    if (curRowIsTagField) {
        boolean ok = false;
        for (byte[] tag : acceptTags) {
            if (UTF8Util.containsCommaSep(base, offset, size, tag)) {
                ok = true;
                break;
            }
        }
        if (!ok) {
            curRowIgnore = true;
        }
        return true;
    }

    int rawValueOffset = curRowRawValueBytes;
    putValue(curRowIndex, (((long) rawValueOffset) << 32) | (long) size);

    Platform.copyMemory(
            base, offset,
            rawValuesBuffer, BYTE_ARRAY_OFFSET + rawValueOffset,
            size);

    curRowRawValueBytes += size;
    return true;
}
 

public void expandPointerArray(LongArray newArray) {
    if (newArray.size() < array.size()) {
        throw new OutOfMemoryError("Not enough memory to grow pointer array");
    }
    Platform.copyMemory(
            array.getBaseObject(),
            array.getBaseOffset(),
            newArray.getBaseObject(),
            newArray.getBaseOffset(),
            array.size() * 8L);
    consumer.freeArray(array);
    array = newArray;
}
 

@Override
public void copyRange(LongArray src, int srcPos, LongArray dst, int dstPos, int length) {
    Platform.copyMemory(
            src.getBaseObject(),
            src.getBaseOffset() + srcPos * 16,
            dst.getBaseObject(),
            dst.getBaseOffset() + dstPos * 16,
            length * 16);
}
 

/**
 * Returns the underlying bytes for this UnsafeRow.
 */
public byte[] getBytes() {
    if (baseObject instanceof byte[] && baseOffset == Platform.BYTE_ARRAY_OFFSET
            && (((byte[]) baseObject).length == sizeInBytes)) {
        return (byte[]) baseObject;
    } else {
        byte[] bytes = new byte[sizeInBytes];
        Platform.copyMemory(baseObject, baseOffset, bytes, Platform.BYTE_ARRAY_OFFSET, sizeInBytes);
        return bytes;
    }
}
 

@Override
public byte[] getRaw(int colId) {
    int offset = colId == 0 ? 0 : sums[colId - 1];
    int to = sums[colId];
    int len = to - offset;
    byte[] bytes = new byte[len];
    Platform.copyMemory(data.array(), BYTE_ARRAY_OFFSET + offset, bytes, BYTE_ARRAY_OFFSET, len);
    return bytes;
}
 

public boolean onStringValue(byte[] value, int offset, int size) {
    if (size > ColumnType.MAX_STRING_UTF8_SIZE) {
        logger.warn("string size overflow, expected less than {}, got {}", ColumnType.MAX_STRING_UTF8_SIZE, size);
        return false;
    }

    if (curRowIsTagField) {
        boolean ok = false;
        for (byte[] tag : acceptTags) {
            if (UTF8Util.containsCommaSep(value, offset, size, tag)) {
                ok = true;
                break;
            }
        }
        if (!ok) {
            curRowIgnore = true;
        }
        return true;
    }

    int rawValueOffset = curRowRawValueBytes;
    putValue(curRowIndex, (((long) rawValueOffset) << 32) | (long) size);
    Platform.copyMemory(
            value, BYTE_ARRAY_OFFSET + offset,
            rawValuesBuffer, BYTE_ARRAY_OFFSET + rawValueOffset, size);
    curRowRawValueBytes += size;
    return true;
}
 

@Override
public byte[] getRaw(int colId) {
    if (rowDataAddr == 0) {
        throw new IllegalStateException("illegal row");
    }
    long offsetAndLen = strOffsetLen(colId);
    int offset = (int) (offsetAndLen >>> 32);
    int len = (int) offsetAndLen & ColumnType.MAX_STRING_UTF8_SIZE_MASK;
    byte[] bytes = new byte[len];
    Platform.copyMemory(null, rowDataAddr + offset, bytes, BYTE_ARRAY_OFFSET, len);
    return bytes;
}
 

/**
 * Writes the content of this row into a memory address, identified by an object and an offset.
 * The target memory address must already been allocated, and have enough space to hold all the
 * bytes in this string.
 */
public void writeToMemory(Object target, long targetOffset) {
    Platform.copyMemory(baseObject, baseOffset, target, targetOffset, sizeInBytes);
}