Skip to content

Commit d99dff2

Browse files
Using Arrow Date32 format to transfer LocalDate across JVMs (#14632)
- #14607 reveals a problem when running ```ruby from Standard.Base import all from Standard.Table import all from Standard.Tableau import all main = node1 = Data.read '/tmp/x.hyper' node2 = node1.read 'Extract' node3 = node2.set (expr 'year([Date])') "Year" node3.get "Year" . to_vector ``` in _dual JVM_ (mock) mode: ``` sbt:enso> runEngineDistribution --vm.D=polyglot.enso.classLoading=Standard.Tableau:guest,hosted --run hyper.enso Error:Argument_Mismatch Unsupported type: class com.oracle.truffle.polyglot.PolyglotMapAndFunction (expected date/time type). ``` - builds upon #14309 - and upon #14471
1 parent 3f1361c commit d99dff2

File tree

4 files changed

+148
-1
lines changed

4 files changed

+148
-1
lines changed

std-bits/table/src/main/java/org/enso/table/data/column/builder/Builder.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,7 @@ static <T> ColumnStorage<T> makeLocal(ColumnStorage<T> storage) {
121121
DoubleBuilder.fromAddress(size, data, validity, type).seal(storage, type);
122122
case TextType type ->
123123
StringBuilder.fromAddress(size, data, validity, type).seal(storage, type);
124+
case DateType type -> DateBuilder.fromAddress(size, data, validity).seal(storage, type);
124125
default -> storage;
125126
};
126127
assert assertSameStorages(storage, localStorage);

std-bits/table/src/main/java/org/enso/table/data/column/builder/DateBuilder.java

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,12 @@
11
package org.enso.table.data.column.builder;
22

3+
import java.lang.foreign.MemorySegment;
4+
import java.nio.ByteOrder;
35
import java.time.LocalDate;
6+
import java.util.BitSet;
47
import java.util.Objects;
58
import org.enso.table.data.column.storage.ColumnStorage;
9+
import org.enso.table.data.column.storage.Storage;
610
import org.enso.table.data.column.storage.TypedStorage;
711
import org.enso.table.data.column.storage.type.DateTimeType;
812
import org.enso.table.data.column.storage.type.DateType;
@@ -18,6 +22,28 @@ final class DateBuilder extends TypedBuilder<LocalDate> {
1822
this.allowDateToDateTimeConversion = allowDateToDateTimeConversion;
1923
}
2024

25+
static DateBuilder fromAddress(int size, long data, long validity) {
26+
var validityBuffer =
27+
MemorySegment.ofAddress(validity).reinterpret((size + 7) / 8).asByteBuffer();
28+
var bits = BitSet.valueOf(validityBuffer);
29+
var buf =
30+
MemorySegment.ofAddress(data)
31+
.reinterpret(Integer.BYTES * size)
32+
.asByteBuffer()
33+
.order(ByteOrder.LITTLE_ENDIAN);
34+
35+
var b = new DateBuilder(size, false);
36+
for (var i = 0; i < size; i++) {
37+
var day = buf.getInt();
38+
if (bits.get(i)) {
39+
b.append(LocalDate.ofEpochDay(day));
40+
} else {
41+
b.appendNulls(1);
42+
}
43+
}
44+
return b;
45+
}
46+
2147
@Override
2248
public DateBuilder append(Object o) {
2349
ensureSpaceToAppend();
@@ -40,7 +66,11 @@ public boolean accepts(Object o) {
4066

4167
@Override
4268
protected ColumnStorage<LocalDate> doSeal() {
43-
return new TypedStorage<>(DateType.INSTANCE, data);
69+
return seal(null, DateType.INSTANCE);
70+
}
71+
72+
final Storage<LocalDate> seal(ColumnStorage<?> otherStorage, DateType type) {
73+
return new TypedStorage<>(type, data, otherStorage);
4474
}
4575

4676
@Override

std-bits/table/src/main/java/org/enso/table/data/column/storage/TypedStorage.java

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,11 @@
44
import java.nio.ByteBuffer;
55
import java.nio.ByteOrder;
66
import java.nio.charset.StandardCharsets;
7+
import java.time.LocalDate;
78
import java.util.Arrays;
89
import java.util.BitSet;
910
import java.util.Iterator;
11+
import org.enso.table.data.column.storage.type.DateType;
1012
import org.enso.table.data.column.storage.type.StorageType;
1113
import org.enso.table.data.column.storage.type.TextType;
1214
import org.enso.table.util.ImmutableBitSet;
@@ -70,6 +72,28 @@ public long addressOfData() {
7072
offheapBuffer = buf;
7173
validitySet = new ImmutableBitSet(validity, data.length);
7274
}
75+
if (offheapBuffer == null && getType() instanceof DateType) {
76+
var fullSize = data.length * Integer.BYTES;
77+
var buf = ByteBuffer.allocateDirect(fullSize).order(ByteOrder.LITTLE_ENDIAN);
78+
var validity = new BitSet();
79+
var at = 0;
80+
for (var value : data) {
81+
if (value instanceof LocalDate s) {
82+
buf.putInt(Math.toIntExact(s.toEpochDay()));
83+
validity.set(at, true);
84+
} else {
85+
buf.putInt(0);
86+
validity.set(at, false);
87+
}
88+
at++;
89+
}
90+
assert buf.limit() == buf.position();
91+
buf.flip();
92+
assert buf.position() == 0;
93+
assert buf.limit() == fullSize;
94+
offheapBuffer = buf;
95+
validitySet = new ImmutableBitSet(validity, data.length);
96+
}
7397
if (offheapBuffer != null) {
7498
return MemorySegment.ofBuffer(offheapBuffer).address();
7599
}
Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
package org.enso.base.polyglot.tests;
2+
3+
import static org.junit.Assert.assertEquals;
4+
import static org.junit.Assert.assertNotEquals;
5+
import static org.junit.Assert.assertNotSame;
6+
7+
import java.time.LocalDate;
8+
import java.util.Objects;
9+
import java.util.Random;
10+
import java.util.stream.IntStream;
11+
import org.enso.table.data.column.builder.Builder;
12+
import org.enso.table.problems.BlackholeProblemAggregator;
13+
import org.enso.table.problems.ProblemAggregator;
14+
import org.enso.test.utils.ContextUtils;
15+
import org.junit.BeforeClass;
16+
import org.junit.ClassRule;
17+
import org.junit.Test;
18+
19+
public class DateStorageTest {
20+
@ClassRule
21+
public static final ContextUtils ctx =
22+
ContextUtils.newBuilder("enso", "arrow").assertGC(false).build();
23+
24+
@BeforeClass
25+
public static void importAll() {
26+
ctx.eval("enso", "from Standard.Base import all");
27+
}
28+
29+
@Test
30+
public void makeLocalFromDateStorage() {
31+
var b = Builder.getForDate(3);
32+
var one = LocalDate.of(1973, 12, 10);
33+
var two = LocalDate.of(1975, 5, 3);
34+
b.append(one).appendNulls(1).append(two);
35+
var storage = b.seal();
36+
var localStorage = Builder.makeLocal(storage);
37+
assertNotSame("local storage is a copy of storage", storage, localStorage);
38+
assertEquals("They have the same size", storage.getSize(), localStorage.getSize());
39+
assertEquals("They have the same type", storage.getType(), localStorage.getType());
40+
for (var i = 0L; i < storage.getSize(); i++) {
41+
var elem = storage.getItemBoxed(i);
42+
var localElem = localStorage.getItemBoxed(i);
43+
assertEquals("At " + i, elem, localElem);
44+
}
45+
}
46+
47+
@Test
48+
public void testCreateViaBuilderAndReadViaArrowSimple16() {
49+
generateAndCompare("Simple 16 values", 16, IntStream.range(0, 16));
50+
}
51+
52+
@Test
53+
public void testCreateViaBuilderAndReadViaArrowRandom() {
54+
generateAndCompareWithSeed(System.currentTimeMillis());
55+
}
56+
57+
private void generateAndCompareWithSeed(long seed) {
58+
var r = new Random(seed);
59+
var size = r.nextInt(256, 4096);
60+
var stream = r.ints(size, 0, 20000);
61+
var msg = "with seed " + seed + " size " + size;
62+
System.err.println(msg);
63+
generateAndCompare(msg, size, stream);
64+
}
65+
66+
private void generateAndCompare(String info, int size, IntStream r) {
67+
var sb = new StringBuilder();
68+
var b = Builder.getForDate(size);
69+
r.mapToObj(LocalDate::ofEpochDay).forEach(b::append);
70+
var storage = b.seal();
71+
assertEquals("Storage has the right size: " + storage, size, storage.getSize());
72+
assertNotEquals("Storage provides acccess to raw data", 0L, storage.addressOfData());
73+
assertNotEquals("Storage provides access to validity bitmap", 0L, storage.addressOfValidity());
74+
75+
var arr =
76+
ctx.eval("arrow", "cast[Date32]")
77+
.execute(storage.addressOfData(), storage.getSize(), storage.addressOfValidity());
78+
for (var i = 0L; i < size; i++) {
79+
var elem0 = storage.getItemBoxed(i);
80+
var value1 = arr.getArrayElement(i);
81+
var elem1 = value1.isNull() ? null : value1.asDate();
82+
if (!Objects.equals(elem0, elem1)) {
83+
sb.append("\n at ").append(i).append(" ").append(elem0).append(" != ").append(elem1);
84+
}
85+
}
86+
assertEquals(info + "\n" + sb.toString(), 0, sb.length());
87+
}
88+
89+
private static ProblemAggregator problemAggregator() {
90+
return BlackholeProblemAggregator.INSTANCE;
91+
}
92+
}

0 commit comments

Comments
 (0)