diff --git a/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/ArrowCastToFixedSizeArrayFactory.java b/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/ArrowCastToFixedSizeArrayFactory.java index e568cf7491..a3e54d5821 100644 --- a/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/ArrowCastToFixedSizeArrayFactory.java +++ b/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/ArrowCastToFixedSizeArrayFactory.java @@ -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 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); } } diff --git a/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/ArrowFixedArrayDate.java b/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/ArrowFixedArrayDate.java index f400afadcf..af848b8b87 100644 --- a/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/ArrowFixedArrayDate.java +++ b/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/ArrowFixedArrayDate.java @@ -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); } } diff --git a/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/ArrowFixedArrayInt.java b/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/ArrowFixedArrayInt.java index bba66eddb9..efd263b9bd 100644 --- a/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/ArrowFixedArrayInt.java +++ b/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/ArrowFixedArrayInt.java @@ -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; } diff --git a/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/ArrowFixedSizeArrayBuilder.java b/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/ArrowFixedSizeArrayBuilder.java index 1cec15ec17..8ad0315941 100644 --- a/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/ArrowFixedSizeArrayBuilder.java +++ b/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/ArrowFixedSizeArrayBuilder.java @@ -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(); } } diff --git a/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/ByteBufferDirect.java b/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/ByteBufferDirect.java index d840a67bed..a5462418b8 100644 --- a/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/ByteBufferDirect.java +++ b/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/ByteBufferDirect.java @@ -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); } diff --git a/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/WriteToBuilderNode.java b/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/WriteToBuilderNode.java new file mode 100644 index 0000000000..acff598d0b --- /dev/null +++ b/engine/runtime-language-arrow/src/main/java/org/enso/interpreter/arrow/runtime/WriteToBuilderNode.java @@ -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()); + } +} diff --git a/engine/runtime-language-arrow/src/test/java/org/enso/interpreter/arrow/VerifyArrowTest.java b/engine/runtime-language-arrow/src/test/java/org/enso/interpreter/arrow/VerifyArrowTest.java index 62a93997df..7578a5f2cd 100644 --- a/engine/runtime-language-arrow/src/test/java/org/enso/interpreter/arrow/VerifyArrowTest.java +++ b/engine/runtime-language-arrow/src/test/java/org/enso/interpreter/arrow/VerifyArrowTest.java @@ -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 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); } } diff --git a/test/Base_Tests/src/Semantic/Arrow_Interop_Spec.enso b/test/Base_Tests/src/Semantic/Arrow_Interop_Spec.enso index 14192b2b1b..4d995e3513 100644 --- a/test/Base_Tests/src/Semantic/Arrow_Interop_Spec.enso +++ b/test/Base_Tests/src/Semantic/Arrow_Interop_Spec.enso @@ -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