org.apache.spark.unsafe.Platform Java Examples

private void checkStrings(DataPack[] dataPacks, int entryCount, ByteBufferReader entryFileReader, ByteBuffer bitmapFileBuffer) throws IOException {
    int packCount = dataPacks.length;
    int bitmapSize = DirectBitMap.bits2words(packCount) << 3;
    byte[] bitmapBuffer = new byte[bitmapSize];
    DirectBitMap bitMap = new DirectBitMap(bitmapBuffer, Platform.BYTE_ARRAY_OFFSET, bitmapSize >>> 3, null);

    StringsStructOnByteBufferReader stringsStruct = new StringsStructOnByteBufferReader(entryCount, entryFileReader);
    byte[] keyBuffer = new byte[ColumnType.MAX_STRING_UTF8_SIZE];
    BytePiece value = new BytePiece();
    for (int packId = 0; packId < packCount; packId++) {
        DataPack dataPack = dataPacks[packId];
        for (int rowId = 0; rowId < dataPack.valueCount(); rowId++) {
            dataPack.rawValueAt(rowId, value);
            int pos = BinarySearch.binarySearchStrings(stringsStruct, value.base, value.addr, value.len, keyBuffer);
            Assert.assertTrue(pos >= 0);
            bitmapFileBuffer.position(pos * bitmapSize);
            bitmapFileBuffer.get(bitmapBuffer);
            Assert.assertTrue(bitMap.get(packId));
        }
    }
}
 

int writeMap(MapColumnVector v, int row) {
  int offset = (int) v.offsets[row];
  int length = (int) v.lengths[row];
  int start = holder.cursor;
  // save room for the key size
  final int KEY_SIZE_BYTES = 8;
  holder.grow(KEY_SIZE_BYTES);
  holder.cursor += KEY_SIZE_BYTES;
  // serialize the keys
  childWriter.initialize(holder, length, keySize);
  for(int c=0; c < length; ++c) {
    keyConvert.convert(childWriter, c, v.keys, offset + c);
  }
  // store the serialized size of the keys
  Platform.putLong(holder.buffer, start, holder.cursor - start - KEY_SIZE_BYTES);
  // serialize the values
  childWriter.initialize(holder, length, valueSize);
  for(int c=0; c < length; ++c) {
    valueConvert.convert(childWriter, c, v.values, offset + c);
  }
  return start;
}
 

@Override
public UnsafeRow next() {
  if (current == null || nextRow >= current.size) {
    current = reader.next();
    nextRow = 0;
  }
  // Reset the holder to start the buffer over again.
  // BufferHolder.reset does the wrong thing...
  holder.cursor = Platform.BYTE_ARRAY_OFFSET;
  writer.reset();
  for(int c=0; c < current.cols.length; ++c) {
    converter[c].convert(writer, c, current.cols[c], nextRow);
  }
  nextRow++;
  return row;
}
 

public static int binarySearchInts(ByteBufferReader reader, int count, int key) throws IOException {
    int from = 0;
    int to = count;

    --to;

    byte[] valBuffer = new byte[4];

    while (from <= to) {
        int mid = from + to >>> 1;
        reader.read(mid << 2, valBuffer, 0, 4);
        int midVal = Platform.getInt(valBuffer, Platform.BYTE_ARRAY_OFFSET);
        if (midVal < key) {
            from = mid + 1;
        } else if (midVal > key) {
            to = mid - 1;
        } else {
            return mid;
        }
    }

    return -(from + 1);
}
 

public static int binarySearchLongs(ByteBufferReader reader, int count, long key) throws IOException {
    int from = 0;
    int to = count;

    --to;

    byte[] valBuffer = new byte[8];

    while (from <= to) {
        int mid = from + to >>> 1;
        reader.read(mid << 3, valBuffer, 0, 8);
        long midVal = Platform.getLong(valBuffer, Platform.BYTE_ARRAY_OFFSET);
        if (midVal < key) {
            from = mid + 1;
        } else if (midVal > key) {
            to = mid - 1;
        } else {
            return mid;
        }
    }

    return -(from + 1);
}
 

public static int binarySearchFloats(ByteBufferReader reader, int count, float key) throws IOException {
    int from = 0;
    int to = count;

    --to;

    byte[] valBuffer = new byte[4];

    while (from <= to) {
        int mid = from + to >>> 1;
        reader.read(mid << 2, valBuffer, 0, 4);
        float midVal = Platform.getFloat(valBuffer, Platform.BYTE_ARRAY_OFFSET);
        if (midVal < key) {
            from = mid + 1;
        } else if (midVal > key) {
            to = mid - 1;
        } else {
            return mid;
        }
    }

    return -(from + 1);
}
 

static void _putValue(long packAddr, UTF8String value) {
    int valueSize = value.numBytes();
    Object valueBase = value.getBaseObject();
    long valueOffset = value.getBaseOffset();

    long offset = packAddr + indexOffsetBySize(valueSize);
    int checkSize = valueSize < MAX_POSISTIONS ? valueSize : MAX_POSISTIONS;

    if (checkSize == 0) {
        // mark empty string exists.
        set(offset, (byte) 1, 0);
    } else {
        for (int pos = 0; pos < checkSize; pos++) {
            set(offset, Platform.getByte(valueBase, valueOffset + pos), pos);
        }
    }
}
 

public static int binarySearchStrings(StringsStructOnByteBufferReader reader,
                                      Object keyBase, long keyAddr, int keyLen,
                                      byte[] valBuffer) throws IOException {
    int from = 0;
    int to = reader.valCount();

    --to;
    int cmp;
    while (from <= to) {
        int mid = from + to >>> 1;
        int valLen = reader.getString(mid, valBuffer);
        cmp = BytesUtil.compareBytes(valBuffer, Platform.BYTE_ARRAY_OFFSET, valLen, keyBase, keyAddr, keyLen);
        if (cmp < 0) {
            from = mid + 1;
        } else if (cmp > 0) {
            to = mid - 1;
        } else {
            return mid;
        }
    }

    return -(from + 1);
}
 

static byte _isValue(long packAddr, UTF8String value) {
    int valueSize = value.numBytes();
    Object valueBase = value.getBaseObject();
    long valueOffset = value.getBaseOffset();

    long offset = packAddr + indexOffsetBySize(valueSize);
    int checkSize = valueSize < MAX_POSISTIONS ? valueSize : MAX_POSISTIONS;

    if (checkSize == 0) {
        if (!isSet(offset, (byte) 1, 0)) {
            return RSValue.None;
        }
    } else {
        for (int pos = 0; pos < checkSize; pos++) {
            if (!isSet(offset, Platform.getByte(valueBase, valueOffset + pos), pos)) {
                return RSValue.None;
            }
        }
    }
    return RSValue.Some;
}
 

public static byte _isLike(long packAddr, LikePattern pattern) {
    // We can exclude cases like "ala%" and "a_l_a%"

    UTF8String original = pattern.original;
    int bytes = original.numBytes();
    Object valueBase = original.getBaseObject();
    long valueOffset = original.getBaseOffset();
    int indexSize = bytes < POSISTIONS ? bytes : POSISTIONS;

    for (int pos = 0; pos < indexSize; pos++) {
        byte c = Platform.getByte(valueBase, valueOffset + pos);
        // The ESCAPE_CHARACTOR case can be optimized. But I'm too lazy...
        if (c == '%' || c == ESCAPE_CHARACTOR) {
            break;
        }
        if (c != '_' && !isSet(packAddr, c, pos)) {
            return RSValue.None;
        }
    }
    return RSValue.Some;
}
 

/**
 * 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);
}
 

/**
 * 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);
}
 

private void checkInts(DataPack[] dataPacks, int entryCount, ByteBufferReader entryFileReader, ByteBuffer bitmapFileBuffer) throws IOException {
    int packCount = dataPacks.length;
    int bitmapSize = DirectBitMap.bits2words(packCount) << 3;
    byte[] bitmapBuffer = new byte[bitmapSize];
    DirectBitMap bitMap = new DirectBitMap(bitmapBuffer, Platform.BYTE_ARRAY_OFFSET, bitmapSize >>> 3, null);

    for (int packId = 0; packId < packCount; packId++) {
        DataPack dataPack = dataPacks[packId];
        for (int rowId = 0; rowId < dataPack.valueCount(); rowId++) {
            int value = dataPack.intValueAt(rowId);
            int pos = BinarySearch.binarySearchInts(entryFileReader, entryCount, value);
            Assert.assertTrue(pos >= 0);
            bitmapFileBuffer.position(pos * bitmapSize);
            bitmapFileBuffer.get(bitmapBuffer);
            Assert.assertTrue(bitMap.get(packId));
        }
    }
}
 

private void checkLongs(DataPack[] dataPacks, int entryCount, ByteBufferReader entryFileReader, ByteBuffer bitmapFileBuffer) throws IOException {
    int packCount = dataPacks.length;
    int bitmapSize = DirectBitMap.bits2words(packCount) << 3;
    byte[] bitmapBuffer = new byte[bitmapSize];
    DirectBitMap bitMap = new DirectBitMap(bitmapBuffer, Platform.BYTE_ARRAY_OFFSET, bitmapSize >>> 3, null);

    for (int packId = 0; packId < packCount; packId++) {
        DataPack dataPack = dataPacks[packId];
        for (int rowId = 0; rowId < dataPack.valueCount(); rowId++) {
            long value = dataPack.longValueAt(rowId);
            int pos = BinarySearch.binarySearchLongs(entryFileReader, entryCount, value);
            Assert.assertTrue(pos >= 0);
            bitmapFileBuffer.position(pos * bitmapSize);
            bitmapFileBuffer.get(bitmapBuffer);
            Assert.assertTrue(bitMap.get(packId));
        }
    }
}
 

/**
 * 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);
}
 

private void checkFloats(DataPack[] dataPacks, int entryCount, ByteBufferReader entryFileReader, ByteBuffer bitmapFileBuffer) throws IOException {
    int packCount = dataPacks.length;
    int bitmapSize = DirectBitMap.bits2words(packCount) << 3;
    byte[] bitmapBuffer = new byte[bitmapSize];
    DirectBitMap bitMap = new DirectBitMap(bitmapBuffer, Platform.BYTE_ARRAY_OFFSET, bitmapSize >>> 3, null);

    for (int packId = 0; packId < packCount; packId++) {
        DataPack dataPack = dataPacks[packId];
        for (int rowId = 0; rowId < dataPack.valueCount(); rowId++) {
            float value = dataPack.floatValueAt(rowId);
            int pos = BinarySearch.binarySearchFloats(entryFileReader, entryCount, value);
            Assert.assertTrue(pos >= 0);
            bitmapFileBuffer.position(pos * bitmapSize);
            bitmapFileBuffer.get(bitmapBuffer);
            Assert.assertTrue(bitMap.get(packId));
        }
    }
}
 

private void checkDoubles(DataPack[] dataPacks, int entryCount, ByteBufferReader entryFileReader, ByteBuffer bitmapFileBuffer) throws IOException {
    int packCount = dataPacks.length;
    int bitmapSize = DirectBitMap.bits2words(packCount) << 3;
    byte[] bitmapBuffer = new byte[bitmapSize];
    DirectBitMap bitMap = new DirectBitMap(bitmapBuffer, Platform.BYTE_ARRAY_OFFSET, bitmapSize >>> 3, null);

    for (int packId = 0; packId < packCount; packId++) {
        DataPack dataPack = dataPacks[packId];
        for (int rowId = 0; rowId < dataPack.valueCount(); rowId++) {
            double value = dataPack.doubleValueAt(rowId);
            int pos = BinarySearch.binarySearchDoubles(entryFileReader, entryCount, value);
            Assert.assertTrue(pos >= 0);
            bitmapFileBuffer.position(pos * bitmapSize);
            bitmapFileBuffer.get(bitmapBuffer);
            Assert.assertTrue(bitMap.get(packId));
        }
    }
}
 

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 writeTo(ByteBuffer buffer) {
    assert (buffer.hasArray());
    byte[] target = buffer.array();
    int offset = buffer.arrayOffset();
    int pos = buffer.position();
    writeToMemory(target, Platform.BYTE_ARRAY_OFFSET + offset + pos);
    buffer.position(pos + sizeInBytes);
}
 

@Override
public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException {
    this.baseOffset = Platform.BYTE_ARRAY_OFFSET;
    this.sizeInBytes = in.readInt();
    this.numFields = in.readInt();
    this.baseObject = new byte[sizeInBytes];
    in.readFully((byte[]) baseObject);
}
 

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 void setFloat(int ordinal, float value) {
    assertIndexIsValid(ordinal);
    if (Float.isNaN(value)) {
        value = Float.NaN;
    }
    Platform.putDouble(baseObject, getFieldOffset(ordinal), value);
}
 

private static void _putValue(long packAddr, UTF8String value) {
    int bytes = value.numBytes();
    Object valueBase = value.getBaseObject();
    long valueOffset = value.getBaseOffset();
    int indexSize = bytes < POSISTIONS ? bytes : POSISTIONS;

    for (int pos = 0; pos < indexSize; pos++) {
        set(packAddr, Platform.getByte(valueBase, valueOffset + pos), pos);
    }
}
 

@Override
public void setDouble(int ordinal, double value) {
    assertIndexIsValid(ordinal);
    if (Double.isNaN(value)) {
        value = Double.NaN;
    }
    Platform.putDouble(baseObject, getFieldOffset(ordinal), value);
}
 

@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 number of set bits in an array of longs. */
public static long pop_array(Object base, long arrAddr, int wordOffset, int numWords) {
    long popCount = 0;
    for (int i = wordOffset, end = wordOffset + numWords; i < end; ++i) {
        popCount += Long.bitCount(Platform.getLong(base, arrAddr + (i << 3)));
    }
    return popCount;
}
 

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;
}
 

/**
 * Returns the popcount or cardinality of A ^ B
 * Neither array is modified.
 */
public static long pop_xor(Object base1, long arr1Addr, Object base2, long arr2Addr, int wordOffset, int numWords) {
    long popCount = 0;
    for (int i = wordOffset, end = wordOffset + numWords; i < end; ++i) {
        popCount += Long.bitCount(Platform.getLong(base1, arr1Addr + (i << 3)) ^ Platform.getLong(base2, arr2Addr + (i << 3)));
    }
    return popCount;
}
 

public UnsafeRow getAggregationBufferFromUnsafeRow(UnsafeRow unsafeGroupingKeyRow) {
    // Probe our map using the serialized key
    final BytesToBytesMap.Location loc = map.lookup(
            unsafeGroupingKeyRow.getBaseObject(),
            unsafeGroupingKeyRow.getBaseOffset(),
            unsafeGroupingKeyRow.getSizeInBytes());
    if (!loc.isDefined()) {
        // This is the first time that we've seen this grouping key, so we'll insert a copy of the
        // empty aggregation buffer into the map:
        boolean putSucceeded = loc.putNewKey(
                unsafeGroupingKeyRow.getBaseObject(),
                unsafeGroupingKeyRow.getBaseOffset(),
                unsafeGroupingKeyRow.getSizeInBytes(),
                emptyAggregationBuffer,
                Platform.BYTE_ARRAY_OFFSET,
                emptyAggregationBuffer.length
        );
        if (!putSucceeded) {
            return null;
        }
    }

    // Reset the pointer to point to the value that we just stored or looked up:
    final MemoryLocation address = loc.getValueAddress();
    currentAggregationBuffer.pointTo(
            address.getBaseObject(),
            address.getBaseOffset(),
            loc.getValueLength()
    );
    return currentAggregationBuffer;
}