org.apache.flink.table.data.binary.BinaryArrayData Maven / Gradle / Ivy
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.flink.table.data.binary;
import org.apache.flink.annotation.Internal;
import org.apache.flink.core.memory.MemorySegment;
import org.apache.flink.core.memory.MemorySegmentFactory;
import org.apache.flink.table.data.ArrayData;
import org.apache.flink.table.data.DecimalData;
import org.apache.flink.table.data.MapData;
import org.apache.flink.table.data.RawValueData;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.data.StringData;
import org.apache.flink.table.data.TimestampData;
import org.apache.flink.table.types.logical.DistinctType;
import org.apache.flink.table.types.logical.LogicalType;
import org.apache.flink.table.types.logical.utils.LogicalTypeUtils;
import java.lang.reflect.Array;
import static org.apache.flink.core.memory.MemoryUtils.UNSAFE;
/**
* A binary implementation of {@link ArrayData} which is backed by {@link MemorySegment}s.
*
* For fields that hold fixed-length primitive types, such as long, double or int, they are
* stored compacted in bytes, just like the original java array.
*
*
The binary layout of {@link BinaryArrayData}:
*
*
* [size(int)] + [null bits(4-byte word boundaries)] + [values or offset&length] + [variable length part].
*
*/
@Internal
public final class BinaryArrayData extends BinarySection implements ArrayData, TypedSetters {
/** Offset for Arrays. */
private static final int BYTE_ARRAY_BASE_OFFSET = UNSAFE.arrayBaseOffset(byte[].class);
private static final int BOOLEAN_ARRAY_OFFSET = UNSAFE.arrayBaseOffset(boolean[].class);
private static final int SHORT_ARRAY_OFFSET = UNSAFE.arrayBaseOffset(short[].class);
private static final int INT_ARRAY_OFFSET = UNSAFE.arrayBaseOffset(int[].class);
private static final int LONG_ARRAY_OFFSET = UNSAFE.arrayBaseOffset(long[].class);
private static final int FLOAT_ARRAY_OFFSET = UNSAFE.arrayBaseOffset(float[].class);
private static final int DOUBLE_ARRAY_OFFSET = UNSAFE.arrayBaseOffset(double[].class);
public static int calculateHeaderInBytes(int numFields) {
return 4 + ((numFields + 31) / 32) * 4;
}
/**
* It store real value when type is primitive. It store the length and offset of variable-length
* part when type is string, map, etc.
*/
public static int calculateFixLengthPartSize(LogicalType type) {
// ordered by type root definition
switch (type.getTypeRoot()) {
case BOOLEAN:
case TINYINT:
return 1;
case CHAR:
case VARCHAR:
case BINARY:
case VARBINARY:
case DECIMAL:
case BIGINT:
case DOUBLE:
case TIMESTAMP_WITHOUT_TIME_ZONE:
case TIMESTAMP_WITH_LOCAL_TIME_ZONE:
case INTERVAL_DAY_TIME:
case ARRAY:
case MULTISET:
case MAP:
case ROW:
case STRUCTURED_TYPE:
case RAW:
// long and double are 8 bytes;
// otherwise it stores the length and offset of the variable-length part for types
// such as is string, map, etc.
return 8;
case TIMESTAMP_WITH_TIME_ZONE:
throw new UnsupportedOperationException();
case SMALLINT:
return 2;
case INTEGER:
case FLOAT:
case DATE:
case TIME_WITHOUT_TIME_ZONE:
case INTERVAL_YEAR_MONTH:
return 4;
case DISTINCT_TYPE:
return calculateFixLengthPartSize(((DistinctType) type).getSourceType());
case NULL:
case SYMBOL:
case UNRESOLVED:
default:
throw new IllegalArgumentException();
}
}
// The number of elements in this array
private int size;
/** The position to start storing array elements. */
private int elementOffset;
public BinaryArrayData() {}
private void assertIndexIsValid(int ordinal) {
assert ordinal >= 0 : "ordinal (" + ordinal + ") should >= 0";
assert ordinal < size : "ordinal (" + ordinal + ") should < " + size;
}
private int getElementOffset(int ordinal, int elementSize) {
return elementOffset + ordinal * elementSize;
}
@Override
public int size() {
return size;
}
@Override
public void pointTo(MemorySegment[] segments, int offset, int sizeInBytes) {
// Read the number of elements from the first 4 bytes.
final int size = BinarySegmentUtils.getInt(segments, offset);
assert size >= 0 : "size (" + size + ") should >= 0";
this.size = size;
this.segments = segments;
this.offset = offset;
this.sizeInBytes = sizeInBytes;
this.elementOffset = offset + calculateHeaderInBytes(this.size);
}
@Override
public boolean isNullAt(int pos) {
assertIndexIsValid(pos);
return BinarySegmentUtils.bitGet(segments, offset + 4, pos);
}
public void setNullAt(int pos) {
assertIndexIsValid(pos);
BinarySegmentUtils.bitSet(segments, offset + 4, pos);
}
public void setNotNullAt(int pos) {
assertIndexIsValid(pos);
BinarySegmentUtils.bitUnSet(segments, offset + 4, pos);
}
@Override
public long getLong(int pos) {
assertIndexIsValid(pos);
return BinarySegmentUtils.getLong(segments, getElementOffset(pos, 8));
}
public void setLong(int pos, long value) {
assertIndexIsValid(pos);
setNotNullAt(pos);
BinarySegmentUtils.setLong(segments, getElementOffset(pos, 8), value);
}
public void setNullLong(int pos) {
assertIndexIsValid(pos);
BinarySegmentUtils.bitSet(segments, offset + 4, pos);
BinarySegmentUtils.setLong(segments, getElementOffset(pos, 8), 0L);
}
@Override
public int getInt(int pos) {
assertIndexIsValid(pos);
return BinarySegmentUtils.getInt(segments, getElementOffset(pos, 4));
}
public void setInt(int pos, int value) {
assertIndexIsValid(pos);
setNotNullAt(pos);
BinarySegmentUtils.setInt(segments, getElementOffset(pos, 4), value);
}
public void setNullInt(int pos) {
assertIndexIsValid(pos);
BinarySegmentUtils.bitSet(segments, offset + 4, pos);
BinarySegmentUtils.setInt(segments, getElementOffset(pos, 4), 0);
}
@Override
public StringData getString(int pos) {
assertIndexIsValid(pos);
int fieldOffset = getElementOffset(pos, 8);
final long offsetAndSize = BinarySegmentUtils.getLong(segments, fieldOffset);
return BinarySegmentUtils.readStringData(segments, offset, fieldOffset, offsetAndSize);
}
@Override
public DecimalData getDecimal(int pos, int precision, int scale) {
assertIndexIsValid(pos);
if (DecimalData.isCompact(precision)) {
return DecimalData.fromUnscaledLong(
BinarySegmentUtils.getLong(segments, getElementOffset(pos, 8)),
precision,
scale);
}
int fieldOffset = getElementOffset(pos, 8);
final long offsetAndSize = BinarySegmentUtils.getLong(segments, fieldOffset);
return BinarySegmentUtils.readDecimalData(
segments, offset, offsetAndSize, precision, scale);
}
@Override
public TimestampData getTimestamp(int pos, int precision) {
assertIndexIsValid(pos);
if (TimestampData.isCompact(precision)) {
return TimestampData.fromEpochMillis(
BinarySegmentUtils.getLong(segments, getElementOffset(pos, 8)));
}
int fieldOffset = getElementOffset(pos, 8);
final long offsetAndNanoOfMilli = BinarySegmentUtils.getLong(segments, fieldOffset);
return BinarySegmentUtils.readTimestampData(segments, offset, offsetAndNanoOfMilli);
}
@Override
public RawValueData getRawValue(int pos) {
assertIndexIsValid(pos);
int fieldOffset = getElementOffset(pos, 8);
final long offsetAndSize = BinarySegmentUtils.getLong(segments, fieldOffset);
return BinarySegmentUtils.readRawValueData(segments, offset, offsetAndSize);
}
@Override
public byte[] getBinary(int pos) {
assertIndexIsValid(pos);
int fieldOffset = getElementOffset(pos, 8);
final long offsetAndSize = BinarySegmentUtils.getLong(segments, fieldOffset);
return BinarySegmentUtils.readBinary(segments, offset, fieldOffset, offsetAndSize);
}
@Override
public ArrayData getArray(int pos) {
assertIndexIsValid(pos);
return BinarySegmentUtils.readArrayData(segments, offset, getLong(pos));
}
@Override
public MapData getMap(int pos) {
assertIndexIsValid(pos);
return BinarySegmentUtils.readMapData(segments, offset, getLong(pos));
}
@Override
public RowData getRow(int pos, int numFields) {
assertIndexIsValid(pos);
int fieldOffset = getElementOffset(pos, 8);
final long offsetAndSize = BinarySegmentUtils.getLong(segments, fieldOffset);
return BinarySegmentUtils.readRowData(segments, numFields, offset, offsetAndSize);
}
@Override
public boolean getBoolean(int pos) {
assertIndexIsValid(pos);
return BinarySegmentUtils.getBoolean(segments, getElementOffset(pos, 1));
}
public void setBoolean(int pos, boolean value) {
assertIndexIsValid(pos);
setNotNullAt(pos);
BinarySegmentUtils.setBoolean(segments, getElementOffset(pos, 1), value);
}
public void setNullBoolean(int pos) {
assertIndexIsValid(pos);
BinarySegmentUtils.bitSet(segments, offset + 4, pos);
BinarySegmentUtils.setBoolean(segments, getElementOffset(pos, 1), false);
}
@Override
public byte getByte(int pos) {
assertIndexIsValid(pos);
return BinarySegmentUtils.getByte(segments, getElementOffset(pos, 1));
}
public void setByte(int pos, byte value) {
assertIndexIsValid(pos);
setNotNullAt(pos);
BinarySegmentUtils.setByte(segments, getElementOffset(pos, 1), value);
}
public void setNullByte(int pos) {
assertIndexIsValid(pos);
BinarySegmentUtils.bitSet(segments, offset + 4, pos);
BinarySegmentUtils.setByte(segments, getElementOffset(pos, 1), (byte) 0);
}
@Override
public short getShort(int pos) {
assertIndexIsValid(pos);
return BinarySegmentUtils.getShort(segments, getElementOffset(pos, 2));
}
public void setShort(int pos, short value) {
assertIndexIsValid(pos);
setNotNullAt(pos);
BinarySegmentUtils.setShort(segments, getElementOffset(pos, 2), value);
}
public void setNullShort(int pos) {
assertIndexIsValid(pos);
BinarySegmentUtils.bitSet(segments, offset + 4, pos);
BinarySegmentUtils.setShort(segments, getElementOffset(pos, 2), (short) 0);
}
@Override
public float getFloat(int pos) {
assertIndexIsValid(pos);
return BinarySegmentUtils.getFloat(segments, getElementOffset(pos, 4));
}
public void setFloat(int pos, float value) {
assertIndexIsValid(pos);
setNotNullAt(pos);
BinarySegmentUtils.setFloat(segments, getElementOffset(pos, 4), value);
}
public void setNullFloat(int pos) {
assertIndexIsValid(pos);
BinarySegmentUtils.bitSet(segments, offset + 4, pos);
BinarySegmentUtils.setFloat(segments, getElementOffset(pos, 4), 0F);
}
@Override
public double getDouble(int pos) {
assertIndexIsValid(pos);
return BinarySegmentUtils.getDouble(segments, getElementOffset(pos, 8));
}
public void setDouble(int pos, double value) {
assertIndexIsValid(pos);
setNotNullAt(pos);
BinarySegmentUtils.setDouble(segments, getElementOffset(pos, 8), value);
}
public void setNullDouble(int pos) {
assertIndexIsValid(pos);
BinarySegmentUtils.bitSet(segments, offset + 4, pos);
BinarySegmentUtils.setDouble(segments, getElementOffset(pos, 8), 0.0);
}
public void setDecimal(int pos, DecimalData value, int precision) {
assertIndexIsValid(pos);
if (DecimalData.isCompact(precision)) {
// compact format
setLong(pos, value.toUnscaledLong());
} else {
int fieldOffset = getElementOffset(pos, 8);
int cursor = (int) (BinarySegmentUtils.getLong(segments, fieldOffset) >>> 32);
assert cursor > 0 : "invalid cursor " + cursor;
// zero-out the bytes
BinarySegmentUtils.setLong(segments, offset + cursor, 0L);
BinarySegmentUtils.setLong(segments, offset + cursor + 8, 0L);
if (value == null) {
setNullAt(pos);
// keep the offset for future update
BinarySegmentUtils.setLong(segments, fieldOffset, ((long) cursor) << 32);
} else {
byte[] bytes = value.toUnscaledBytes();
assert (bytes.length <= 16);
// Write the bytes to the variable length portion.
BinarySegmentUtils.copyFromBytes(segments, offset + cursor, bytes, 0, bytes.length);
setLong(pos, ((long) cursor << 32) | ((long) bytes.length));
}
}
}
public void setTimestamp(int pos, TimestampData value, int precision) {
assertIndexIsValid(pos);
if (TimestampData.isCompact(precision)) {
setLong(pos, value.getMillisecond());
} else {
int fieldOffset = getElementOffset(pos, 8);
int cursor = (int) (BinarySegmentUtils.getLong(segments, fieldOffset) >>> 32);
assert cursor > 0 : "invalid cursor " + cursor;
if (value == null) {
setNullAt(pos);
// zero-out the bytes
BinarySegmentUtils.setLong(segments, offset + cursor, 0L);
// keep the offset for future update
BinarySegmentUtils.setLong(segments, fieldOffset, ((long) cursor) << 32);
} else {
// write millisecond to the variable length portion.
BinarySegmentUtils.setLong(segments, offset + cursor, value.getMillisecond());
// write nanoOfMillisecond to the fixed-length portion.
setLong(pos, ((long) cursor << 32) | (long) value.getNanoOfMillisecond());
}
}
}
public boolean anyNull() {
for (int i = offset + 4; i < elementOffset; i += 4) {
if (BinarySegmentUtils.getInt(segments, i) != 0) {
return true;
}
}
return false;
}
private void checkNoNull() {
if (anyNull()) {
throw new RuntimeException("Primitive array must not contain a null value.");
}
}
@Override
public boolean[] toBooleanArray() {
checkNoNull();
boolean[] values = new boolean[size];
BinarySegmentUtils.copyToUnsafe(
segments, elementOffset, values, BOOLEAN_ARRAY_OFFSET, size);
return values;
}
@Override
public byte[] toByteArray() {
checkNoNull();
byte[] values = new byte[size];
BinarySegmentUtils.copyToUnsafe(
segments, elementOffset, values, BYTE_ARRAY_BASE_OFFSET, size);
return values;
}
@Override
public short[] toShortArray() {
checkNoNull();
short[] values = new short[size];
BinarySegmentUtils.copyToUnsafe(
segments, elementOffset, values, SHORT_ARRAY_OFFSET, size * 2);
return values;
}
@Override
public int[] toIntArray() {
checkNoNull();
int[] values = new int[size];
BinarySegmentUtils.copyToUnsafe(
segments, elementOffset, values, INT_ARRAY_OFFSET, size * 4);
return values;
}
@Override
public long[] toLongArray() {
checkNoNull();
long[] values = new long[size];
BinarySegmentUtils.copyToUnsafe(
segments, elementOffset, values, LONG_ARRAY_OFFSET, size * 8);
return values;
}
@Override
public float[] toFloatArray() {
checkNoNull();
float[] values = new float[size];
BinarySegmentUtils.copyToUnsafe(
segments, elementOffset, values, FLOAT_ARRAY_OFFSET, size * 4);
return values;
}
@Override
public double[] toDoubleArray() {
checkNoNull();
double[] values = new double[size];
BinarySegmentUtils.copyToUnsafe(
segments, elementOffset, values, DOUBLE_ARRAY_OFFSET, size * 8);
return values;
}
@SuppressWarnings("unchecked")
public T[] toObjectArray(LogicalType elementType) {
Class elementClass = (Class) LogicalTypeUtils.toInternalConversionClass(elementType);
ArrayData.ElementGetter elementGetter = ArrayData.createElementGetter(elementType);
T[] values = (T[]) Array.newInstance(elementClass, size);
for (int i = 0; i < size; i++) {
if (!isNullAt(i)) {
values[i] = (T) elementGetter.getElementOrNull(this, i);
}
}
return values;
}
public BinaryArrayData copy() {
return copy(new BinaryArrayData());
}
public BinaryArrayData copy(BinaryArrayData reuse) {
byte[] bytes = BinarySegmentUtils.copyToBytes(segments, offset, sizeInBytes);
reuse.pointTo(MemorySegmentFactory.wrap(bytes), 0, sizeInBytes);
return reuse;
}
@Override
public int hashCode() {
return BinarySegmentUtils.hashByWords(segments, offset, sizeInBytes);
}
// ------------------------------------------------------------------------------------------
// Construction Utilities
// ------------------------------------------------------------------------------------------
public static BinaryArrayData fromPrimitiveArray(boolean[] arr) {
return fromPrimitiveArray(arr, BOOLEAN_ARRAY_OFFSET, arr.length, 1);
}
public static BinaryArrayData fromPrimitiveArray(byte[] arr) {
return fromPrimitiveArray(arr, BYTE_ARRAY_BASE_OFFSET, arr.length, 1);
}
public static BinaryArrayData fromPrimitiveArray(short[] arr) {
return fromPrimitiveArray(arr, SHORT_ARRAY_OFFSET, arr.length, 2);
}
public static BinaryArrayData fromPrimitiveArray(int[] arr) {
return fromPrimitiveArray(arr, INT_ARRAY_OFFSET, arr.length, 4);
}
public static BinaryArrayData fromPrimitiveArray(long[] arr) {
return fromPrimitiveArray(arr, LONG_ARRAY_OFFSET, arr.length, 8);
}
public static BinaryArrayData fromPrimitiveArray(float[] arr) {
return fromPrimitiveArray(arr, FLOAT_ARRAY_OFFSET, arr.length, 4);
}
public static BinaryArrayData fromPrimitiveArray(double[] arr) {
return fromPrimitiveArray(arr, DOUBLE_ARRAY_OFFSET, arr.length, 8);
}
private static BinaryArrayData fromPrimitiveArray(
Object arr, int offset, int length, int elementSize) {
final long headerInBytes = calculateHeaderInBytes(length);
final long valueRegionInBytes = elementSize * length;
// must align by 8 bytes
long totalSizeInLongs = (headerInBytes + valueRegionInBytes + 7) / 8;
if (totalSizeInLongs > Integer.MAX_VALUE / 8) {
throw new UnsupportedOperationException(
"Cannot convert this array to unsafe format as " + "it's too big.");
}
long totalSize = totalSizeInLongs * 8;
final byte[] data = new byte[(int) totalSize];
UNSAFE.putInt(data, (long) BYTE_ARRAY_BASE_OFFSET, length);
UNSAFE.copyMemory(
arr, offset, data, BYTE_ARRAY_BASE_OFFSET + headerInBytes, valueRegionInBytes);
BinaryArrayData result = new BinaryArrayData();
result.pointTo(MemorySegmentFactory.wrap(data), 0, (int) totalSize);
return result;
}
}
© 2015 - 2025 Weber Informatics LLC | Privacy Policy