Skip to content

Commit fee9e42

Browse files
dfa1claude
andcommitted
feat: add LongArray#fold and ScanResult#column
fold(identity, LongBinaryOperator) enables zero-boxing aggregation (sum, min, max) without the long[]{0} trick. column(name) returns a typed array via inferred generic, throwing VortexException on unknown column names. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent e16181b commit fee9e42

4 files changed

Lines changed: 139 additions & 0 deletions

File tree

core/src/main/java/io/github/dfa1/vortex/core/array/LongArray.java

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import java.lang.foreign.MemorySegment;
88
import java.lang.foreign.ValueLayout;
99
import java.nio.ByteOrder;
10+
import java.util.function.LongBinaryOperator;
1011
import java.util.function.LongConsumer;
1112

1213
/// Concrete [Array] for I64/U64 primitive columns. Hoists a single
@@ -70,4 +71,14 @@ public void forEachLong(LongConsumer c) {
7071
c.accept(buf.getAtIndex(LAYOUT, i));
7172
}
7273
}
74+
75+
public long fold(long identity, LongBinaryOperator op) {
76+
MemorySegment buf = buffer;
77+
long n = length;
78+
long result = identity;
79+
for (long i = 0; i < n; i++) {
80+
result = op.applyAsLong(result, buf.getAtIndex(LAYOUT, i));
81+
}
82+
return result;
83+
}
7384
}
Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
package io.github.dfa1.vortex.core.array;
2+
3+
import io.github.dfa1.vortex.core.ArrayStats;
4+
import io.github.dfa1.vortex.core.DType;
5+
import io.github.dfa1.vortex.core.PType;
6+
import org.junit.jupiter.api.Test;
7+
8+
import java.lang.foreign.Arena;
9+
import java.lang.foreign.MemorySegment;
10+
import java.lang.foreign.ValueLayout;
11+
import java.nio.ByteOrder;
12+
13+
import static org.assertj.core.api.Assertions.assertThat;
14+
15+
class LongArrayTest {
16+
17+
private static final ValueLayout.OfLong LE_LONG =
18+
ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
19+
20+
private static LongArray of(long... values) {
21+
MemorySegment seg = Arena.ofAuto().allocate((long) values.length * 8, 8);
22+
for (int i = 0; i < values.length; i++) {
23+
seg.setAtIndex(LE_LONG, i, values[i]);
24+
}
25+
DType dtype = new DType.Primitive(PType.I64, false);
26+
return new LongArray(dtype, values.length, seg, ArrayStats.empty());
27+
}
28+
29+
@Test
30+
void fold_sum_returnsCorrectTotal() {
31+
// Given
32+
LongArray sut = of(1L, 2L, 3L, 4L, 5L);
33+
34+
// When
35+
long result = sut.fold(0L, Long::sum);
36+
37+
// Then
38+
assertThat(result).isEqualTo(15L);
39+
}
40+
41+
@Test
42+
void fold_max_returnsLargestValue() {
43+
// Given
44+
LongArray sut = of(3L, 1L, 4L, 1L, 5L, 9L, 2L);
45+
46+
// When
47+
long result = sut.fold(Long.MIN_VALUE, Math::max);
48+
49+
// Then
50+
assertThat(result).isEqualTo(9L);
51+
}
52+
53+
@Test
54+
void fold_min_returnsSmallestValue() {
55+
// Given
56+
LongArray sut = of(3L, 1L, 4L, 1L, 5L, 9L, 2L);
57+
58+
// When
59+
long result = sut.fold(Long.MAX_VALUE, Math::min);
60+
61+
// Then
62+
assertThat(result).isEqualTo(1L);
63+
}
64+
65+
@Test
66+
void fold_emptyArray_returnsIdentity() {
67+
// Given
68+
LongArray sut = of();
69+
70+
// When
71+
long result = sut.fold(42L, Long::sum);
72+
73+
// Then
74+
assertThat(result).isEqualTo(42L);
75+
}
76+
}

reader/src/main/java/io/github/dfa1/vortex/scan/ScanResult.java

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package io.github.dfa1.vortex.scan;
22

33
import io.github.dfa1.vortex.core.Array;
4+
import io.github.dfa1.vortex.core.VortexException;
45

56
import java.util.Map;
67

@@ -9,4 +10,12 @@ public record ScanResult(
910
long rowCount,
1011
Map<String, Array> columns
1112
) {
13+
@SuppressWarnings("unchecked")
14+
public <T extends Array> T column(String name) {
15+
Array arr = columns.get(name);
16+
if (arr == null) {
17+
throw new VortexException("unknown column: " + name);
18+
}
19+
return (T) arr;
20+
}
1221
}

writer/src/test/java/io/github/dfa1/vortex/writer/VortexWriterTest.java

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import io.github.dfa1.vortex.core.Array;
44
import io.github.dfa1.vortex.core.DType;
55
import io.github.dfa1.vortex.core.PType;
6+
import io.github.dfa1.vortex.core.array.LongArray;
67
import io.github.dfa1.vortex.encoding.CodecRegistry;
78
import io.github.dfa1.vortex.encoding.PrimitiveCodec;
89
import io.github.dfa1.vortex.io.VortexReader;
@@ -138,6 +139,48 @@ void writeAndRead_idValues_decodedCorrectly(@TempDir Path tmp) throws IOExceptio
138139
}
139140
}
140141

142+
@Test
143+
void scanResult_column_returnsTypedArray(@TempDir Path tmp) throws IOException {
144+
// Given
145+
Path file = tmp.resolve("typed.vtx");
146+
long[] ids = {10L, 20L, 30L};
147+
148+
try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE);
149+
var sut = VortexWriter.create(ch, SCHEMA, WriteOptions.defaults())) {
150+
sut.writeChunk(Map.of("id", ids, "value", new double[]{1.0, 2.0, 3.0}));
151+
}
152+
153+
// When
154+
var registry = primitiveRegistry();
155+
try (var vf = VortexReader.open(file, registry)) {
156+
List<ScanResult> results = scanAll(vf, ScanOptions.all());
157+
LongArray idArray = results.get(0).column("id");
158+
159+
// Then
160+
assertThat(idArray.fold(0L, Long::sum)).isEqualTo(60L);
161+
}
162+
}
163+
164+
@Test
165+
void scanResult_column_unknownName_throwsVortexException(@TempDir Path tmp) throws IOException {
166+
// Given
167+
Path file = tmp.resolve("unknown.vtx");
168+
169+
try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE);
170+
var sut = VortexWriter.create(ch, SCHEMA, WriteOptions.defaults())) {
171+
sut.writeChunk(Map.of("id", new long[]{1L}, "value", new double[]{1.0}));
172+
}
173+
174+
// When / Then
175+
var registry = primitiveRegistry();
176+
try (var vf = VortexReader.open(file, registry)) {
177+
List<ScanResult> results = scanAll(vf, ScanOptions.all());
178+
ScanResult sut = results.get(0);
179+
assertThatThrownBy(() -> sut.column("nonexistent"))
180+
.hasMessageContaining("unknown column: nonexistent");
181+
}
182+
}
183+
141184
@Test
142185
void writeAndRead_columnProjection_returnsOnlyRequestedColumns(@TempDir Path tmp)
143186
throws IOException {

0 commit comments

Comments
 (0)