Skip to content

Commit

Permalink
Arrow builder is not an Array (#9358)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
hubertp authored Mar 13, 2024
1 parent 2322b40 commit f82e802
Show file tree
Hide file tree
Showing 8 changed files with 304 additions and 307 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
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 {
Expand Down Expand Up @@ -43,7 +44,8 @@ static Object doDate32(
@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")
Expand All @@ -53,7 +55,8 @@ static Object doDate64(
@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")
Expand All @@ -63,7 +66,8 @@ static Object doInt8(
@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")
Expand All @@ -73,7 +77,8 @@ static Object doInt16(
@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")
Expand All @@ -83,7 +88,8 @@ static Object doInt32(
@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")
Expand All @@ -93,11 +99,13 @@ static Object doInt64(
@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);
Expand All @@ -117,10 +125,13 @@ private static ByteBufferDirect pointer(Object[] args, InteropLibrary interop, S
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);
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,8 @@ public ArrowFixedArrayDate(int size, LogicalLayout unit) {
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;
}
Expand Down Expand Up @@ -164,7 +163,7 @@ private static ZonedDateTime zonedDateTimeFromSeconds(long seconds, long nano, Z

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);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
Loading

0 comments on commit f82e802

Please sign in to comment.