mirror of
https://github.com/enso-org/enso.git
synced 2024-11-22 22:10:15 +03:00
Arrow builder is not an Array (#9358)
Follow up on #9150 - making sure that Arrow builder is not accidentally treated as an Array by disallowing reading elements. # Important Notes Also making sure that the length of the resulting Arrow Array is consistent with what user requested.
This commit is contained in:
Hubert Plociniczak
GitHub
parent
2322b40a22
commit
f82e8020fe
@ -14,6 +14,7 @@ import com.oracle.truffle.api.library.CachedLibrary;
|
||||
import com.oracle.truffle.api.library.ExportLibrary;
|
||||
import com.oracle.truffle.api.library.ExportMessage;
|
||||
import org.enso.interpreter.arrow.LogicalLayout;
|
||||
import org.graalvm.collections.Pair;
|
||||
|
||||
@ExportLibrary(InteropLibrary.class)
|
||||
public class ArrowCastToFixedSizeArrayFactory implements TruffleObject {
|
||||
@ -43,7 +44,8 @@ public class ArrowCastToFixedSizeArrayFactory implements TruffleObject {
|
||||
@Cached.Shared("interop") @CachedLibrary(limit = "1") InteropLibrary iop)
|
||||
throws UnsupportedMessageException, ArityException, UnsupportedTypeException {
|
||||
var unit = LogicalLayout.Date32;
|
||||
return new ArrowFixedArrayDate(pointer(args, iop, unit), unit);
|
||||
var pair = pointer(args, iop, unit);
|
||||
return new ArrowFixedArrayDate(pair.getLeft(), pair.getRight(), unit);
|
||||
}
|
||||
|
||||
@Specialization(guards = "receiver.getLayout() == Date64")
|
||||
@ -53,7 +55,8 @@ public class ArrowCastToFixedSizeArrayFactory implements TruffleObject {
|
||||
@Cached.Shared("interop") @CachedLibrary(limit = "1") InteropLibrary iop)
|
||||
throws UnsupportedMessageException, ArityException, UnsupportedTypeException {
|
||||
var unit = LogicalLayout.Date64;
|
||||
return new ArrowFixedArrayDate(pointer(args, iop, unit), unit);
|
||||
var pair = pointer(args, iop, unit);
|
||||
return new ArrowFixedArrayDate(pair.getLeft(), pair.getRight(), unit);
|
||||
}
|
||||
|
||||
@Specialization(guards = "receiver.getLayout() == Int8")
|
||||
@ -63,7 +66,8 @@ public class ArrowCastToFixedSizeArrayFactory implements TruffleObject {
|
||||
@Cached.Shared("interop") @CachedLibrary(limit = "1") InteropLibrary iop)
|
||||
throws UnsupportedMessageException, ArityException, UnsupportedTypeException {
|
||||
var unit = LogicalLayout.Int8;
|
||||
return new ArrowFixedArrayInt(pointer(args, iop, unit), unit);
|
||||
var pair = pointer(args, iop, unit);
|
||||
return new ArrowFixedArrayInt(pair.getLeft(), pair.getRight(), unit);
|
||||
}
|
||||
|
||||
@Specialization(guards = "receiver.getLayout() == Int16")
|
||||
@ -73,7 +77,8 @@ public class ArrowCastToFixedSizeArrayFactory implements TruffleObject {
|
||||
@Cached.Shared("interop") @CachedLibrary(limit = "1") InteropLibrary iop)
|
||||
throws UnsupportedMessageException, ArityException, UnsupportedTypeException {
|
||||
var unit = LogicalLayout.Int16;
|
||||
return new ArrowFixedArrayInt(pointer(args, iop, unit), unit);
|
||||
var pair = pointer(args, iop, unit);
|
||||
return new ArrowFixedArrayInt(pair.getLeft(), pair.getRight(), unit);
|
||||
}
|
||||
|
||||
@Specialization(guards = "receiver.getLayout() == Int32")
|
||||
@ -83,7 +88,8 @@ public class ArrowCastToFixedSizeArrayFactory implements TruffleObject {
|
||||
@Cached.Shared("interop") @CachedLibrary(limit = "1") InteropLibrary iop)
|
||||
throws UnsupportedMessageException, ArityException, UnsupportedTypeException {
|
||||
var unit = LogicalLayout.Int32;
|
||||
return new ArrowFixedArrayInt(pointer(args, iop, unit), unit);
|
||||
var pair = pointer(args, iop, unit);
|
||||
return new ArrowFixedArrayInt(pair.getLeft(), pair.getRight(), unit);
|
||||
}
|
||||
|
||||
@Specialization(guards = "receiver.getLayout() == Int64")
|
||||
@ -93,11 +99,13 @@ public class ArrowCastToFixedSizeArrayFactory implements TruffleObject {
|
||||
@Cached.Shared("interop") @CachedLibrary(limit = "1") InteropLibrary iop)
|
||||
throws UnsupportedMessageException, ArityException, UnsupportedTypeException {
|
||||
var unit = LogicalLayout.Int64;
|
||||
return new ArrowFixedArrayInt(pointer(args, iop, unit), unit);
|
||||
var pair = pointer(args, iop, unit);
|
||||
return new ArrowFixedArrayInt(pair.getLeft(), pair.getRight(), unit);
|
||||
}
|
||||
|
||||
@CompilerDirectives.TruffleBoundary
|
||||
private static ByteBufferDirect pointer(Object[] args, InteropLibrary interop, SizeInBytes unit)
|
||||
private static Pair<ByteBufferDirect, Integer> pointer(
|
||||
Object[] args, InteropLibrary interop, SizeInBytes unit)
|
||||
throws ArityException, UnsupportedTypeException, UnsupportedMessageException {
|
||||
if (args.length < 2) {
|
||||
throw ArityException.create(2, 3, args.length);
|
||||
@ -117,10 +125,13 @@ public class ArrowCastToFixedSizeArrayFactory implements TruffleObject {
|
||||
throw UnsupportedTypeException.create(
|
||||
new Object[] {args[2]}, "Address of non-null bitmap is invalid");
|
||||
}
|
||||
return ByteBufferDirect.fromAddress(
|
||||
interop.asLong(args[0]), interop.asLong(args[2]), capacity, unit);
|
||||
return Pair.create(
|
||||
ByteBufferDirect.fromAddress(
|
||||
interop.asLong(args[0]), interop.asLong(args[2]), capacity, unit),
|
||||
capacity);
|
||||
} else {
|
||||
return ByteBufferDirect.fromAddress(interop.asLong(args[0]), capacity, unit);
|
||||
return Pair.create(
|
||||
ByteBufferDirect.fromAddress(interop.asLong(args[0]), capacity, unit), capacity);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -27,9 +27,8 @@ public final class ArrowFixedArrayDate implements TruffleObject {
|
||||
this.buffer = ByteBufferDirect.forSize(size, unit);
|
||||
}
|
||||
|
||||
public ArrowFixedArrayDate(ByteBufferDirect buffer, LogicalLayout unit)
|
||||
throws UnsupportedMessageException {
|
||||
this.size = buffer.capacity() / unit.sizeInBytes();
|
||||
public ArrowFixedArrayDate(ByteBufferDirect buffer, int size, LogicalLayout unit) {
|
||||
this.size = size;
|
||||
this.unit = unit;
|
||||
this.buffer = buffer;
|
||||
}
|
||||
@ -164,7 +163,7 @@ public final class ArrowFixedArrayDate implements TruffleObject {
|
||||
|
||||
static final ZoneId UTC = ZoneId.of("UTC");
|
||||
|
||||
static int typeAdjustedIndex(long index, int daySizeInBytes) {
|
||||
public static int typeAdjustedIndex(long index, int daySizeInBytes) {
|
||||
return Math.toIntExact(index * daySizeInBytes);
|
||||
}
|
||||
}
|
||||
|
||||
@ -16,9 +16,8 @@ public final class ArrowFixedArrayInt implements TruffleObject {
|
||||
private final ByteBufferDirect buffer;
|
||||
private final LogicalLayout unit;
|
||||
|
||||
public ArrowFixedArrayInt(ByteBufferDirect buffer, LogicalLayout unit)
|
||||
throws UnsupportedMessageException {
|
||||
this.size = buffer.capacity() / unit.sizeInBytes();
|
||||
public ArrowFixedArrayInt(ByteBufferDirect buffer, int size, LogicalLayout unit) {
|
||||
this.size = size;
|
||||
this.unit = unit;
|
||||
this.buffer = buffer;
|
||||
}
|
||||
|
||||
@ -1,15 +1,13 @@
|
||||
package org.enso.interpreter.arrow.runtime;
|
||||
|
||||
import com.oracle.truffle.api.CompilerDirectives;
|
||||
import com.oracle.truffle.api.dsl.Cached;
|
||||
import com.oracle.truffle.api.dsl.Fallback;
|
||||
import com.oracle.truffle.api.dsl.ImportStatic;
|
||||
import com.oracle.truffle.api.dsl.Specialization;
|
||||
import com.oracle.truffle.api.interop.*;
|
||||
import com.oracle.truffle.api.library.CachedLibrary;
|
||||
import com.oracle.truffle.api.interop.InteropLibrary;
|
||||
import com.oracle.truffle.api.interop.TruffleObject;
|
||||
import com.oracle.truffle.api.interop.UnknownIdentifierException;
|
||||
import com.oracle.truffle.api.interop.UnsupportedMessageException;
|
||||
import com.oracle.truffle.api.interop.UnsupportedTypeException;
|
||||
import com.oracle.truffle.api.library.ExportLibrary;
|
||||
import com.oracle.truffle.api.library.ExportMessage;
|
||||
import java.time.*;
|
||||
import org.enso.interpreter.arrow.LogicalLayout;
|
||||
|
||||
@ExportLibrary(InteropLibrary.class)
|
||||
@ -18,9 +16,11 @@ public final class ArrowFixedSizeArrayBuilder implements TruffleObject {
|
||||
private final LogicalLayout unit;
|
||||
private final int size;
|
||||
private int index;
|
||||
|
||||
private boolean sealed;
|
||||
|
||||
private static final String APPEND_OP = "append";
|
||||
private static final String BUILD_OP = "build";
|
||||
|
||||
public ArrowFixedSizeArrayBuilder(int size, LogicalLayout unit) {
|
||||
this.size = size;
|
||||
this.unit = unit;
|
||||
@ -33,9 +33,16 @@ public final class ArrowFixedSizeArrayBuilder implements TruffleObject {
|
||||
return unit;
|
||||
}
|
||||
|
||||
@ExportMessage
|
||||
public boolean hasArrayElements() {
|
||||
return true;
|
||||
public boolean isSealed() {
|
||||
return sealed;
|
||||
}
|
||||
|
||||
public ByteBufferDirect getBuffer() {
|
||||
return buffer;
|
||||
}
|
||||
|
||||
public int getSize() {
|
||||
return size;
|
||||
}
|
||||
|
||||
@ExportMessage
|
||||
@ -46,8 +53,8 @@ public final class ArrowFixedSizeArrayBuilder implements TruffleObject {
|
||||
@ExportMessage
|
||||
public boolean isMemberInvocable(String member) {
|
||||
return switch (member) {
|
||||
case "append" -> !this.sealed;
|
||||
case "build" -> true;
|
||||
case APPEND_OP -> !this.sealed;
|
||||
case BUILD_OP -> true;
|
||||
default -> false;
|
||||
};
|
||||
}
|
||||
@ -61,25 +68,22 @@ public final class ArrowFixedSizeArrayBuilder implements TruffleObject {
|
||||
Object invokeMember(
|
||||
String name,
|
||||
Object[] args,
|
||||
@Cached.Shared("interop") @CachedLibrary(limit = "1") InteropLibrary iop)
|
||||
@Cached(value = "buildWriterOrNull(name)", neverDefault = true)
|
||||
WriteToBuilderNode writeToBuilderNode)
|
||||
throws UnsupportedMessageException, UnknownIdentifierException, UnsupportedTypeException {
|
||||
switch (name) {
|
||||
case "build":
|
||||
case BUILD_OP:
|
||||
sealed = true;
|
||||
return switch (unit) {
|
||||
case Date32, Date64 -> new ArrowFixedArrayDate(buffer, unit);
|
||||
case Int8, Int16, Int32, Int64 -> new ArrowFixedArrayInt(buffer, unit);
|
||||
case Date32, Date64 -> new ArrowFixedArrayDate(buffer, size, unit);
|
||||
case Int8, Int16, Int32, Int64 -> new ArrowFixedArrayInt(buffer, size, unit);
|
||||
};
|
||||
case "append":
|
||||
case APPEND_OP:
|
||||
if (sealed) {
|
||||
throw UnsupportedMessageException.create();
|
||||
}
|
||||
var current = index;
|
||||
try {
|
||||
iop.writeArrayElement(this, current, args[0]);
|
||||
} catch (InvalidArrayIndexException e) {
|
||||
throw UnsupportedMessageException.create(e);
|
||||
}
|
||||
writeToBuilderNode.executeWrite(this, current, args[0]);
|
||||
index += 1;
|
||||
return NullValue.get();
|
||||
default:
|
||||
@ -87,242 +91,7 @@ public final class ArrowFixedSizeArrayBuilder implements TruffleObject {
|
||||
}
|
||||
}
|
||||
|
||||
@ExportMessage
|
||||
@ImportStatic(LogicalLayout.class)
|
||||
static class ReadArrayElement {
|
||||
|
||||
@Specialization(guards = "receiver.getUnit() == Date32")
|
||||
static Object doDay(ArrowFixedSizeArrayBuilder receiver, long index)
|
||||
throws UnsupportedMessageException {
|
||||
return ArrowFixedArrayDate.readDay(receiver.buffer, index);
|
||||
}
|
||||
|
||||
@Specialization(guards = "receiver.getUnit() == Date64")
|
||||
static Object doMilliseconds(ArrowFixedSizeArrayBuilder receiver, long index)
|
||||
throws UnsupportedMessageException {
|
||||
return ArrowFixedArrayDate.readMilliseconds(receiver.buffer, index);
|
||||
}
|
||||
|
||||
@Specialization(guards = "receiver.getUnit() == Int8")
|
||||
public static Object doByte(ArrowFixedSizeArrayBuilder receiver, long index)
|
||||
throws UnsupportedMessageException, InvalidArrayIndexException {
|
||||
var at =
|
||||
ArrowFixedArrayInt.adjustedIndex(receiver.buffer, receiver.unit, receiver.size, index);
|
||||
if (receiver.buffer.isNull((int) index)) {
|
||||
return NullValue.get();
|
||||
}
|
||||
return receiver.buffer.get(at);
|
||||
}
|
||||
|
||||
@Specialization(guards = "receiver.getUnit() == Int16")
|
||||
public static Object doShort(ArrowFixedSizeArrayBuilder receiver, long index)
|
||||
throws UnsupportedMessageException, InvalidArrayIndexException {
|
||||
var at =
|
||||
ArrowFixedArrayInt.adjustedIndex(receiver.buffer, receiver.unit, receiver.size, index);
|
||||
if (receiver.buffer.isNull((int) index)) {
|
||||
return NullValue.get();
|
||||
}
|
||||
return receiver.buffer.getShort(at);
|
||||
}
|
||||
|
||||
@Specialization(guards = "receiver.getUnit() == Int32")
|
||||
public static Object doInt(ArrowFixedSizeArrayBuilder receiver, long index)
|
||||
throws UnsupportedMessageException, InvalidArrayIndexException {
|
||||
var at =
|
||||
ArrowFixedArrayInt.adjustedIndex(receiver.buffer, receiver.unit, receiver.size, index);
|
||||
if (receiver.buffer.isNull((int) index)) {
|
||||
return NullValue.get();
|
||||
}
|
||||
return receiver.buffer.getInt(at);
|
||||
}
|
||||
|
||||
@Specialization(guards = "receiver.getUnit() == Int64")
|
||||
public static Object doLong(ArrowFixedSizeArrayBuilder receiver, long index)
|
||||
throws UnsupportedMessageException, InvalidArrayIndexException {
|
||||
var at =
|
||||
ArrowFixedArrayInt.adjustedIndex(receiver.buffer, receiver.unit, receiver.size, index);
|
||||
if (receiver.buffer.isNull((int) index)) {
|
||||
return NullValue.get();
|
||||
}
|
||||
return receiver.buffer.getLong(at);
|
||||
}
|
||||
}
|
||||
|
||||
@ExportMessage
|
||||
@ImportStatic(LogicalLayout.class)
|
||||
static class WriteArrayElement {
|
||||
|
||||
@Specialization(guards = "receiver.getUnit() == Date32")
|
||||
static void doDay(
|
||||
ArrowFixedSizeArrayBuilder receiver,
|
||||
long index,
|
||||
Object value,
|
||||
@Cached.Shared("interop") @CachedLibrary(limit = "1") InteropLibrary iop)
|
||||
throws UnsupportedMessageException, UnsupportedTypeException {
|
||||
if (receiver.sealed) {
|
||||
throw UnsupportedMessageException.create();
|
||||
}
|
||||
if (!iop.isDate(value)) {
|
||||
throw UnsupportedTypeException.create(new Object[] {value}, "value is not a date");
|
||||
}
|
||||
var at = ArrowFixedArrayDate.typeAdjustedIndex(index, 4);
|
||||
var time = iop.asDate(value).toEpochDay();
|
||||
receiver.buffer.putInt(at, Math.toIntExact(time));
|
||||
}
|
||||
|
||||
@Specialization(guards = {"receiver.getUnit() == Date64", "!iop.isNull(value)"})
|
||||
static void doMilliseconds(
|
||||
ArrowFixedSizeArrayBuilder receiver,
|
||||
long index,
|
||||
Object value,
|
||||
@Cached.Shared("interop") @CachedLibrary(limit = "1") InteropLibrary iop)
|
||||
throws UnsupportedMessageException, UnsupportedTypeException {
|
||||
if (receiver.sealed) {
|
||||
throw UnsupportedMessageException.create();
|
||||
}
|
||||
if (!iop.isDate(value) || !iop.isTime(value)) {
|
||||
throw UnsupportedTypeException.create(
|
||||
new Object[] {value}, "value is not a date and a time");
|
||||
}
|
||||
|
||||
var at = ArrowFixedArrayDate.typeAdjustedIndex(index, 8);
|
||||
if (iop.isTimeZone(value)) {
|
||||
var zoneDateTimeInstant =
|
||||
instantForZone(
|
||||
iop.asDate(value),
|
||||
iop.asTime(value),
|
||||
iop.asTimeZone(value),
|
||||
ArrowFixedArrayDate.UTC);
|
||||
var secondsPlusNano =
|
||||
zoneDateTimeInstant.getEpochSecond() * ArrowFixedArrayDate.NANO_DIV
|
||||
+ zoneDateTimeInstant.getNano();
|
||||
receiver.buffer.putLong(at, secondsPlusNano);
|
||||
} else {
|
||||
var dateTime = instantForOffset(iop.asDate(value), iop.asTime(value), ZoneOffset.UTC);
|
||||
var secondsPlusNano =
|
||||
dateTime.getEpochSecond() * ArrowFixedArrayDate.NANO_DIV + dateTime.getNano();
|
||||
receiver.buffer.putLong(at, secondsPlusNano);
|
||||
}
|
||||
}
|
||||
|
||||
@CompilerDirectives.TruffleBoundary
|
||||
private static Instant instantForZone(
|
||||
LocalDate date, LocalTime time, ZoneId zone, ZoneId target) {
|
||||
return date.atTime(time).atZone(zone).withZoneSameLocal(target).toInstant();
|
||||
}
|
||||
|
||||
@CompilerDirectives.TruffleBoundary
|
||||
private static Instant instantForOffset(LocalDate date, LocalTime time, ZoneOffset offset) {
|
||||
return date.atTime(time).toInstant(offset);
|
||||
}
|
||||
|
||||
@Specialization(guards = "iop.isNull(value)")
|
||||
static void doNull(
|
||||
ArrowFixedSizeArrayBuilder receiver,
|
||||
long index,
|
||||
Object value,
|
||||
@Cached.Shared("interop") @CachedLibrary(limit = "1") InteropLibrary iop)
|
||||
throws UnsupportedMessageException {
|
||||
if (receiver.sealed) {
|
||||
throw UnsupportedMessageException.create();
|
||||
}
|
||||
receiver.buffer.setNull((int) index);
|
||||
}
|
||||
|
||||
@Specialization(guards = "receiver.getUnit() == Int8")
|
||||
public static void doByte(
|
||||
ArrowFixedSizeArrayBuilder receiver,
|
||||
long index,
|
||||
Object value,
|
||||
@Cached.Shared("interop") @CachedLibrary(limit = "1") InteropLibrary iop)
|
||||
throws UnsupportedMessageException, InvalidArrayIndexException, UnsupportedTypeException {
|
||||
validAccess(receiver, index);
|
||||
if (!iop.fitsInByte(value)) {
|
||||
throw UnsupportedTypeException.create(new Object[] {value}, "value does not fit a byte");
|
||||
}
|
||||
receiver.buffer.put(typeAdjustedIndex(index, receiver.unit), (iop.asByte(value)));
|
||||
}
|
||||
|
||||
@Specialization(guards = "receiver.getUnit() == Int16")
|
||||
public static void doShort(
|
||||
ArrowFixedSizeArrayBuilder receiver,
|
||||
long index,
|
||||
Object value,
|
||||
@Cached.Shared("interop") @CachedLibrary(limit = "1") InteropLibrary iop)
|
||||
throws UnsupportedMessageException, InvalidArrayIndexException, UnsupportedTypeException {
|
||||
validAccess(receiver, index);
|
||||
if (!iop.fitsInShort(value)) {
|
||||
throw UnsupportedTypeException.create(
|
||||
new Object[] {value}, "value does not fit a 2 byte short");
|
||||
}
|
||||
receiver.buffer.putShort(typeAdjustedIndex(index, receiver.unit), (iop.asShort(value)));
|
||||
}
|
||||
|
||||
@Specialization(guards = "receiver.getUnit() == Int32")
|
||||
public static void doInt(
|
||||
ArrowFixedSizeArrayBuilder receiver,
|
||||
long index,
|
||||
int value,
|
||||
@Cached.Shared("interop") @CachedLibrary(limit = "1") InteropLibrary iop)
|
||||
throws UnsupportedMessageException, InvalidArrayIndexException, UnsupportedTypeException {
|
||||
validAccess(receiver, index);
|
||||
if (!iop.fitsInInt(value)) {
|
||||
throw UnsupportedTypeException.create(
|
||||
new Object[] {value}, "value does not fit a 4 byte int");
|
||||
}
|
||||
receiver.buffer.putInt(typeAdjustedIndex(index, receiver.unit), (iop.asInt(value)));
|
||||
}
|
||||
|
||||
@Specialization(guards = "receiver.getUnit() == Int64")
|
||||
public static void doLong(ArrowFixedSizeArrayBuilder receiver, long index, long value)
|
||||
throws UnsupportedMessageException, InvalidArrayIndexException {
|
||||
validAccess(receiver, index);
|
||||
receiver.buffer.putLong(typeAdjustedIndex(index, receiver.unit), value);
|
||||
}
|
||||
|
||||
@Fallback
|
||||
public static void doOther(
|
||||
ArrowFixedSizeArrayBuilder receiver,
|
||||
long index,
|
||||
Object value,
|
||||
@Cached.Shared("interop") @CachedLibrary(limit = "1") InteropLibrary iop)
|
||||
throws UnsupportedTypeException {
|
||||
throw UnsupportedTypeException.create(
|
||||
new Object[] {index, value}, "unknown type of receiver");
|
||||
}
|
||||
|
||||
private static void validAccess(ArrowFixedSizeArrayBuilder receiver, long index)
|
||||
throws InvalidArrayIndexException, UnsupportedMessageException {
|
||||
if (receiver.sealed) {
|
||||
throw UnsupportedMessageException.create();
|
||||
}
|
||||
if (index >= receiver.size || index < 0) {
|
||||
throw InvalidArrayIndexException.create(index);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ExportMessage
|
||||
long getArraySize() {
|
||||
return size;
|
||||
}
|
||||
|
||||
@ExportMessage
|
||||
boolean isArrayElementReadable(long index) {
|
||||
return index >= 0 && index < size && !buffer.isNull((int) index);
|
||||
}
|
||||
|
||||
@ExportMessage
|
||||
boolean isArrayElementModifiable(long index) {
|
||||
return !sealed && index >= 0 && index < size;
|
||||
}
|
||||
|
||||
@ExportMessage
|
||||
boolean isArrayElementInsertable(long index) {
|
||||
return !sealed && index >= 0 && index < size;
|
||||
}
|
||||
|
||||
private static int typeAdjustedIndex(long index, SizeInBytes unit) {
|
||||
return ArrowFixedArrayDate.typeAdjustedIndex(index, unit.sizeInBytes());
|
||||
static WriteToBuilderNode buildWriterOrNull(String op) {
|
||||
return APPEND_OP.equals(op) ? WriteToBuilderNode.build() : WriteToBuilderNodeGen.getUncached();
|
||||
}
|
||||
}
|
||||
|
||||
@ -19,6 +19,7 @@ final class ByteBufferDirect implements AutoCloseable {
|
||||
private ByteBufferDirect(int valueCount, SizeInBytes unit) {
|
||||
var padded = RoundingUtil.forValueCount(valueCount, unit);
|
||||
var buffer = ByteBuffer.allocate(padded.getTotalSizeInBytes());
|
||||
|
||||
this.allocated = buffer;
|
||||
this.dataBuffer = buffer.slice(0, padded.getDataBufferSizeInBytes());
|
||||
this.bitmapBuffer = buffer.slice(dataBuffer.capacity(), padded.getValidityBitmapSizeInBytes());
|
||||
@ -129,7 +130,7 @@ final class ByteBufferDirect implements AutoCloseable {
|
||||
return dataBuffer.getInt(index);
|
||||
}
|
||||
|
||||
public void putInt(int index, int value) throws UnsupportedMessageException {
|
||||
public void putInt(int index, int value) {
|
||||
setValidityBitmap(index, 4);
|
||||
dataBuffer.putInt(index, value);
|
||||
}
|
||||
@ -143,7 +144,7 @@ final class ByteBufferDirect implements AutoCloseable {
|
||||
return dataBuffer.getLong(index);
|
||||
}
|
||||
|
||||
public void putLong(int index, long value) throws UnsupportedMessageException {
|
||||
public void putLong(int index, long value) {
|
||||
setValidityBitmap(index, 8);
|
||||
dataBuffer.putLong(index, value);
|
||||
}
|
||||
|
||||
@ -0,0 +1,217 @@
|
||||
package org.enso.interpreter.arrow.runtime;
|
||||
|
||||
import com.oracle.truffle.api.CompilerDirectives;
|
||||
import com.oracle.truffle.api.dsl.*;
|
||||
import com.oracle.truffle.api.interop.InteropLibrary;
|
||||
import com.oracle.truffle.api.interop.UnsupportedMessageException;
|
||||
import com.oracle.truffle.api.interop.UnsupportedTypeException;
|
||||
import com.oracle.truffle.api.library.CachedLibrary;
|
||||
import com.oracle.truffle.api.nodes.Node;
|
||||
import java.time.Instant;
|
||||
import java.time.LocalDate;
|
||||
import java.time.LocalTime;
|
||||
import java.time.ZoneId;
|
||||
import java.time.ZoneOffset;
|
||||
import org.enso.interpreter.arrow.LogicalLayout;
|
||||
|
||||
@ImportStatic(LogicalLayout.class)
|
||||
@GenerateUncached
|
||||
@GenerateInline(value = false)
|
||||
abstract class WriteToBuilderNode extends Node {
|
||||
|
||||
public abstract void executeWrite(ArrowFixedSizeArrayBuilder receiver, long index, Object value)
|
||||
throws UnsupportedTypeException;
|
||||
|
||||
@NeverDefault
|
||||
static WriteToBuilderNode build() {
|
||||
return WriteToBuilderNodeGen.create();
|
||||
}
|
||||
|
||||
@Specialization(guards = "receiver.getUnit() == Date32")
|
||||
void doWriteDay(
|
||||
ArrowFixedSizeArrayBuilder receiver,
|
||||
long index,
|
||||
Object value,
|
||||
@Cached.Shared("interop") @CachedLibrary(limit = "1") InteropLibrary iop)
|
||||
throws UnsupportedTypeException {
|
||||
validAccess(receiver, index);
|
||||
if (iop.isNull(value)) {
|
||||
receiver.getBuffer().setNull((int) index);
|
||||
return;
|
||||
}
|
||||
if (!iop.isDate(value)) {
|
||||
throw UnsupportedTypeException.create(new Object[] {value}, "value is not a date");
|
||||
}
|
||||
var at = ArrowFixedArrayDate.typeAdjustedIndex(index, 4);
|
||||
long time;
|
||||
try {
|
||||
time = iop.asDate(value).toEpochDay();
|
||||
} catch (UnsupportedMessageException e) {
|
||||
throw UnsupportedTypeException.create(new Object[] {value}, "value is not a date");
|
||||
}
|
||||
receiver.getBuffer().putInt(at, Math.toIntExact(time));
|
||||
}
|
||||
|
||||
@Specialization(guards = {"receiver.getUnit() == Date64"})
|
||||
void doWriteMilliseconds(
|
||||
ArrowFixedSizeArrayBuilder receiver,
|
||||
long index,
|
||||
Object value,
|
||||
@Cached.Shared("interop") @CachedLibrary(limit = "1") InteropLibrary iop)
|
||||
throws UnsupportedTypeException {
|
||||
validAccess(receiver, index);
|
||||
if (iop.isNull(value)) {
|
||||
receiver.getBuffer().setNull((int) index);
|
||||
return;
|
||||
}
|
||||
if (!iop.isDate(value) || !iop.isTime(value)) {
|
||||
throw UnsupportedTypeException.create(new Object[] {value}, "value is not a date and a time");
|
||||
}
|
||||
|
||||
var at = ArrowFixedArrayDate.typeAdjustedIndex(index, 8);
|
||||
if (iop.isTimeZone(value)) {
|
||||
Instant zoneDateTimeInstant;
|
||||
try {
|
||||
zoneDateTimeInstant =
|
||||
instantForZone(
|
||||
iop.asDate(value),
|
||||
iop.asTime(value),
|
||||
iop.asTimeZone(value),
|
||||
ArrowFixedArrayDate.UTC);
|
||||
} catch (UnsupportedMessageException e) {
|
||||
throw UnsupportedTypeException.create(new Object[] {value}, "value is not a date");
|
||||
}
|
||||
var secondsPlusNano =
|
||||
zoneDateTimeInstant.getEpochSecond() * ArrowFixedArrayDate.NANO_DIV
|
||||
+ zoneDateTimeInstant.getNano();
|
||||
receiver.getBuffer().putLong(at, secondsPlusNano);
|
||||
} else {
|
||||
Instant dateTime;
|
||||
try {
|
||||
dateTime = instantForOffset(iop.asDate(value), iop.asTime(value), ZoneOffset.UTC);
|
||||
} catch (UnsupportedMessageException e) {
|
||||
throw UnsupportedTypeException.create(new Object[] {value}, "value is not a date");
|
||||
}
|
||||
var secondsPlusNano =
|
||||
dateTime.getEpochSecond() * ArrowFixedArrayDate.NANO_DIV + dateTime.getNano();
|
||||
receiver.getBuffer().putLong(at, secondsPlusNano);
|
||||
}
|
||||
}
|
||||
|
||||
@CompilerDirectives.TruffleBoundary
|
||||
private static Instant instantForZone(
|
||||
LocalDate date, LocalTime time, ZoneId zone, ZoneId target) {
|
||||
return date.atTime(time).atZone(zone).withZoneSameLocal(target).toInstant();
|
||||
}
|
||||
|
||||
@CompilerDirectives.TruffleBoundary
|
||||
private static Instant instantForOffset(LocalDate date, LocalTime time, ZoneOffset offset) {
|
||||
return date.atTime(time).toInstant(offset);
|
||||
}
|
||||
|
||||
@Specialization(guards = "receiver.getUnit() == Int8")
|
||||
void doWriteByte(
|
||||
ArrowFixedSizeArrayBuilder receiver,
|
||||
long index,
|
||||
Object value,
|
||||
@Cached.Shared("interop") @CachedLibrary(limit = "1") InteropLibrary iop)
|
||||
throws UnsupportedTypeException {
|
||||
validAccess(receiver, index);
|
||||
if (iop.isNull(value)) {
|
||||
receiver.getBuffer().setNull((int) index);
|
||||
return;
|
||||
}
|
||||
if (!iop.fitsInByte(value)) {
|
||||
throw UnsupportedTypeException.create(new Object[] {value}, "value does not fit a byte");
|
||||
}
|
||||
try {
|
||||
receiver.getBuffer().put(typeAdjustedIndex(index, receiver.getUnit()), (iop.asByte(value)));
|
||||
} catch (UnsupportedMessageException e) {
|
||||
throw UnsupportedTypeException.create(new Object[] {value}, "value is not a byte");
|
||||
}
|
||||
}
|
||||
|
||||
@Specialization(guards = "receiver.getUnit() == Int16")
|
||||
void doWriteShort(
|
||||
ArrowFixedSizeArrayBuilder receiver,
|
||||
long index,
|
||||
Object value,
|
||||
@Cached.Shared("interop") @CachedLibrary(limit = "1") InteropLibrary iop)
|
||||
throws UnsupportedTypeException {
|
||||
validAccess(receiver, index);
|
||||
if (iop.isNull(value)) {
|
||||
receiver.getBuffer().setNull((int) index);
|
||||
return;
|
||||
}
|
||||
if (!iop.fitsInShort(value)) {
|
||||
throw UnsupportedTypeException.create(
|
||||
new Object[] {value}, "value does not fit a 2 byte short");
|
||||
}
|
||||
try {
|
||||
receiver
|
||||
.getBuffer()
|
||||
.putShort(typeAdjustedIndex(index, receiver.getUnit()), (iop.asShort(value)));
|
||||
} catch (UnsupportedMessageException e) {
|
||||
throw UnsupportedTypeException.create(new Object[] {value}, "value is not a short");
|
||||
}
|
||||
}
|
||||
|
||||
@Specialization(guards = "receiver.getUnit() == Int32")
|
||||
void doWriteInt(
|
||||
ArrowFixedSizeArrayBuilder receiver,
|
||||
long index,
|
||||
int value,
|
||||
@Cached.Shared("interop") @CachedLibrary(limit = "1") InteropLibrary iop)
|
||||
throws UnsupportedTypeException {
|
||||
validAccess(receiver, index);
|
||||
if (iop.isNull(value)) {
|
||||
receiver.getBuffer().setNull((int) index);
|
||||
return;
|
||||
}
|
||||
if (!iop.fitsInInt(value)) {
|
||||
throw UnsupportedTypeException.create(
|
||||
new Object[] {value}, "value does not fit a 4 byte int");
|
||||
}
|
||||
try {
|
||||
receiver.getBuffer().putInt(typeAdjustedIndex(index, receiver.getUnit()), (iop.asInt(value)));
|
||||
} catch (UnsupportedMessageException e) {
|
||||
throw UnsupportedTypeException.create(new Object[] {value}, "value is not an int");
|
||||
}
|
||||
}
|
||||
|
||||
@Specialization(guards = "receiver.getUnit() == Int64")
|
||||
public static void doWriteLong(
|
||||
ArrowFixedSizeArrayBuilder receiver,
|
||||
long index,
|
||||
long value,
|
||||
@Cached.Shared("interop") @CachedLibrary(limit = "1") InteropLibrary iop)
|
||||
throws UnsupportedTypeException {
|
||||
validAccess(receiver, index);
|
||||
if (iop.isNull(value)) {
|
||||
receiver.getBuffer().setNull((int) index);
|
||||
return;
|
||||
}
|
||||
receiver.getBuffer().putLong(typeAdjustedIndex(index, receiver.getUnit()), value);
|
||||
}
|
||||
|
||||
@Fallback
|
||||
void doWriteOther(ArrowFixedSizeArrayBuilder receiver, long index, Object value)
|
||||
throws UnsupportedTypeException {
|
||||
throw UnsupportedTypeException.create(new Object[] {index, value}, "unknown type of receiver");
|
||||
}
|
||||
|
||||
private static void validAccess(ArrowFixedSizeArrayBuilder receiver, long index)
|
||||
throws UnsupportedTypeException {
|
||||
if (receiver.isSealed()) {
|
||||
throw UnsupportedTypeException.create(
|
||||
new Object[] {receiver}, "receiver is not an unsealed buffer");
|
||||
}
|
||||
if (index >= receiver.getSize() || index < 0) {
|
||||
throw UnsupportedTypeException.create(new Object[] {index}, "index is out of range");
|
||||
}
|
||||
}
|
||||
|
||||
private static int typeAdjustedIndex(long index, SizeInBytes unit) {
|
||||
return ArrowFixedArrayDate.typeAdjustedIndex(index, unit.sizeInBytes());
|
||||
}
|
||||
}
|
||||
@ -8,6 +8,7 @@ import java.time.ZonedDateTime;
|
||||
import java.time.temporal.Temporal;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.Set;
|
||||
import java.util.logging.Formatter;
|
||||
import java.util.logging.Handler;
|
||||
import java.util.logging.LogRecord;
|
||||
@ -55,11 +56,12 @@ public class VerifyArrowTest {
|
||||
assertNotNull("Arrow is available", arrow);
|
||||
var date32Constr = ctx.eval("arrow", "new[Date32]");
|
||||
|
||||
Value date32ArrayBuilder = date32Constr.newInstance(10);
|
||||
var arrLength = 10;
|
||||
Value date32ArrayBuilder = date32Constr.newInstance(arrLength);
|
||||
assertNotNull("allocated value should not be null", date32ArrayBuilder);
|
||||
assertTrue("allocated value should be an array", date32ArrayBuilder.hasArrayElements());
|
||||
assertTrue("allocated value should not be an array", !date32ArrayBuilder.hasArrayElements());
|
||||
var startDate = LocalDate.now();
|
||||
populateArrayWithConsecutiveDays(date32ArrayBuilder, startDate);
|
||||
populateBuilderWithConsecutiveDays(date32ArrayBuilder, startDate, arrLength, Set.of());
|
||||
Value date32Array = date32ArrayBuilder.invokeMember("build");
|
||||
var rawDayPlus2 = date32Array.getArrayElement(2);
|
||||
var dayPlus2 = rawDayPlus2.asDate();
|
||||
@ -68,7 +70,7 @@ public class VerifyArrowTest {
|
||||
|
||||
date32ArrayBuilder = date32Constr.newInstance(10);
|
||||
var startDateTime = ZonedDateTime.now();
|
||||
populateArrayWithConsecutiveDays(date32ArrayBuilder, startDateTime);
|
||||
populateBuilderWithConsecutiveDays(date32ArrayBuilder, startDateTime, arrLength, Set.of());
|
||||
date32Array = date32ArrayBuilder.invokeMember("build");
|
||||
rawDayPlus2 = date32Array.getArrayElement(2);
|
||||
assertFalse(rawDayPlus2.isTime());
|
||||
@ -93,12 +95,13 @@ public class VerifyArrowTest {
|
||||
assertNotNull("Arrow is available", arrow);
|
||||
var date64Constr = ctx.eval("arrow", "new[Date64]");
|
||||
|
||||
Value date64ArrayBuilder = date64Constr.newInstance(10);
|
||||
var arrLength = 10;
|
||||
Value date64ArrayBuilder = date64Constr.newInstance(arrLength);
|
||||
assertNotNull("allocated value should not be null", date64ArrayBuilder);
|
||||
assertTrue("allocated value should be an array", date64ArrayBuilder.hasArrayElements());
|
||||
assertTrue("allocated value should not be an array", !date64ArrayBuilder.hasArrayElements());
|
||||
var startDate = ZonedDateTime.now(ZoneId.of("Europe/Paris"));
|
||||
var startDateZone = startDate.getZone();
|
||||
populateArrayWithConsecutiveDays(date64ArrayBuilder, startDate);
|
||||
populateBuilderWithConsecutiveDays(date64ArrayBuilder, startDate, arrLength, Set.of());
|
||||
Value date64Array = date64ArrayBuilder.invokeMember("build");
|
||||
var rawZonedDateTime = date64Array.getArrayElement(2);
|
||||
var dayPlus2 =
|
||||
@ -111,9 +114,7 @@ public class VerifyArrowTest {
|
||||
var startDate2 = ZonedDateTime.parse("2023-11-01T02:00:01+01:00[Europe/Paris]");
|
||||
var startDate2Zone = startDate2.getZone();
|
||||
var startDate2Pnf = ZonedDateTime.parse("2023-11-01T02:00:01-07:00[US/Pacific]");
|
||||
populateArrayWithConsecutiveDays(date64ArrayBuilder, startDate2);
|
||||
date64ArrayBuilder.setArrayElement(5, null);
|
||||
date64ArrayBuilder.setArrayElement(9, null);
|
||||
populateBuilderWithConsecutiveDays(date64ArrayBuilder, startDate2, arrLength, Set.of(5, 9));
|
||||
assertTrue(date64ArrayBuilder.canInvokeMember("append"));
|
||||
date64Array = date64ArrayBuilder.invokeMember("build");
|
||||
rawZonedDateTime = date64Array.getArrayElement(2);
|
||||
@ -137,18 +138,19 @@ public class VerifyArrowTest {
|
||||
var int8Constr = ctx.eval("arrow", "new[Int8]");
|
||||
assertNotNull(int8Constr);
|
||||
|
||||
Value int8ArrayBuilder = int8Constr.newInstance(10);
|
||||
var arrLength = 10;
|
||||
Value int8ArrayBuilder = int8Constr.newInstance(arrLength);
|
||||
assertNotNull(int8ArrayBuilder);
|
||||
populateIntArray(int8ArrayBuilder, (byte) 42, int8ArrayBuilder.getArraySize() - 1);
|
||||
populateIntBuilder(int8ArrayBuilder, (byte) 42, arrLength - 1);
|
||||
assertThrows(RuntimeException.class, () -> int8ArrayBuilder.setArrayElement(5, 300));
|
||||
assertThrows(RuntimeException.class, () -> int8ArrayBuilder.invokeMember("append", 300));
|
||||
int8ArrayBuilder.setArrayElement(5, 4);
|
||||
assertThrows(UnsupportedOperationException.class, () -> int8ArrayBuilder.getArrayElement(5));
|
||||
var int8Array = int8ArrayBuilder.invokeMember("build");
|
||||
var v = int8Array.getArrayElement(5);
|
||||
assertEquals((byte) 4, v.asByte());
|
||||
assertEquals((byte) 5, v.asByte());
|
||||
assertThrows(UnsupportedOperationException.class, () -> int8Array.setArrayElement(5, 21));
|
||||
v = int8Array.getArrayElement(5);
|
||||
assertEquals((byte) 4, v.asByte());
|
||||
assertEquals((byte) 5, v.asByte());
|
||||
}
|
||||
|
||||
@Test
|
||||
@ -262,21 +264,18 @@ public class VerifyArrowTest {
|
||||
}
|
||||
}
|
||||
|
||||
private void populateArrayWithConsecutiveDays(Value arr, Temporal startDate) {
|
||||
var len = arr.getArraySize();
|
||||
private void populateBuilderWithConsecutiveDays(
|
||||
Value builder, Temporal startDate, int len, Set<Integer> skip) {
|
||||
for (int i = 0; i < len; i++) {
|
||||
arr.invokeMember("append", startDate.plus(2, java.time.temporal.ChronoUnit.DAYS));
|
||||
var date = skip.contains(i) ? null : startDate.plus(i, java.time.temporal.ChronoUnit.DAYS);
|
||||
builder.invokeMember("append", date);
|
||||
}
|
||||
}
|
||||
|
||||
private void populateIntArray(Value arr, byte startValue) {
|
||||
populateIntArray(arr, startValue, arr.getArraySize());
|
||||
}
|
||||
|
||||
private void populateIntArray(Value arr, byte startValue, long until) {
|
||||
private void populateIntBuilder(Value builder, byte startValue, long until) {
|
||||
for (int i = 0; i < until; i++) {
|
||||
var v = startValue + i;
|
||||
arr.setArrayElement(i, (byte) v);
|
||||
builder.invokeMember("append", i, (byte) v);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -19,17 +19,19 @@ add_specs suite_builder =
|
||||
arrow_array_constructor = new_arrow
|
||||
builder = arrow_array_constructor.new 10
|
||||
|
||||
builder.length . should_equal 10
|
||||
builder.at 1 . should_equal Nothing
|
||||
builder.at 10 . should_fail_with Index_Out_Of_Bounds
|
||||
Test.expect_panic_with (builder.length) No_Such_Method
|
||||
Test.expect_panic_with (builder.at 1) No_Such_Method
|
||||
builder.append 2
|
||||
builder.append 10
|
||||
builder.at 1 . should_equal 10
|
||||
builder.append 127
|
||||
|
||||
# Check that the value that will not fit into a byte is rejected
|
||||
Test.expect_panic_with (builder.append 128) Unsupported_Argument_Types
|
||||
v = builder.build
|
||||
v.length . should_equal 10
|
||||
v.at 1 . should_equal 10
|
||||
v.at 11 . should_fail_with Index_Out_Of_Bounds
|
||||
v.at 25 . should_fail_with Index_Out_Of_Bounds # Capacity of data buffer after padding will be 24
|
||||
v.at 2 . should_equal 127
|
||||
Test.expect_panic_with (builder.append 21) No_Such_Method
|
||||
v.at 4 . should_equal Nothing
|
||||
|
||||
Loading…
Reference in New Issue
Block a user