Skip to content

Commit 6fc84c7

Browse files
dfa1dfa1
committed
experiment: passing arena top-down
1 parent 0fb136d commit 6fc84c7

3 files changed

Lines changed: 296 additions & 274 deletions

File tree

‎integration/src/test/java/io/github/dfa1/vortex/integration/RustWritesJavaReadsIntegrationTest.java‎

Lines changed: 163 additions & 164 deletions
Original file line numberDiff line numberDiff line change
@@ -34,168 +34,167 @@
3434
/// Cross-compatibility: Rust (JNI) writer → Java reader.
3535
class RustWritesJavaReadsIntegrationTest {
3636

37-
static {
38-
NativeLoader.loadJni();
39-
}
40-
41-
private static final Session SESSION = Session.create();
42-
private static final BufferAllocator ALLOCATOR = ArrowAllocation.rootAllocator();
43-
44-
private static final Schema JNI_SCHEMA = new Schema(List.of(
45-
Field.notNullable("id", new ArrowType.Int(64, true)),
46-
Field.notNullable("value", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE))
47-
));
48-
49-
@Test
50-
void jniWriter_javaReader_singleChunk(@TempDir Path tmp) throws IOException {
51-
// Given
52-
Path file = tmp.resolve("jni_single.vtx");
53-
long[] ids = {1L, 2L, 3L};
54-
double[] vals = {1.1, 2.2, 3.3};
55-
writeJni(file, ids, vals);
56-
57-
// When / Then
58-
try (var vf = VortexReader.open(file, CodecRegistry.loadAll())) {
59-
List<ScanResult> results = scanAll(vf);
60-
assertThat(results).hasSize(1);
61-
assertThat(results.get(0).rowCount()).isEqualTo(3L);
62-
assertThat(toLongs(results.get(0))).containsExactly(1L, 2L, 3L);
63-
assertThat(toDoubles(results.get(0))).containsExactly(1.1, 2.2, 3.3);
64-
}
65-
}
66-
67-
@Test
68-
void jniWriter_javaReader_multipleChunks(@TempDir Path tmp) throws IOException {
69-
// Given
70-
Path file = tmp.resolve("jni_multi.vtx");
71-
String uri = file.toAbsolutePath().toUri().toString();
72-
try (VortexWriter writer = VortexWriter.create(SESSION, uri, JNI_SCHEMA, new HashMap<>(), ALLOCATOR)) {
73-
flushBatch(writer, new long[]{1L, 2L}, new double[]{1.1, 2.2});
74-
flushBatch(writer, new long[]{3L, 4L, 5L}, new double[]{3.3, 4.4, 5.5});
75-
}
76-
77-
// When / Then — JNI may merge small batches; verify total rows and values
78-
try (var vf = VortexReader.open(file, CodecRegistry.loadAll())) {
79-
List<ScanResult> results = scanAll(vf);
80-
long totalRows = results.stream().mapToLong(ScanResult::rowCount).sum();
81-
assertThat(totalRows).isEqualTo(5L);
82-
long[] allIds = results.stream()
83-
.flatMapToLong(r -> java.util.Arrays.stream(toLongs(r)))
84-
.toArray();
85-
assertThat(allIds).containsExactly(1L, 2L, 3L, 4L, 5L);
86-
}
87-
}
88-
89-
@Test
90-
void jniWriter_javaReader_columnProjection(@TempDir Path tmp) throws IOException {
91-
// Given
92-
Path file = tmp.resolve("jni_proj.vtx");
93-
writeJni(file, new long[]{10L, 20L}, new double[]{0.1, 0.2});
94-
95-
// When / Then
96-
try (var vf = VortexReader.open(file, CodecRegistry.loadAll())) {
97-
List<ScanResult> results = scanAll(vf, io.github.dfa1.vortex.scan.ScanOptions.columns("id"));
98-
assertThat(results).hasSize(1);
99-
assertThat(results.get(0).columns()).containsKey("id");
100-
assertThat(results.get(0).columns()).doesNotContainKey("value");
101-
assertThat(toLongs(results.get(0))).containsExactly(10L, 20L);
102-
}
103-
}
104-
105-
@Test
106-
void jniWriter_javaReader_fewUniqueF64Values(@TempDir Path tmp) throws IOException {
107-
// Given — 10_000 rows cycling through only 3 unique F64 values to trigger dict encoding
108-
int n = 10_000;
109-
long[] ids = new long[n];
110-
double[] vals = new double[n];
111-
double[] unique = {1.1, 2.2, 3.3};
112-
for (int i = 0; i < n; i++) {
113-
ids[i] = i;
114-
vals[i] = unique[i % unique.length];
115-
}
116-
Path file = tmp.resolve("jni_dict.vtx");
117-
writeJni(file, ids, vals);
118-
119-
// When / Then
120-
try (var vf = VortexReader.open(file, CodecRegistry.loadAll())) {
121-
List<ScanResult> results = scanAll(vf, io.github.dfa1.vortex.scan.ScanOptions.columns("value"));
122-
long total = results.stream().mapToLong(ScanResult::rowCount).sum();
123-
assertThat(total).isEqualTo(n);
124-
var layout = ValueLayout.JAVA_DOUBLE_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
125-
double sum = 0;
126-
for (ScanResult r : results) {
127-
var arr = r.columns().get("value");
128-
for (long j = 0; j < arr.length(); j++) {
129-
sum += arr.buffer(0).get(layout, j * Double.BYTES);
130-
}
131-
}
132-
// 10_000 rows: 3333 full cycles of [1.1,2.2,3.3] (=6.6 each) + one 1.1 remainder
133-
assertThat(sum).isCloseTo(21_998.9, org.assertj.core.data.Offset.offset(0.1));
134-
}
135-
}
136-
137-
// ── JNI write helpers ─────────────────────────────────────────────────────
138-
139-
private static void writeJni(Path file, long[] ids, double[] vals) throws IOException {
140-
String uri = file.toAbsolutePath().toUri().toString();
141-
try (VortexWriter writer = VortexWriter.create(SESSION, uri, JNI_SCHEMA, new HashMap<>(), ALLOCATOR)) {
142-
flushBatch(writer, ids, vals);
143-
}
144-
}
145-
146-
private static void flushBatch(VortexWriter writer, long[] ids, double[] vals) throws IOException {
147-
int n = ids.length;
148-
try (VectorSchemaRoot root = VectorSchemaRoot.create(JNI_SCHEMA, ALLOCATOR)) {
149-
BigIntVector idVec = (BigIntVector) root.getVector("id");
150-
Float8Vector valVec = (Float8Vector) root.getVector("value");
151-
idVec.allocateNew(n);
152-
valVec.allocateNew(n);
153-
for (int i = 0; i < n; i++) {
154-
idVec.setSafe(i, ids[i]);
155-
valVec.setSafe(i, vals[i]);
156-
}
157-
root.setRowCount(n);
158-
try (ArrowArray arr = ArrowArray.allocateNew(ALLOCATOR);
159-
ArrowSchema schema = ArrowSchema.allocateNew(ALLOCATOR)) {
160-
Data.exportVectorSchemaRoot(ALLOCATOR, root, null, arr, schema);
161-
writer.writeBatch(arr.memoryAddress(), schema.memoryAddress());
162-
}
163-
}
164-
}
165-
166-
// ── Java read helpers ─────────────────────────────────────────────────────
167-
168-
private static List<ScanResult> scanAll(VortexReader vf) throws IOException {
169-
return scanAll(vf, io.github.dfa1.vortex.scan.ScanOptions.all());
170-
}
171-
172-
private static List<ScanResult> scanAll(VortexReader vf,
173-
io.github.dfa1.vortex.scan.ScanOptions opts) throws IOException {
174-
var results = new ArrayList<ScanResult>();
175-
var iter = vf.scan(opts);
176-
while (iter.hasNext()) {
177-
results.add(iter.next());
178-
}
179-
return results;
180-
}
181-
182-
private static long[] toLongs(ScanResult chunk) {
183-
var layout = ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
184-
var arr = chunk.columns().get("id");
185-
long[] out = new long[(int) arr.length()];
186-
for (int i = 0; i < out.length; i++) {
187-
out[i] = arr.buffer(0).get(layout, (long) i * Long.BYTES);
188-
}
189-
return out;
190-
}
191-
192-
private static double[] toDoubles(ScanResult chunk) {
193-
var layout = ValueLayout.JAVA_DOUBLE_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
194-
var arr = chunk.columns().get("value");
195-
double[] out = new double[(int) arr.length()];
196-
for (int i = 0; i < out.length; i++) {
197-
out[i] = arr.buffer(0).get(layout, (long) i * Double.BYTES);
198-
}
199-
return out;
200-
}
37+
private static final Session SESSION = Session.create();
38+
private static final BufferAllocator ALLOCATOR = ArrowAllocation.rootAllocator();
39+
private static final Schema JNI_SCHEMA = new Schema(List.of(
40+
Field.notNullable("id", new ArrowType.Int(64, true)),
41+
Field.notNullable("value", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE))
42+
));
43+
44+
static {
45+
NativeLoader.loadJni();
46+
}
47+
48+
private static void writeJni(Path file, long[] ids, double[] vals) throws IOException {
49+
String uri = file.toAbsolutePath().toUri().toString();
50+
try (VortexWriter writer = VortexWriter.create(SESSION, uri, JNI_SCHEMA, new HashMap<>(), ALLOCATOR)) {
51+
flushBatch(writer, ids, vals);
52+
}
53+
}
54+
55+
private static void flushBatch(VortexWriter writer, long[] ids, double[] vals) throws IOException {
56+
int n = ids.length;
57+
try (VectorSchemaRoot root = VectorSchemaRoot.create(JNI_SCHEMA, ALLOCATOR)) {
58+
BigIntVector idVec = (BigIntVector) root.getVector("id");
59+
Float8Vector valVec = (Float8Vector) root.getVector("value");
60+
idVec.allocateNew(n);
61+
valVec.allocateNew(n);
62+
for (int i = 0; i < n; i++) {
63+
idVec.setSafe(i, ids[i]);
64+
valVec.setSafe(i, vals[i]);
65+
}
66+
root.setRowCount(n);
67+
try (ArrowArray arr = ArrowArray.allocateNew(ALLOCATOR);
68+
ArrowSchema schema = ArrowSchema.allocateNew(ALLOCATOR)) {
69+
Data.exportVectorSchemaRoot(ALLOCATOR, root, null, arr, schema);
70+
writer.writeBatch(arr.memoryAddress(), schema.memoryAddress());
71+
}
72+
}
73+
}
74+
75+
private static List<ScanResult> scanAll(VortexReader vf) {
76+
return scanAll(vf, io.github.dfa1.vortex.scan.ScanOptions.all());
77+
}
78+
79+
private static List<ScanResult> scanAll(VortexReader vf,
80+
io.github.dfa1.vortex.scan.ScanOptions opts) {
81+
var results = new ArrayList<ScanResult>();
82+
var iter = vf.scan(opts);
83+
while (iter.hasNext()) {
84+
results.add(iter.next());
85+
}
86+
return results;
87+
}
88+
89+
// ── JNI write helpers ─────────────────────────────────────────────────────
90+
91+
private static long[] toLongs(ScanResult chunk) {
92+
var layout = ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
93+
var arr = chunk.columns().get("id");
94+
long[] out = new long[(int) arr.length()];
95+
for (int i = 0; i < out.length; i++) {
96+
out[i] = arr.buffer(0).get(layout, (long) i * Long.BYTES);
97+
}
98+
return out;
99+
}
100+
101+
private static double[] toDoubles(ScanResult chunk) {
102+
var layout = ValueLayout.JAVA_DOUBLE_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
103+
var arr = chunk.columns().get("value");
104+
double[] out = new double[(int) arr.length()];
105+
for (int i = 0; i < out.length; i++) {
106+
out[i] = arr.buffer(0).get(layout, (long) i * Double.BYTES);
107+
}
108+
return out;
109+
}
110+
111+
// ── Java read helpers ─────────────────────────────────────────────────────
112+
113+
@Test
114+
void jniWriter_javaReader_singleChunk(@TempDir Path tmp) throws IOException {
115+
// Given
116+
Path file = tmp.resolve("jni_single.vtx");
117+
long[] ids = {1L, 2L, 3L};
118+
double[] vals = {1.1, 2.2, 3.3};
119+
writeJni(file, ids, vals);
120+
121+
// When / Then
122+
try (var vf = VortexReader.open(file, CodecRegistry.loadAll())) {
123+
List<ScanResult> results = scanAll(vf);
124+
assertThat(results).hasSize(1);
125+
assertThat(results.get(0).rowCount()).isEqualTo(3L);
126+
assertThat(toLongs(results.get(0))).containsExactly(1L, 2L, 3L);
127+
assertThat(toDoubles(results.get(0))).containsExactly(1.1, 2.2, 3.3);
128+
}
129+
}
130+
131+
@Test
132+
void jniWriter_javaReader_multipleChunks(@TempDir Path tmp) throws IOException {
133+
// Given
134+
Path file = tmp.resolve("jni_multi.vtx");
135+
String uri = file.toAbsolutePath().toUri().toString();
136+
try (VortexWriter writer = VortexWriter.create(SESSION, uri, JNI_SCHEMA, new HashMap<>(), ALLOCATOR)) {
137+
flushBatch(writer, new long[]{1L, 2L}, new double[]{1.1, 2.2});
138+
flushBatch(writer, new long[]{3L, 4L, 5L}, new double[]{3.3, 4.4, 5.5});
139+
}
140+
141+
// When / Then — JNI may merge small batches; verify total rows and values
142+
try (var vf = VortexReader.open(file, CodecRegistry.loadAll())) {
143+
List<ScanResult> results = scanAll(vf);
144+
long totalRows = results.stream().mapToLong(ScanResult::rowCount).sum();
145+
assertThat(totalRows).isEqualTo(5L);
146+
long[] allIds = results.stream()
147+
.flatMapToLong(r -> java.util.Arrays.stream(toLongs(r)))
148+
.toArray();
149+
assertThat(allIds).containsExactly(1L, 2L, 3L, 4L, 5L);
150+
}
151+
}
152+
153+
@Test
154+
void jniWriter_javaReader_columnProjection(@TempDir Path tmp) throws IOException {
155+
// Given
156+
Path file = tmp.resolve("jni_proj.vtx");
157+
writeJni(file, new long[]{10L, 20L}, new double[]{0.1, 0.2});
158+
159+
// When / Then
160+
try (var vf = VortexReader.open(file, CodecRegistry.loadAll())) {
161+
List<ScanResult> results = scanAll(vf, io.github.dfa1.vortex.scan.ScanOptions.columns("id"));
162+
assertThat(results).hasSize(1);
163+
assertThat(results.get(0).columns()).containsKey("id");
164+
assertThat(results.get(0).columns()).doesNotContainKey("value");
165+
assertThat(toLongs(results.get(0))).containsExactly(10L, 20L);
166+
}
167+
}
168+
169+
@Test
170+
void jniWriter_javaReader_fewUniqueF64Values(@TempDir Path tmp) throws IOException {
171+
// Given — 10_000 rows cycling through only 3 unique F64 values to trigger dict encoding
172+
int n = 10_000;
173+
long[] ids = new long[n];
174+
double[] vals = new double[n];
175+
double[] unique = {1.1, 2.2, 3.3};
176+
for (int i = 0; i < n; i++) {
177+
ids[i] = i;
178+
vals[i] = unique[i % unique.length];
179+
}
180+
Path file = tmp.resolve("jni_dict.vtx");
181+
writeJni(file, ids, vals);
182+
183+
// When / Then
184+
try (var vf = VortexReader.open(file, CodecRegistry.loadAll())) {
185+
List<ScanResult> results = scanAll(vf, io.github.dfa1.vortex.scan.ScanOptions.columns("value"));
186+
long total = results.stream().mapToLong(ScanResult::rowCount).sum();
187+
assertThat(total).isEqualTo(n);
188+
var layout = ValueLayout.JAVA_DOUBLE_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
189+
double sum = 0;
190+
for (ScanResult r : results) {
191+
var arr = r.columns().get("value");
192+
for (long j = 0; j < arr.length(); j++) {
193+
sum += arr.buffer(0).get(layout, j * Double.BYTES);
194+
}
195+
}
196+
// 10_000 rows: 3333 full cycles of [1.1,2.2,3.3] (=6.6 each) + one 1.1 remainder
197+
assertThat(sum).isCloseTo(21_998.9, org.assertj.core.data.Offset.offset(0.1));
198+
}
199+
}
201200
}

0 commit comments

Comments
 (0)