Skip to content

Commit 649a076

Browse files
dfa1claude
dfa1
and
claude
committed
feat: compress Parquet timestamps via datetimeparts encoding
Detect TIMESTAMP INT64 columns in ParquetImporter and map them to DType.Extension("vortex.timestamp") + DateTimePartsData so the cascading compressor decomposes microseconds into days/seconds/subseconds and bitpacks each component independently. Wire up the full path: - ParquetImporter: TimestampType → timestampExtension(), buildChunk wraps long[] in DateTimePartsData - VortexWriter: handle DateTimePartsData in arrayLength, DType.Extension in serializeDType, add DateTimePartsEncoding to CASCADE_CODECS - DateTimePartsEncoding: add encodeCascade exposing 3 I64 child slots - CascadingCompressor: use spliceResult for non-primitive dtypes so cascaded children are recursively compressed NYC taxi 2024-01 result: 76.0 MB → 53.0 MB (Rust reference: 42.8 MB) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent f402c04 commit 649a076

4 files changed

Lines changed: 105 additions & 7 deletions

File tree

‎core/src/main/java/io/github/dfa1/vortex/encoding/CascadingCompressor.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,9 +33,11 @@ private EncodeResult encodeWithCtx(DType dtype, Object data, CompressorContext c
3333
if (dtype instanceof DType.Struct structDtype) {
3434
return encodeStruct(structDtype, (StructData) data, ctx);
3535
}
36-
// Cascading only targets primitives (integers + floats); other types use first-match.
36+
// Non-primitives (extension types): find the accepting encoding and splice
37+
// through it so its cascaded children (e.g. datetimeparts → days/seconds/subseconds)
38+
// are recursively compressed rather than stored as raw primitives.
3739
if (!(dtype instanceof DType.Primitive)) {
38-
return findPrimitiveEncoding(dtype).encode(dtype, data);
40+
return spliceResult(findPrimitiveEncoding(dtype), dtype, data, ctx);
3941
}
4042
int n = dataLength(data);
4143

‎core/src/main/java/io/github/dfa1/vortex/encoding/DateTimePartsEncoding.java‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,11 @@ public EncodeResult encode(DType dtype, Object data) {
4848
return Encoder.encode((DType.Extension) dtype, (DateTimePartsData) data);
4949
}
5050

51+
@Override
52+
public CascadeStep encodeCascade(DType dtype, Object data, CompressorContext ctx) {
53+
return Encoder.encodeCascade((DType.Extension) dtype, (DateTimePartsData) data);
54+
}
55+
5156
@Override
5257
public Array decode(DecodeContext ctx) {
5358
return Decoder.decode(ctx);
@@ -125,6 +130,53 @@ static EncodeResult encode(DType.Extension dtype, DateTimePartsData data) {
125130
new int[]{});
126131
return new EncodeResult(root, List.copyOf(allBuffers), null, null);
127132
}
133+
134+
static CascadeStep encodeCascade(DType.Extension dtype, DateTimePartsData data) {
135+
ByteBuffer extMeta = dtype.metadata();
136+
byte[] extBytes = new byte[extMeta.remaining()];
137+
extMeta.duplicate().get(extBytes);
138+
TimeUnit unit = TimeUnit.fromTag(extBytes[0]);
139+
140+
long divisor = unit.divisor();
141+
long ticksPerDay = SECONDS_PER_DAY * divisor;
142+
int n = data.timestamps().length;
143+
144+
long[] days = new long[n];
145+
long[] seconds = new long[n];
146+
long[] subseconds = new long[n];
147+
148+
for (int i = 0; i < n; i++) {
149+
long ts = data.timestamps()[i];
150+
long d = ts / ticksPerDay;
151+
long rem = ts % ticksPerDay;
152+
if (rem < 0) {
153+
rem += ticksPerDay;
154+
d--;
155+
}
156+
days[i] = d;
157+
seconds[i] = rem / divisor;
158+
subseconds[i] = rem % divisor;
159+
}
160+
161+
byte[] metaBytes = EncodingProtos.DateTimePartsMetadata.newBuilder()
162+
.setDaysPtype(I64_PROTO).setSecondsPtype(I64_PROTO).setSubsecondsPtype(I64_PROTO)
163+
.build().toByteArray();
164+
165+
// 3 null slots filled by the cascading compressor (days, seconds, subseconds)
166+
EncodeNode partialRoot = new EncodeNode(
167+
EncodingId.VORTEX_DATETIMEPARTS,
168+
ByteBuffer.wrap(metaBytes),
169+
new EncodeNode[3],
170+
new int[0]);
171+
172+
DType daysDtype = data.nullable() ? I64_NULLABLE : I64;
173+
List<ChildSlot> children = List.of(
174+
new ChildSlot(daysDtype, days, 0),
175+
new ChildSlot(I64, seconds, 1),
176+
new ChildSlot(I64, subseconds, 2));
177+
178+
return new CascadeStep(partialRoot, List.of(), children, null, null);
179+
}
128180
}
129181

130182
private static final class Decoder {

‎parquet/src/main/java/io/github/dfa1/vortex/parquet/ParquetImporter.java‎

Lines changed: 32 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,13 @@
1010
import dev.hardwood.schema.FileSchema;
1111
import io.github.dfa1.vortex.core.DType;
1212
import io.github.dfa1.vortex.core.PType;
13+
import io.github.dfa1.vortex.encoding.DateTimePartsData;
14+
import io.github.dfa1.vortex.encoding.TimeUnit;
1315
import io.github.dfa1.vortex.writer.VortexWriter;
1416

1517
import java.io.IOException;
18+
import java.nio.ByteBuffer;
19+
import java.nio.ByteOrder;
1620
import java.nio.channels.FileChannel;
1721
import java.nio.file.Path;
1822
import java.nio.file.StandardOpenOption;
@@ -94,7 +98,7 @@ public static void importParquet(Path parquetPath, Path vortexPath, ImportOption
9498
rowsDone++;
9599

96100
if (chunkPos == chunkSize) {
97-
writer.writeChunk(buildChunk(columns, buffers, chunkPos));
101+
writer.writeChunk(buildChunk(columns, types, buffers, chunkPos));
98102
buffers = allocateBuffers(columns, chunkSize);
99103
chunkPos = 0;
100104
if (options.progressListener() != null) {
@@ -104,7 +108,7 @@ public static void importParquet(Path parquetPath, Path vortexPath, ImportOption
104108
}
105109

106110
if (chunkPos > 0) {
107-
writer.writeChunk(buildChunk(columns, buffers, chunkPos));
111+
writer.writeChunk(buildChunk(columns, types, buffers, chunkPos));
108112
if (options.progressListener() != null) {
109113
options.progressListener().onProgress(rowsDone, totalRows);
110114
}
@@ -143,9 +147,26 @@ private static DType mapInt64(LogicalType logical, boolean nullable) {
143147
if (logical instanceof LogicalType.IntType it) {
144148
return new DType.Primitive(it.isSigned() ? PType.I64 : PType.U64, nullable);
145149
}
150+
if (logical instanceof LogicalType.TimestampType ts) {
151+
return timestampExtension(ts.unit(), nullable);
152+
}
146153
return new DType.Primitive(PType.I64, nullable);
147154
}
148155

156+
private static DType.Extension timestampExtension(LogicalType.TimeUnit parquetUnit, boolean nullable) {
157+
TimeUnit unit = switch (parquetUnit) {
158+
case MILLIS -> TimeUnit.Milliseconds;
159+
case MICROS -> TimeUnit.Microseconds;
160+
case NANOS -> TimeUnit.Nanoseconds;
161+
};
162+
// Extension metadata: byte[0]=unit tag, bytes[1-2]=tz_len u16 LE (0 = no tz)
163+
ByteBuffer meta = ByteBuffer.allocate(3).order(ByteOrder.LITTLE_ENDIAN);
164+
meta.put((byte) unit.ordinal());
165+
meta.putShort((short) 0);
166+
meta.flip();
167+
return new DType.Extension("vortex.timestamp", new DType.Primitive(PType.I64, nullable), meta, nullable);
168+
}
169+
149170
private static DType mapByteArray(LogicalType logical, boolean nullable, String colName) {
150171
if (logical instanceof LogicalType.StringType
151172
|| logical instanceof LogicalType.EnumType
@@ -181,7 +202,7 @@ yield switch (it.bitWidth()) {
181202
}
182203
yield new int[chunkSize];
183204
}
184-
case INT64 -> new long[chunkSize];
205+
case INT64 -> new long[chunkSize]; // covers plain I64 and timestamps
185206
default -> throw new UnsupportedOperationException("unsupported type: " + col.type());
186207
};
187208
}
@@ -206,10 +227,16 @@ private static void fillRow(RowReader reader, List<ColumnSchema> columns, Object
206227
}
207228
}
208229

209-
private static Map<String, Object> buildChunk(List<ColumnSchema> columns, Object[] buffers, int size) {
230+
private static Map<String, Object> buildChunk(List<ColumnSchema> columns, List<DType> types,
231+
Object[] buffers, int size) {
210232
Map<String, Object> chunk = new LinkedHashMap<>();
211233
for (int c = 0; c < columns.size(); c++) {
212-
chunk.put(columns.get(c).name(), trimBuffer(buffers[c], size));
234+
Object buf = trimBuffer(buffers[c], size);
235+
if (types.get(c) instanceof DType.Extension) {
236+
boolean nullable = types.get(c).nullable();
237+
buf = new DateTimePartsData((long[]) buf, nullable);
238+
}
239+
chunk.put(columns.get(c).name(), buf);
213240
}
214241
return chunk;
215242
}

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
import io.github.dfa1.vortex.encoding.BoolEncoding;
88
import io.github.dfa1.vortex.encoding.CascadingCompressor;
99
import io.github.dfa1.vortex.encoding.CompressorContext;
10+
import io.github.dfa1.vortex.encoding.DateTimePartsData;
11+
import io.github.dfa1.vortex.encoding.DateTimePartsEncoding;
1012
import io.github.dfa1.vortex.encoding.DictEncoding;
1113
import io.github.dfa1.vortex.encoding.Encoding;
1214
import io.github.dfa1.vortex.encoding.FrameOfReferenceEncoding;
@@ -18,6 +20,7 @@
1820
import io.github.dfa1.vortex.encoding.EncodingId;
1921
import io.github.dfa1.vortex.encoding.PrimitiveEncoding;
2022
import io.github.dfa1.vortex.fbs.ArraySpec;
23+
import io.github.dfa1.vortex.fbs.Extension;
2124
import io.github.dfa1.vortex.fbs.Footer;
2225
import io.github.dfa1.vortex.fbs.Layout;
2326
import io.github.dfa1.vortex.fbs.LayoutSpec;
@@ -65,6 +68,7 @@ public final class VortexWriter implements Closeable {
6568
new AlpEncoding(), new PrimitiveEncoding(), new BoolEncoding(), new DictEncoding(), new VarBinEncoding());
6669

6770
private static final List<Encoding> CASCADE_CODECS = List.of(
71+
new DateTimePartsEncoding(),
6872
new AlpEncoding(), new FrameOfReferenceEncoding(), new DictEncoding(),
6973
new BitpackedEncoding(), new VarBinEncoding(), new PrimitiveEncoding(), new BoolEncoding());
7074

@@ -113,6 +117,7 @@ private static long arrayLength(Object data) {
113117
case String[] a -> a.length;
114118
case ListData d -> d.outerLen();
115119
case ListViewData d -> d.outerLen();
120+
case DateTimePartsData d -> d.timestamps().length;
116121
default -> throw new UnsupportedOperationException(
117122
"unsupported data type: " + data.getClass());
118123
};
@@ -166,6 +171,18 @@ private static int serializeDType(FlatBufferBuilder fbb, DType dtype) {
166171
int inner = io.github.dfa1.vortex.fbs.List.createList(fbb, elemTypeOff, l.nullable());
167172
yield io.github.dfa1.vortex.fbs.DType.createDType(fbb, Type.List, inner);
168173
}
174+
case DType.Extension e -> {
175+
int idOff = fbb.createString(e.extensionId());
176+
int storageDtypeOff = serializeDType(fbb, e.storageDType());
177+
int metaOff = 0;
178+
if (e.metadata() != null && e.metadata().remaining() > 0) {
179+
byte[] metaBytes = new byte[e.metadata().remaining()];
180+
e.metadata().duplicate().get(metaBytes);
181+
metaOff = Extension.createMetadataVector(fbb, metaBytes);
182+
}
183+
int inner = Extension.createExtension(fbb, idOff, storageDtypeOff, metaOff);
184+
yield io.github.dfa1.vortex.fbs.DType.createDType(fbb, Type.Extension, inner);
185+
}
169186
default -> throw new UnsupportedOperationException("unsupported DType: " + dtype);
170187
};
171188
}

0 commit comments

Comments
 (0)