Skip to content

Commit 1800dbf

Browse files
dfa1claude
andcommitted
feat(writer): add VarBin encode + Utf8 dtype serialization; add symbol column to CSV benchmark
- VarBinEncoding.encode() handles String[] → bytes buf + I64 offsets - VortexWriter.DEFAULT_CODECS includes VarBinEncoding - VortexWriter.serializeDType supports DType.Utf8 - VortexWriter.arrayLength handles String[] - VortexVsCsvFileSizeTest: 30 real ticker symbols (AAPL, MSFT...) as a Utf8 column Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent becdde9 commit 1800dbf

3 files changed

Lines changed: 66 additions & 8 deletions

File tree

core/src/main/java/io/github/dfa1/vortex/encoding/VarBinEncoding.java

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,13 @@
99
import io.github.dfa1.vortex.core.array.VarBinArray;
1010
import io.github.dfa1.vortex.core.VortexException;
1111

12+
import java.lang.foreign.Arena;
1213
import java.lang.foreign.MemorySegment;
14+
import java.lang.foreign.ValueLayout;
1315
import java.nio.ByteBuffer;
16+
import java.nio.ByteOrder;
17+
import java.nio.charset.StandardCharsets;
18+
import java.util.List;
1419

1520
/// Decoder for {@code vortex.varbin} — variable-length binary / UTF-8 string arrays.
1621
///
@@ -33,9 +38,45 @@ public EncodingId encodingId() {
3338
return EncodingId.VORTEX_VARBIN;
3439
}
3540

41+
@Override
42+
public boolean accepts(DType dtype) {
43+
return dtype instanceof DType.Utf8 || dtype instanceof DType.Binary;
44+
}
45+
3646
@Override
3747
public EncodeResult encode(DType dtype, Object data) {
38-
throw new UnsupportedOperationException("encode not supported by " + encodingId());
48+
String[] strings = (String[]) data;
49+
int n = strings.length;
50+
51+
byte[][] byteArrays = new byte[n][];
52+
int totalBytes = 0;
53+
for (int i = 0; i < n; i++) {
54+
byteArrays[i] = strings[i].getBytes(StandardCharsets.UTF_8);
55+
totalBytes += byteArrays[i].length;
56+
}
57+
58+
Arena arena = Arena.ofAuto();
59+
MemorySegment bytesBuf = arena.allocate(totalBytes > 0 ? totalBytes : 1);
60+
MemorySegment offsetsBuf = arena.allocate((long) (n + 1) * Long.BYTES, Long.BYTES);
61+
ValueLayout.OfLong leI64 = ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
62+
63+
long pos = 0;
64+
offsetsBuf.setAtIndex(leI64, 0, 0L);
65+
for (int i = 0; i < n; i++) {
66+
MemorySegment.copy(MemorySegment.ofArray(byteArrays[i]), 0, bytesBuf, pos, byteArrays[i].length);
67+
pos += byteArrays[i].length;
68+
offsetsBuf.setAtIndex(leI64, i + 1, pos);
69+
}
70+
71+
byte[] metaBytes = EncodingProtos.VarBinMetadata.newBuilder()
72+
.setOffsetsPtype(dev.vortex.proto.DTypeProtos.PType.forNumber(PType.I64.ordinal()))
73+
.build()
74+
.toByteArray();
75+
76+
EncodeNode offsetsNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 1);
77+
EncodeNode root = new EncodeNode(encodingId(), ByteBuffer.wrap(metaBytes),
78+
new EncodeNode[]{offsetsNode}, new int[]{0});
79+
return new EncodeResult(root, List.of(bytesBuf, offsetsBuf), null, null);
3980
}
4081

4182
@Override

integration/src/test/java/io/github/dfa1/vortex/integration/VortexVsCsvFileSizeTest.java

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,16 @@ class VortexVsCsvFileSizeTest {
2929
private static final int TOTAL_ROWS = 100_000;
3030
private static final int BATCH_SIZE = 50_000;
3131

32+
private static final String[] SYMBOLS = {
33+
"AAPL", "MSFT", "GOOGL", "AMZN", "NVDA", "META", "TSLA", "BRK.B", "JPM", "V",
34+
"UNH", "XOM", "LLY", "JNJ", "MA", "PG", "HD", "MRK", "AVGO", "CVX",
35+
"PEP", "ABBV", "KO", "COST", "WMT", "MCD", "CSCO", "ACN", "BAC", "CRM"
36+
};
37+
3238
private static final DType.Struct SCHEMA = new DType.Struct(
33-
List.of("date", "open", "high", "low", "close", "volume"),
39+
List.of("symbol", "date", "open", "high", "low", "close", "volume"),
3440
List.of(
41+
new DType.Utf8(false),
3542
new DType.Primitive(PType.I32, false),
3643
new DType.Primitive(PType.F64, false),
3744
new DType.Primitive(PType.F64, false),
@@ -42,7 +49,7 @@ class VortexVsCsvFileSizeTest {
4249
false
4350
);
4451

45-
private record OhlcBatch(int[] dates, double[] open, double[] high,
52+
private record OhlcBatch(String[] symbols, int[] dates, double[] open, double[] high,
4653
double[] low, double[] close, long[] volume) {}
4754

4855
private static List<OhlcBatch> generateBatches() {
@@ -55,6 +62,7 @@ private static List<OhlcBatch> generateBatches() {
5562

5663
while (rowsLeft > 0) {
5764
int n = Math.min(rowsLeft, BATCH_SIZE);
65+
String[] symbols = new String[n];
5866
int[] dates = new int[n];
5967
double[] open = new double[n], high = new double[n], low = new double[n], close = new double[n];
6068
long[] volume = new long[n];
@@ -65,6 +73,7 @@ private static List<OhlcBatch> generateBatches() {
6573
double o = Math.round(px * (1 + ret * 0.3) * 100.0) / 100.0;
6674
double c = Math.round(px * (1 + ret) * 100.0) / 100.0;
6775
double spread = Math.abs(px * rng.nextDouble() * 0.03);
76+
symbols[i] = SYMBOLS[ticker];
6877
dates[i] = startDay + (rowsDone + i) / 30;
6978
open[i] = o;
7079
high[i] = Math.round((Math.max(o, c) + spread) * 100.0) / 100.0;
@@ -73,7 +82,7 @@ private static List<OhlcBatch> generateBatches() {
7382
volume[i] = Math.max(100_000L, Math.round(1_000_000 + rng.nextGaussian() * 200_000));
7483
prices[ticker] = c;
7584
}
76-
batches.add(new OhlcBatch(dates, open, high, low, close, volume));
85+
batches.add(new OhlcBatch(symbols, dates, open, high, low, close, volume));
7786
rowsLeft -= n;
7887
rowsDone += n;
7988
}
@@ -92,18 +101,19 @@ void vortexSmallerThanCsv(@TempDir Path tmp) throws IOException {
92101
VortexWriter writer = VortexWriter.create(ch, SCHEMA, WriteOptions.defaults())) {
93102
for (OhlcBatch b : batches) {
94103
writer.writeChunk(Map.of(
95-
"date", b.dates(), "open", b.open(), "high", b.high(),
104+
"symbol", b.symbols(), "date", b.dates(), "open", b.open(), "high", b.high(),
96105
"low", b.low(), "close", b.close(), "volume", b.volume()
97106
));
98107
}
99108
}
100109

101110
// When — write CSV
102111
try (BufferedWriter csv = Files.newBufferedWriter(csvFile)) {
103-
csv.write("date,open,high,low,close,volume\n");
112+
csv.write("symbol,date,open,high,low,close,volume\n");
104113
for (OhlcBatch b : batches) {
105114
for (int i = 0; i < b.dates().length; i++) {
106-
csv.write(LocalDate.ofEpochDay(b.dates()[i]) + "," + b.open()[i] + "," + b.high()[i] + ","
115+
csv.write(b.symbols()[i] + "," + LocalDate.ofEpochDay(b.dates()[i]) + ","
116+
+ b.open()[i] + "," + b.high()[i] + ","
107117
+ b.low()[i] + "," + b.close()[i] + "," + b.volume()[i] + "\n");
108118
}
109119
}

writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import io.github.dfa1.vortex.encoding.AlpEncoding;
66
import io.github.dfa1.vortex.encoding.BoolEncoding;
77
import io.github.dfa1.vortex.encoding.Encoding;
8+
import io.github.dfa1.vortex.encoding.VarBinEncoding;
89
import io.github.dfa1.vortex.encoding.EncodeNode;
910
import io.github.dfa1.vortex.encoding.EncodeResult;
1011
import io.github.dfa1.vortex.encoding.EncodingId;
@@ -53,7 +54,8 @@ public final class VortexWriter implements Closeable {
5354
private static final ByteBuffer MAGIC = ByteBuffer.wrap(new byte[]{'V', 'T', 'X', 'F'})
5455
.asReadOnlyBuffer();
5556

56-
private static final List<Encoding> DEFAULT_CODECS = List.of(new AlpEncoding(), new PrimitiveEncoding(), new BoolEncoding());
57+
private static final List<Encoding> DEFAULT_CODECS = List.of(
58+
new AlpEncoding(), new VarBinEncoding(), new PrimitiveEncoding(), new BoolEncoding());
5759

5860
private final WritableByteChannel channel;
5961
private final DType.Struct schema;
@@ -97,6 +99,7 @@ private static long arrayLength(Object data) {
9799
case float[] a -> a.length;
98100
case double[] a -> a.length;
99101
case boolean[] a -> a.length;
102+
case String[] a -> a.length;
100103
default -> throw new UnsupportedOperationException(
101104
"unsupported data type: " + data.getClass());
102105
};
@@ -141,6 +144,10 @@ private static int serializeDType(FlatBufferBuilder fbb, DType dtype) {
141144
fbb, namesVec, dtypesVec, s.nullable());
142145
yield io.github.dfa1.vortex.fbs.DType.createDType(fbb, Type.Struct_, inner);
143146
}
147+
case DType.Utf8 u -> {
148+
int inner = io.github.dfa1.vortex.fbs.Utf8.createUtf8(fbb, u.nullable());
149+
yield io.github.dfa1.vortex.fbs.DType.createDType(fbb, Type.Utf8, inner);
150+
}
144151
default -> throw new UnsupportedOperationException("unsupported DType: " + dtype);
145152
};
146153
}

0 commit comments

Comments
 (0)