diff --git a/CHANGELOG.md b/CHANGELOG.md index 86bfeba45..ae9b68d05 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- `vortex.json` and the `vortex.st.*` geometry columns open without a decoder and read as their storage, as vortex-jni reads them. ([#547](https://github.com/dfa1/vortex-java/pull/547)) +- The CLI takes `--allow-unknown` to open files with extension types it has no decoder for, and CSV `ExportOptions.withReadRegistry` picks the registry an export opens the file with. ([#547](https://github.com/dfa1/vortex-java/pull/547)) +- Custom extension types: register an `ExtensionDecoder` on `ReadRegistry` and an `ExtensionEncoder` on `WriteRegistry`, with a `validate` hook on both sides, and `Chunk.as` decodes any registered type. ([#530](https://github.com/dfa1/vortex-java/issues/530)) + +### Changed +- **Breaking:** opening a file with an extension id that has no registered decoder fails unless `ReadRegistry.Builder#allowUnknown()` is set, as in Rust; the column then reads as its storage array. ([#530](https://github.com/dfa1/vortex-java/issues/530)) +- **Breaking:** `vortex.date`, `vortex.time`, `vortex.timestamp` and `vortex.uuid` dtypes are checked as Rust checks them (time unit, storage type, metadata) when writing and reading, and a file whose extension dtype has no metadata field is refused, as Rust refuses it. ([#547](https://github.com/dfa1/vortex-java/pull/547)) +- **Breaking:** `DateExtensionDecoder.decode` and `decodeAll` take the column's `DType.Extension`, which carries the time unit. ([#547](https://github.com/dfa1/vortex-java/pull/547)) +- **Breaking:** `ExtensionDecoder` is generic in its domain type and declares `domainType()` and `decodeAll(DType.Extension, Array)`. ([#530](https://github.com/dfa1/vortex-java/issues/530)) +- **Breaking:** `ExtensionId` is a sealed `WellKnown`/`Custom` type like `EncodingId`, so third-party extension ids can be named, and `ExtensionId.parse` returns the id instead of an `Optional`. ([#530](https://github.com/dfa1/vortex-java/issues/530)) + +### Fixed +- `vortex count`, `schema`, `stats` and `export` report a malformed file as an error instead of crashing with a stack trace. ([#547](https://github.com/dfa1/vortex-java/pull/547)) +- `vortex.date` columns in milliseconds, as Rust writes an Arrow `Date64`, decode to the right dates instead of failing or reading the count as days. ([#547](https://github.com/dfa1/vortex-java/pull/547)) +- An extension column with no metadata no longer writes a file Rust refuses to read. ([#547](https://github.com/dfa1/vortex-java/pull/547)) +- `VortexWriter.create(…, WriteRegistry)` uses the registry's extension encoders instead of the built-in ones, and a `Collection` for an extension column with no encoder fails with a `VortexException`. ([#530](https://github.com/dfa1/vortex-java/issues/530)) + ## [0.17.2] — 2026-10-17 ### Changed diff --git a/CLAUDE.md b/CLAUDE.md index 9730f608b..09422ee36 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -192,11 +192,13 @@ Add an `EncodingId.WellKnown` constant `VORTEX_FOO("vortex.foo")` (re-exported o ### Adding an extension type -Add `ExtensionId` constant, then per side: -- **Decode:** singleton `FooExtensionDecoder implements ExtensionDecoder` in `reader.extension` + - a `case VORTEX_FOO` in `Chunk.as()` — not registry-managed. -- **Encode:** `FooExtensionEncoder implements ExtensionEncoder` in `writer` + a - `.register(FooExtensionEncoder.INSTANCE)` call in `WriteRegistry.Builder#registerDefaults()` +Add an `ExtensionId.WellKnown` constant `VORTEX_FOO("vortex.foo")` (re-exported on the interface), then per side: +- **Decode:** singleton `FooExtensionDecoder implements ExtensionDecoder` in `reader.extension` + + a `.register(FooExtensionDecoder.INSTANCE)` call in `ReadRegistry.Builder#registerDefaults()`. +- **Encode:** `FooExtensionEncoder implements ExtensionEncoder` in `writer.encode` + a + `.register(FooExtensionEncoder.INSTANCE)` call in `WriteRegistry.Builder#registerDefaultExtensions()`. +- Both may override `validate(DType.Extension)` (Rust's `validate_dtype`). Third-party types skip the + constant: `ExtensionId.parse("acme.foo")` gives a `Custom` id, registered the same way. ## Memory model @@ -267,8 +269,9 @@ When stuck on encode/decode behavior, consult **in this order**: - **DType is pluggable only via `Extension`.** `DType` is a sealed interface; downstream code must not add variants. Use `new DType.Extension("ip.address", new DType.Primitive(PType.I32, false), - null, false)` and register decoders/encoders on the registries. - Mirrors Rust (`vortex.date`, `vortex.uuid`, …). No SPI for DType variants planned. + null, false)` and register an `ExtensionDecoder`/`ExtensionEncoder` on `ReadRegistry`/`WriteRegistry`. + Mirrors Rust's `ExtVTable` (`vortex.date`, `vortex.uuid`, …): an unregistered id fails the read + unless `allowUnknown()`, then reads as storage. No SPI for DType variants planned. - **Encoding and layout decode are both pluggable via builder-only registries** — no `ServiceLoader`. Encodings register on `ReadRegistry`/`WriteRegistry` (`ReadRegistry.builder().registerDefaults().register(custom).build()`); layouts register on @@ -418,6 +421,11 @@ Generated proto records live in `io.github.dfa1.vortex.core.proto`; the runtime one for every encoding round-trip and file-format boundary. Include **filtered and row-range** vortex-jni reads, not just full scans: a full scan never prunes or slices, which hid a JVM abort (buffer alignment) and silently empty filtered results (zone stride) until 0.15.x. +- **Every fix and feature gets a cross-check test against vortex-jni**, not only Java unit tests: + vortex-jni writes and Java reads, and/or Java writes and vortex-jni reads, asserting both agree. + Write it first, so it fails without the change. Java-only tests share Java's assumptions; + cross-checks are what caught unknown extension ids that Rust refuses, absent extension metadata + that made files unreadable, and `vortex.date` milliseconds read as days (#547). - JUnit 5 + Mockito (BDDMockito) + AssertJ. Class under test named `sut`. Every test has `// Given` / `// When` / `// Then`. BDDMockito only: `given(mock.m()).willReturn(v)` / `then(...)` (static-import only `given`/`then`, never `willReturn`/`willThrow`). diff --git a/cli/src/main/java/io/github/dfa1/vortex/cli/CliHandles.java b/cli/src/main/java/io/github/dfa1/vortex/cli/CliHandles.java index aa8e2da92..650dd12e1 100644 --- a/cli/src/main/java/io/github/dfa1/vortex/cli/CliHandles.java +++ b/cli/src/main/java/io/github/dfa1/vortex/cli/CliHandles.java @@ -1,6 +1,7 @@ package io.github.dfa1.vortex.cli; import io.github.dfa1.vortex.cli.tui.IoWorker; +import io.github.dfa1.vortex.reader.ReadRegistry; import io.github.dfa1.vortex.reader.VortexHandle; import io.github.dfa1.vortex.reader.VortexHttpReader; import io.github.dfa1.vortex.reader.VortexReader; @@ -88,9 +89,10 @@ static String describe(Throwable t) { /// @return the opened handle, or `null` if the target is missing or malformed /// @throws IOException if opening the file or URL fails static VortexHandle openTarget(String target) throws IOException { + ReadRegistry registry = readRegistry(); if (target.startsWith("http://") || target.startsWith("https://")) { try { - return VortexHttpReader.open(new URI(target)); + return VortexHttpReader.open(new URI(target), registry); } catch (URISyntaxException _) { System.err.println("invalid URL: " + target); return null; @@ -101,6 +103,25 @@ static VortexHandle openTarget(String target) throws IOException { System.err.println("file not found: " + path); return null; } - return VortexReader.open(path); + return VortexReader.open(path, registry); + } + + // One CLI invocation per process; VortexCli.execute sets it before every dispatch. + private static ReadRegistry readRegistry = ReadRegistry.loadAll(); + + /// Sets whether files open with unknown extension ids allowed (the `--allow-unknown` flag). + /// + /// @param allowUnknown `true` to read unregistered extension types as their storage + static void allowUnknown(boolean allowUnknown) { + readRegistry = allowUnknown + ? ReadRegistry.builder().registerDefaults().allowUnknown().build() + : ReadRegistry.loadAll(); + } + + /// The registry every subcommand opens Vortex files with. + /// + /// @return the read registry for this invocation + static ReadRegistry readRegistry() { + return readRegistry; } } diff --git a/cli/src/main/java/io/github/dfa1/vortex/cli/CountCommand.java b/cli/src/main/java/io/github/dfa1/vortex/cli/CountCommand.java index e9b380e4c..fd9ab7c8c 100644 --- a/cli/src/main/java/io/github/dfa1/vortex/cli/CountCommand.java +++ b/cli/src/main/java/io/github/dfa1/vortex/cli/CountCommand.java @@ -22,10 +22,10 @@ static int run(String[] args) { System.err.println("file not found: " + path); return ExitStatus.FILE_NOT_FOUND; } - try (VortexReader reader = VortexReader.open(path)) { + try (VortexReader reader = VortexReader.open(path, CliHandles.readRegistry())) { System.out.println(reader.layout().rowCount()); return ExitStatus.OK; - } catch (IOException e) { + } catch (IOException | io.github.dfa1.vortex.core.error.VortexException e) { System.err.println("error: " + e.getMessage()); return ExitStatus.ERROR; } diff --git a/cli/src/main/java/io/github/dfa1/vortex/cli/ExportCommand.java b/cli/src/main/java/io/github/dfa1/vortex/cli/ExportCommand.java index 28fdfa245..485ce147b 100644 --- a/cli/src/main/java/io/github/dfa1/vortex/cli/ExportCommand.java +++ b/cli/src/main/java/io/github/dfa1/vortex/cli/ExportCommand.java @@ -5,6 +5,7 @@ import io.github.dfa1.vortex.csv.ExportOptions; import io.github.dfa1.vortex.parquet.ParquetExporter; import io.github.dfa1.vortex.reader.VortexHandle; +import io.github.dfa1.vortex.reader.VortexReader; import java.io.IOException; import java.io.OutputStreamWriter; @@ -43,7 +44,7 @@ static int run(String[] args) { return runParquet(inputPath, outputPath); } return runCsv(inputPath, outputPath, toStdout); - } catch (IOException e) { + } catch (IOException | io.github.dfa1.vortex.core.error.VortexException e) { ProgressBar.clear(); System.err.println("error: " + e.getMessage()); return ExitStatus.ERROR; @@ -69,7 +70,7 @@ private static int runRemote(String target, String[] args, boolean toStdout) { ProgressBar.clear(); System.out.printf("written: %s (%s)%n", outputPath, ByteSize.format(Files.size(outputPath))); return ExitStatus.OK; - } catch (IOException e) { + } catch (IOException | io.github.dfa1.vortex.core.error.VortexException e) { ProgressBar.clear(); System.err.println("error: " + e.getMessage()); return ExitStatus.ERROR; @@ -78,7 +79,8 @@ private static int runRemote(String target, String[] args, boolean toStdout) { private static int runCsv(Path inputPath, Path outputPath, boolean toStdout) throws IOException { ExportOptions options = ExportOptions.defaults() - .withProgressListener(ProgressBar::render); + .withProgressListener(ProgressBar::render) + .withReadRegistry(CliHandles.readRegistry()); if (toStdout) { Writer stdout = new OutputStreamWriter(System.out, StandardCharsets.UTF_8); CsvExporter.exportCsv(inputPath, stdout, options); @@ -96,7 +98,9 @@ private static int runParquet(Path inputPath, Path outputPath) throws IOExceptio io.github.dfa1.vortex.parquet.ExportOptions options = io.github.dfa1.vortex.parquet.ExportOptions.defaults() .withProgressListener(ProgressBar::render); - ParquetExporter.exportParquet(inputPath, outputPath, options); + try (VortexReader reader = VortexReader.open(inputPath, CliHandles.readRegistry())) { + ParquetExporter.exportParquet(reader, outputPath, options); + } ProgressBar.clear(); printResult(inputPath, outputPath); return ExitStatus.OK; diff --git a/cli/src/main/java/io/github/dfa1/vortex/cli/FilterCommand.java b/cli/src/main/java/io/github/dfa1/vortex/cli/FilterCommand.java index ca215ed8a..c8bd3bd07 100644 --- a/cli/src/main/java/io/github/dfa1/vortex/cli/FilterCommand.java +++ b/cli/src/main/java/io/github/dfa1/vortex/cli/FilterCommand.java @@ -58,7 +58,8 @@ static int run(String[] args) { RowPredicate rowPred = toRowPredicate(filter); try { Writer stdout = new OutputStreamWriter(System.out, StandardCharsets.UTF_8); - CsvExporter.exportCsvFiltered(path, stdout, ExportOptions.defaults(), scanOptions, rowPred); + CsvExporter.exportCsvFiltered(path, stdout, + ExportOptions.defaults().withReadRegistry(CliHandles.readRegistry()), scanOptions, rowPred); stdout.flush(); return ExitStatus.OK; } catch (IOException | io.github.dfa1.vortex.core.error.VortexException e) { diff --git a/cli/src/main/java/io/github/dfa1/vortex/cli/SchemaCommand.java b/cli/src/main/java/io/github/dfa1/vortex/cli/SchemaCommand.java index 5813dad53..da40dfab1 100644 --- a/cli/src/main/java/io/github/dfa1/vortex/cli/SchemaCommand.java +++ b/cli/src/main/java/io/github/dfa1/vortex/cli/SchemaCommand.java @@ -27,10 +27,10 @@ static int run(String[] args) { System.err.println("file not found: " + path); return ExitStatus.FILE_NOT_FOUND; } - try (VortexReader reader = VortexReader.open(path)) { + try (VortexReader reader = VortexReader.open(path, CliHandles.readRegistry())) { printSchema(System.out, path, reader.dtype(), reader.layout().rowCount()); return ExitStatus.OK; - } catch (IOException e) { + } catch (IOException | io.github.dfa1.vortex.core.error.VortexException e) { System.err.println("error: " + e.getMessage()); return ExitStatus.ERROR; } diff --git a/cli/src/main/java/io/github/dfa1/vortex/cli/SelectCommand.java b/cli/src/main/java/io/github/dfa1/vortex/cli/SelectCommand.java index 50ff246b6..bee4af6c9 100644 --- a/cli/src/main/java/io/github/dfa1/vortex/cli/SelectCommand.java +++ b/cli/src/main/java/io/github/dfa1/vortex/cli/SelectCommand.java @@ -48,7 +48,8 @@ static int run(String[] args) { System.err.println("error: " + e.getMessage()); return ExitStatus.USAGE_ERROR; } - ExportOptions options = ExportOptions.defaults().withColumns(parsed.columns()); + ExportOptions options = ExportOptions.defaults().withColumns(parsed.columns()) + .withReadRegistry(CliHandles.readRegistry()); try { Writer stdout = new OutputStreamWriter(System.out, StandardCharsets.UTF_8); if (parsed.conditions().isEmpty()) { diff --git a/cli/src/main/java/io/github/dfa1/vortex/cli/StatsCommand.java b/cli/src/main/java/io/github/dfa1/vortex/cli/StatsCommand.java index 292e466ea..a0adbecee 100644 --- a/cli/src/main/java/io/github/dfa1/vortex/cli/StatsCommand.java +++ b/cli/src/main/java/io/github/dfa1/vortex/cli/StatsCommand.java @@ -27,7 +27,7 @@ static int run(String[] args) { System.err.println("file not found: " + path); return ExitStatus.FILE_NOT_FOUND; } - try (VortexReader reader = VortexReader.open(path)) { + try (VortexReader reader = VortexReader.open(path, CliHandles.readRegistry())) { if (!(reader.dtype() instanceof DType.Struct schema)) { System.err.println("error: root dtype is not a struct"); return ExitStatus.ERROR; @@ -50,7 +50,7 @@ static int run(String[] args) { System.out.printf("%-20s %-12s %15s %15s%n", name.value(), type, min, max); } return ExitStatus.OK; - } catch (IOException e) { + } catch (IOException | io.github.dfa1.vortex.core.error.VortexException e) { System.err.println("error: " + e.getMessage()); return ExitStatus.ERROR; } diff --git a/cli/src/main/java/io/github/dfa1/vortex/cli/VortexCli.java b/cli/src/main/java/io/github/dfa1/vortex/cli/VortexCli.java index 4c57aaf0b..3bfb2732e 100644 --- a/cli/src/main/java/io/github/dfa1/vortex/cli/VortexCli.java +++ b/cli/src/main/java/io/github/dfa1/vortex/cli/VortexCli.java @@ -22,6 +22,9 @@ private VortexCli() { /// Global flag, accepted anywhere on the command line: prints how long the operation took. private static final String TIMING_FLAG = "--timing"; + /// Global flag, accepted anywhere on the command line: opens files whose extension types have no + /// decoder, reading those columns as their storage (`ReadRegistry.Builder#allowUnknown()`). + private static final String ALLOW_UNKNOWN_FLAG = "--allow-unknown"; public static void main(String[] args) { System.exit(execute(args)); @@ -33,7 +36,10 @@ public static void main(String[] args) { static int execute(String[] args) { long start = System.nanoTime(); boolean timing = Arrays.asList(args).contains(TIMING_FLAG); - String[] commandArgs = Arrays.stream(args).filter(arg -> !TIMING_FLAG.equals(arg)).toArray(String[]::new); + CliHandles.allowUnknown(Arrays.asList(args).contains(ALLOW_UNKNOWN_FLAG)); + String[] commandArgs = Arrays.stream(args) + .filter(arg -> !TIMING_FLAG.equals(arg) && !ALLOW_UNKNOWN_FLAG.equals(arg)) + .toArray(String[]::new); int exit = dispatch(commandArgs); if (timing) { Duration elapsed = Duration.ofNanos(System.nanoTime() - start); @@ -67,7 +73,7 @@ private static int dispatch(String[] args) { } static void printUsage(PrintStream out) { - out.println("Usage: java -jar vortex-cli--all.jar [args] [--timing]"); + out.println("Usage: java -jar vortex-cli--all.jar [args] [--timing] [--allow-unknown]"); out.println(" inspect [--html] print file structure; --html writes an HTML report"); out.println(" tui open interactive inspector; url is http(s)://"); out.println(" view open scrollable data grid; url is http(s)://"); @@ -80,5 +86,7 @@ static void printUsage(PrintStream out) { out.println(" stats print per-column min/max statistics"); out.println(" filter filter rows to CSV (e.g. \"price >= 100\")"); out.println("Any subcommand also takes --timing: print the elapsed time to stderr, JVM start-up excluded."); + out.println("Any subcommand also takes --allow-unknown: open a file whose extension types have no decoder,"); + out.println(" reading those columns as their storage."); } } diff --git a/cli/src/main/java/io/github/dfa1/vortex/cli/tui/GridRender.java b/cli/src/main/java/io/github/dfa1/vortex/cli/tui/GridRender.java index 80b490b45..8e602f8d3 100644 --- a/cli/src/main/java/io/github/dfa1/vortex/cli/tui/GridRender.java +++ b/cli/src/main/java/io/github/dfa1/vortex/cli/tui/GridRender.java @@ -99,16 +99,17 @@ private static boolean isUnsigned(Array arr) { } private static Optional formatExtension(DType.Extension ext, Array storage, long i) { - Optional idOpt = ExtensionId.parse(ext.extensionId()); - if (idOpt.isEmpty()) { + if (!(ExtensionId.parse(ext.extensionId()) instanceof ExtensionId.WellKnown id)) { return Optional.empty(); } - return Optional.of(switch (idOpt.get()) { - case VORTEX_DATE -> DateExtensionDecoder.INSTANCE.decode(storage, i).toString(); - case VORTEX_TIME -> TimeExtensionDecoder.INSTANCE.decode(ext, storage, i).toString(); - case VORTEX_TIMESTAMP -> TimestampExtensionDecoder.INSTANCE.instant(ext, storage, i).toString(); - case VORTEX_UUID -> UuidExtensionDecoder.INSTANCE.decode(storage, i).toString(); - }); + return switch (id) { + case VORTEX_DATE -> Optional.of(DateExtensionDecoder.INSTANCE.decode(ext, storage, i).toString()); + case VORTEX_TIME -> Optional.of(TimeExtensionDecoder.INSTANCE.decode(ext, storage, i).toString()); + case VORTEX_TIMESTAMP -> Optional.of(TimestampExtensionDecoder.INSTANCE.instant(ext, storage, i).toString()); + case VORTEX_UUID -> Optional.of(UuidExtensionDecoder.INSTANCE.decode(storage, i).toString()); + // JSON text and geometries have no Java domain type yet: render the storage + default -> Optional.empty(); + }; } private static String bytesToHex(byte[] bytes) { diff --git a/cli/src/main/java/io/github/dfa1/vortex/cli/tui/InspectorRender.java b/cli/src/main/java/io/github/dfa1/vortex/cli/tui/InspectorRender.java index ba67b323b..20cde0413 100644 --- a/cli/src/main/java/io/github/dfa1/vortex/cli/tui/InspectorRender.java +++ b/cli/src/main/java/io/github/dfa1/vortex/cli/tui/InspectorRender.java @@ -39,11 +39,9 @@ private InspectorRender() { /// @return formatted value, or an angle-bracketed fallback for unknown shapes static String formatValue(Array array, int i, DType declared) { if (declared instanceof DType.Extension ext - && ExtensionId.parse(ext.extensionId()) - .filter(id -> id == ExtensionId.VORTEX_DATE) - .isPresent()) { + && ExtensionId.parse(ext.extensionId()) == ExtensionId.VORTEX_DATE) { try { - return DateExtensionDecoder.INSTANCE.decode(array, i).toString(); + return DateExtensionDecoder.INSTANCE.decode(ext, array, i).toString(); } catch (RuntimeException _) { // fall through to generic rendering on shape mismatch } diff --git a/cli/src/test/java/io/github/dfa1/vortex/cli/VortexCliTest.java b/cli/src/test/java/io/github/dfa1/vortex/cli/VortexCliTest.java index d0e758ff6..037b62952 100644 --- a/cli/src/test/java/io/github/dfa1/vortex/cli/VortexCliTest.java +++ b/cli/src/test/java/io/github/dfa1/vortex/cli/VortexCliTest.java @@ -1,13 +1,21 @@ package io.github.dfa1.vortex.cli; +import io.github.dfa1.vortex.core.model.ColumnName; +import io.github.dfa1.vortex.core.model.DType; +import io.github.dfa1.vortex.writer.VortexWriter; +import io.github.dfa1.vortex.writer.WriteOptions; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.PrintStream; +import java.nio.channels.FileChannel; import java.nio.charset.StandardCharsets; import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.List; +import java.util.Map; import static io.github.dfa1.vortex.cli.CliTestSupport.capture; import static io.github.dfa1.vortex.cli.CliTestSupport.writeSmallVortex; @@ -39,7 +47,8 @@ void printUsage_listsEverySubcommand() { .contains("select") .contains("stats") .contains("filter") - .contains("--timing"); + .contains("--timing") + .contains("--allow-unknown"); } @Test @@ -105,4 +114,58 @@ void timingAlone_isAUsageError() { assertThat(result.status()).isEqualTo(ExitStatus.USAGE_ERROR); assertThat(result.stderr()).contains("Usage:"); } + + @Test + void customExtension_withoutAllowUnknown_isRefused(@TempDir Path tmp) throws IOException { + // Given — a column of an extension type with no decoder; Rust refuses it unless the session + // allows unknown ids, and so does the CLI by default + Path file = writeCustomExtensionVortex(tmp); + + // When + CliTestSupport.Captured result = capture(() -> VortexCli.execute(new String[]{"count", file.toString()})); + + // Then + assertThat(result.status()).isNotEqualTo(ExitStatus.OK); + assertThat(result.stderr()).contains("no extension decoder registered for acme.money"); + } + + @Test + void customExtension_withAllowUnknown_readsTheStorage(@TempDir Path tmp) throws IOException { + // Given + Path file = writeCustomExtensionVortex(tmp); + + // When — the flag reaches the CSV path too, which opens the file itself + CliTestSupport.Captured result = capture(() -> VortexCli.execute( + new String[]{"select", file.toString(), "price", "--allow-unknown"})); + + // Then — the column reads as its I64 storage + assertThat(result.status()).isEqualTo(ExitStatus.OK); + assertThat(result.stdout()).contains("999").contains("12050"); + } + + @Test + void allowUnknown_doesNotLeakIntoTheNextInvocation(@TempDir Path tmp) throws IOException { + // Given — the registry is process state, reset by every execute + Path file = writeCustomExtensionVortex(tmp); + capture(() -> VortexCli.execute(new String[]{"count", file.toString(), "--allow-unknown"})); + + // When + CliTestSupport.Captured result = capture(() -> VortexCli.execute(new String[]{"count", file.toString()})); + + // Then + assertThat(result.status()).isNotEqualTo(ExitStatus.OK); + } + + private static Path writeCustomExtensionVortex(Path dir) throws IOException { + Path file = dir.resolve("money.vortex"); + DType.Struct schema = new DType.Struct(List.of(ColumnName.of("price")), + List.of(new DType.Extension("acme.money", DType.I64, + java.lang.foreign.MemorySegment.ofArray("EUR".getBytes(StandardCharsets.US_ASCII)), false)), + false); + try (FileChannel ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE); + VortexWriter writer = VortexWriter.create(ch, schema, WriteOptions.defaults())) { + writer.writeChunk(Map.of(ColumnName.of("price"), new long[]{999L, 12_050L})); + } + return file; + } } diff --git a/cli/src/test/java/io/github/dfa1/vortex/cli/tui/GridRenderTest.java b/cli/src/test/java/io/github/dfa1/vortex/cli/tui/GridRenderTest.java index 809c0cbf8..423c1e4e8 100644 --- a/cli/src/test/java/io/github/dfa1/vortex/cli/tui/GridRenderTest.java +++ b/cli/src/test/java/io/github/dfa1/vortex/cli/tui/GridRenderTest.java @@ -101,8 +101,7 @@ void rendersDecimal() { void rendersDateExtensionFromIntStorage() { try (Arena arena = Arena.ofConfined()) { // Given — vortex.date over I32 epoch-day storage: day 0 = 1970-01-01 - DType dateExt = new DType.Extension("vortex.date", - DType.I32, null, false); + DType dateExt = io.github.dfa1.vortex.reader.extension.DateExtensionDecoder.INSTANCE.dtype(false); IntArray storage = ArrayFixtures.ints(arena, 0, 1); // When / Then diff --git a/cli/src/test/java/io/github/dfa1/vortex/cli/tui/InspectorRenderTest.java b/cli/src/test/java/io/github/dfa1/vortex/cli/tui/InspectorRenderTest.java index 8e2e3d134..b47ee4f64 100644 --- a/cli/src/test/java/io/github/dfa1/vortex/cli/tui/InspectorRenderTest.java +++ b/cli/src/test/java/io/github/dfa1/vortex/cli/tui/InspectorRenderTest.java @@ -111,8 +111,7 @@ void genericArrayDecimalBadShapeRendersFallback() { void rendersDateExtension() { try (Arena arena = Arena.ofConfined()) { // Given — vortex.date over I32 epoch-day storage - DType dateExt = new DType.Extension("vortex.date", - DType.I32, null, false); + DType dateExt = io.github.dfa1.vortex.reader.extension.DateExtensionDecoder.INSTANCE.dtype(false); IntArray storage = ArrayFixtures.ints(arena, 0); // When / Then — day 0 = 1970-01-01 diff --git a/core/src/main/java/io/github/dfa1/vortex/core/model/ExtensionId.java b/core/src/main/java/io/github/dfa1/vortex/core/model/ExtensionId.java index baa112fe7..d182e58ca 100644 --- a/core/src/main/java/io/github/dfa1/vortex/core/model/ExtensionId.java +++ b/core/src/main/java/io/github/dfa1/vortex/core/model/ExtensionId.java @@ -1,54 +1,312 @@ package io.github.dfa1.vortex.core.model; +import io.github.dfa1.vortex.core.error.VortexException; + +import java.io.Serializable; +import java.lang.foreign.ValueLayout; +import java.nio.ByteBuffer; +import java.nio.charset.CharacterCodingException; +import java.nio.charset.StandardCharsets; import java.util.Map; -import java.util.Optional; +import java.util.Objects; import java.util.function.Function; import java.util.stream.Collectors; import java.util.stream.Stream; -/// Strongly-typed identifier for the spec-defined Vortex extension types. -/// Mirrors [io.github.dfa1.vortex.core.model.EncodingId] in shape so callers -/// reach for the same pattern across both registries. +/// Identity of an extension type — either a [WellKnown] constant or a third-party [Custom] id. /// -/// Unknown wire ids (e.g. `"acme.geopoint"`) are not represented here — -/// they flow through the registry's unknown-passthrough path. -public enum ExtensionId { - /// `vortex.date`. - VORTEX_DATE("vortex.date"), - /// `vortex.time`. - VORTEX_TIME("vortex.time"), - /// `vortex.timestamp`. - VORTEX_TIMESTAMP("vortex.timestamp"), - /// `vortex.uuid`. - VORTEX_UUID("vortex.uuid"); - - private static final Map LOOKUP = Stream.of(values()) - .collect(Collectors.toUnmodifiableMap(ExtensionId::id, Function.identity())); - - private final String id; - - ExtensionId(String id) { - this.id = id; - } +/// [WellKnown] holds every id the vortex-jni reference session registers by default: the +/// temporal types and `vortex.uuid`, plus `vortex.json` and the `vortex.st.*` geometry types. A +/// file with one of them opens without a registered decoder, the column reading as its storage. +/// Mirrors [EncodingId] in shape so callers reach for the same pattern across both registries. +/// +/// The wire representation is always a string (e.g. `"vortex.timestamp"`); [#parse(String)] maps any +/// such string to a typed value, and [#id()] recovers the wire string from a typed value. +public sealed interface ExtensionId extends Serializable permits ExtensionId.WellKnown, ExtensionId.Custom { - /// Parses a raw extension id string into the matching spec constant. - /// Callers that demand a known id chain `.orElseThrow(...)`. + /// Returns the wire string of this extension id (e.g. `"vortex.timestamp"`). /// - /// @param id raw extension id string (e.g. `"vortex.timestamp"`) - /// @return matching constant, or empty if not a known spec extension - public static Optional parse(String id) { - return Optional.ofNullable(LOOKUP.get(id)); - } + /// @return the wire string of this extension id + String id(); - /// Returns the canonical wire-format id string. + /// Parses a wire string into its typed representation: the matching [WellKnown] constant, + /// else a [Custom] wrapping the raw string. Total over every valid id; callers parsing + /// untrusted input must guard invalid ids and raise their own domain error. /// - /// @return canonical id (e.g. `"vortex.timestamp"`) - public String id() { - return id; + /// @param raw the raw extension id string (e.g. `"vortex.timestamp"`) + /// @return the matching [WellKnown] constant, or a [Custom] wrapping `raw` if none matches + /// @throws NullPointerException if `raw` is `null` + /// @throws IllegalArgumentException if `raw` is blank or contains an ISO control character + static ExtensionId parse(String raw) { + WellKnown known = WellKnown.byId(raw); + return known != null ? known : new Custom(raw); } - @Override - public String toString() { - return id; + /// Extension ids the vortex-jni reference session registers by default. + enum WellKnown implements ExtensionId { + /// `vortex.date`. + VORTEX_DATE("vortex.date"), + /// `vortex.time`. + VORTEX_TIME("vortex.time"), + /// `vortex.timestamp`. + VORTEX_TIMESTAMP("vortex.timestamp"), + /// `vortex.uuid`. + VORTEX_UUID("vortex.uuid"), + /// `vortex.json`: a JSON document per row over `Utf8` storage (Rust's `vortex-json`). + VORTEX_JSON("vortex.json"), + /// `vortex.st.wkb`: a Well-Known Binary geometry per row over `Binary` storage (Rust's `vortex-spatial`). + VORTEX_ST_WKB("vortex.st.wkb"), + /// `vortex.st.point` (Rust's `vortex-spatial`). + VORTEX_ST_POINT("vortex.st.point"), + /// `vortex.st.linestring` (Rust's `vortex-spatial`). + VORTEX_ST_LINESTRING("vortex.st.linestring"), + /// `vortex.st.polygon` (Rust's `vortex-spatial`). + VORTEX_ST_POLYGON("vortex.st.polygon"), + /// `vortex.st.multipoint` (Rust's `vortex-spatial`). + VORTEX_ST_MULTIPOINT("vortex.st.multipoint"), + /// `vortex.st.multilinestring` (Rust's `vortex-spatial`). + VORTEX_ST_MULTILINESTRING("vortex.st.multilinestring"), + /// `vortex.st.multipolygon` (Rust's `vortex-spatial`). + VORTEX_ST_MULTIPOLYGON("vortex.st.multipolygon"), + /// `vortex.st.box` (Rust's `vortex-spatial`). + VORTEX_ST_BOX("vortex.st.box"); + + // GeoArrow field names per dimension (vortex-spatial coordinate.rs and rect.rs) + private static final java.util.List> COORDINATE_FIELDS = java.util.List.of( + java.util.List.of("x", "y"), java.util.List.of("x", "y", "z"), + java.util.List.of("x", "y", "m"), java.util.List.of("x", "y", "z", "m")); + private static final java.util.List> BOX_FIELDS = java.util.List.of( + java.util.List.of("xmin", "ymin", "xmax", "ymax"), + java.util.List.of("xmin", "ymin", "zmin", "xmax", "ymax", "zmax"), + java.util.List.of("xmin", "ymin", "mmin", "xmax", "ymax", "mmax"), + java.util.List.of("xmin", "ymin", "zmin", "mmin", "xmax", "ymax", "zmax", "mmax")); + + private static final Map LOOKUP = Stream.of(values()) + .collect(Collectors.toUnmodifiableMap(WellKnown::id, Function.identity())); + + private final String id; + + WellKnown(String id) { + this.id = id; + } + + /// Returns the well-known constant whose wire string is `id`, or `null` if none matches. + /// + /// @param id the wire string to look up (may be `null`) + /// @return the matching constant, or `null` if unrecognized + static WellKnown byId(String id) { + return LOOKUP.get(id); + } + + @Override + public String id() { + return id; + } + + @Override + public String toString() { + return id; + } + + /// Rejects a dtype this extension cannot represent, with the rules of Rust's + /// `deserialize_metadata` and `validate_dtype` for it: the time-unit byte and storage ptype + /// of `vortex.date`/`vortex.time`, the metadata layout and `I64` storage of + /// `vortex.timestamp`, the version byte and `FixedSizeList` storage of `vortex.uuid`, + /// empty metadata over `Utf8` for `vortex.json`, `Binary` storage for `vortex.st.wkb`, and + /// the GeoArrow coordinate-struct storage of the other `vortex.st.*` geometries. + /// + /// @param dtype an extension dtype with this id + /// @throws VortexException if the metadata or storage dtype is invalid for this extension + public void validate(DType.Extension dtype) { + byte[] meta = dtype.metadata() == null ? new byte[0] : dtype.metadata().toArray(ValueLayout.JAVA_BYTE); + DType storage = dtype.storageDType(); + switch (this) { + case VORTEX_DATE -> { + TimeUnit unit = unit(meta); + PType expected = switch (unit) { + case Days -> PType.I32; + case Milliseconds -> PType.I64; + case Nanoseconds, Microseconds, Seconds -> throw new VortexException( + id + " does not support time unit " + unit); + }; + requirePtype(storage, expected, unit); + } + case VORTEX_TIME -> { + TimeUnit unit = unit(meta); + PType expected = switch (unit) { + case Seconds, Milliseconds -> PType.I32; + case Microseconds, Nanoseconds -> PType.I64; + case Days -> throw new VortexException(id + " does not support time unit " + unit); + }; + requirePtype(storage, expected, unit); + } + case VORTEX_TIMESTAMP -> { + if (meta.length < 3) { + throw new VortexException(id + " metadata must have at least 3 bytes, got " + meta.length); + } + unit(meta); + int tzLen = Byte.toUnsignedInt(meta[1]) | Byte.toUnsignedInt(meta[2]) << 8; + if (meta.length < 3 + tzLen) { + throw new VortexException(id + " metadata is truncated: declared timezone length " + + tzLen + " but only " + (meta.length - 3) + " bytes available"); + } + try { + StandardCharsets.UTF_8.newDecoder().decode(ByteBuffer.wrap(meta, 3, tzLen)); + } catch (CharacterCodingException e) { + throw new VortexException(id + " timezone is not valid UTF-8", e); + } + if (!(storage instanceof DType.Primitive p && p.ptype() == PType.I64)) { + throw new VortexException(id + " storage dtype must be I64, got " + storage); + } + } + case VORTEX_UUID -> { + if (meta.length > 1) { + throw new VortexException(id + " metadata must be 0 or 1 bytes, got " + meta.length); + } + // uuid crate versions 0-8, plus Max as 0x0f or the older 0xff + if (meta.length == 1 && !(meta[0] >= 0 && meta[0] <= 8 || meta[0] == 0x0f || meta[0] == (byte) 0xff)) { + throw new VortexException("unknown UUID version discriminant: " + Byte.toUnsignedInt(meta[0])); + } + if (!(storage instanceof DType.FixedSizeList list && list.fixedSize() == 16 + && list.elementType() instanceof DType.Primitive p && p.ptype() == PType.U8 && !p.nullable())) { + throw new VortexException(id + " storage dtype must be FixedSizeList, got " + storage); + } + } + case VORTEX_JSON -> { + if (meta.length != 0) { + throw new VortexException(id + " metadata must be empty"); + } + if (!(storage instanceof DType.Utf8)) { + throw new VortexException(id + " storage dtype must be Utf8, got " + storage); + } + } + // ponytail: the spatial metadata (a SpatialMetadata protobuf: CRS, edges) is not decoded, + // so malformed metadata Rust refuses still opens here; decode it when Java reads geometries + case VORTEX_ST_WKB -> { + if (!(storage instanceof DType.Binary)) { + throw new VortexException(id + " storage dtype must be Binary, got " + storage); + } + } + case VORTEX_ST_POINT -> requireF64Struct(storage, COORDINATE_FIELDS); + case VORTEX_ST_LINESTRING, VORTEX_ST_MULTIPOINT -> requireF64Struct(listElement(storage, 1), COORDINATE_FIELDS); + case VORTEX_ST_POLYGON, VORTEX_ST_MULTILINESTRING -> requireF64Struct(listElement(storage, 2), COORDINATE_FIELDS); + case VORTEX_ST_MULTIPOLYGON -> requireF64Struct(listElement(storage, 3), COORDINATE_FIELDS); + case VORTEX_ST_BOX -> requireF64Struct(storage, BOX_FIELDS); + } + } + + /// Unwraps `depth` levels of `List`: a linestring is a list of coordinates, a polygon a list + /// of rings, a multipolygon a list of polygons. + private DType listElement(DType storage, int depth) { + DType current = storage; + for (int i = 0; i < depth; i++) { + if (!(current instanceof DType.List list)) { + throw new VortexException(id + " storage must nest " + depth + " List level(s) over coordinates, got " + storage); + } + current = list.elementType(); + } + return current; + } + + /// A coordinate or box: a `Struct` of non-nullable `F64` fields named for one GeoArrow + /// dimension (xy, xyz, xym or xyzm). + private void requireF64Struct(DType dtype, java.util.List> dimensions) { + if (!(dtype instanceof DType.Struct struct)) { + throw new VortexException(id + " coordinates must be a Struct, got " + dtype); + } + for (DType field : struct.fieldTypes()) { + if (!(field instanceof DType.Primitive p && p.ptype() == PType.F64 && !p.nullable())) { + throw new VortexException(id + " coordinate fields must be non-nullable F64, got " + field); + } + } + java.util.List names = struct.fieldNames().stream().map(ColumnName::value).toList(); + if (!dimensions.contains(names)) { + throw new VortexException(id + " is not a valid GeoArrow dimension: " + names); + } + } + + private TimeUnit unit(byte[] meta) { + if (meta.length == 0) { + throw new VortexException(id + " metadata must not be empty"); + } + try { + return TimeUnit.fromTag(meta[0]); + } catch (IllegalArgumentException e) { + throw new VortexException(id + " metadata has an invalid time unit", e); + } + } + + private void requirePtype(DType storage, PType expected, TimeUnit unit) { + if (!(storage instanceof DType.Primitive p && p.ptype() == expected)) { + throw new VortexException(id + " storage dtype for " + unit + " must be " + expected + ", got " + storage); + } + } } + + /// A third-party extension id whose wire string is not part of the [WellKnown] set. + /// + /// @param id the wire string of this extension id; must be non-blank, free of ISO control + /// characters, and must not collide with a [WellKnown] wire string + record Custom(String id) implements ExtensionId { + + /// Validates that `id` is a usable custom extension id, with the same policy as + /// [EncodingId.Custom]. + /// + /// @param id the wire string of this extension id + /// @throws NullPointerException if `id` is `null` + /// @throws IllegalArgumentException if `id` is blank, contains an ISO control character, or + /// matches a [WellKnown] wire string + public Custom { + Objects.requireNonNull(id, "id"); + if (id.isBlank()) { + throw new IllegalArgumentException("extension id must not be blank"); + } + for (int i = 0; i < id.length(); i++) { + if (Character.isISOControl(id.charAt(i))) { + throw new IllegalArgumentException( + "extension id contains control character U+%04X".formatted((int) id.charAt(i))); + } + } + WellKnown wellKnown = WellKnown.byId(id); + if (wellKnown != null) { + throw new IllegalArgumentException( + "\"" + id + "\" is a well-known extension id; use ExtensionId." + wellKnown.name() + " instead"); + } + } + + @Override + public String toString() { + return id; + } + } + + // Re-export every WellKnown constant, typed as WellKnown, so `ExtensionId.VORTEX_FOO` call + // sites keep compiling and remain usable wherever a WellKnown is required. + + /// Well-known `vortex.date` id. + WellKnown VORTEX_DATE = WellKnown.VORTEX_DATE; + /// Well-known `vortex.time` id. + WellKnown VORTEX_TIME = WellKnown.VORTEX_TIME; + /// Well-known `vortex.timestamp` id. + WellKnown VORTEX_TIMESTAMP = WellKnown.VORTEX_TIMESTAMP; + /// Well-known `vortex.uuid` id. + WellKnown VORTEX_UUID = WellKnown.VORTEX_UUID; + /// Well-known `vortex.json` id. + WellKnown VORTEX_JSON = WellKnown.VORTEX_JSON; + /// Well-known `vortex.st.wkb` id. + WellKnown VORTEX_ST_WKB = WellKnown.VORTEX_ST_WKB; + /// Well-known `vortex.st.point` id. + WellKnown VORTEX_ST_POINT = WellKnown.VORTEX_ST_POINT; + /// Well-known `vortex.st.linestring` id. + WellKnown VORTEX_ST_LINESTRING = WellKnown.VORTEX_ST_LINESTRING; + /// Well-known `vortex.st.polygon` id. + WellKnown VORTEX_ST_POLYGON = WellKnown.VORTEX_ST_POLYGON; + /// Well-known `vortex.st.multipoint` id. + WellKnown VORTEX_ST_MULTIPOINT = WellKnown.VORTEX_ST_MULTIPOINT; + /// Well-known `vortex.st.multilinestring` id. + WellKnown VORTEX_ST_MULTILINESTRING = WellKnown.VORTEX_ST_MULTILINESTRING; + /// Well-known `vortex.st.multipolygon` id. + WellKnown VORTEX_ST_MULTIPOLYGON = WellKnown.VORTEX_ST_MULTIPOLYGON; + /// Well-known `vortex.st.box` id. + WellKnown VORTEX_ST_BOX = WellKnown.VORTEX_ST_BOX; } diff --git a/core/src/test/java/io/github/dfa1/vortex/core/model/ExtensionIdTest.java b/core/src/test/java/io/github/dfa1/vortex/core/model/ExtensionIdTest.java index 390c58901..31e9ba75b 100644 --- a/core/src/test/java/io/github/dfa1/vortex/core/model/ExtensionIdTest.java +++ b/core/src/test/java/io/github/dfa1/vortex/core/model/ExtensionIdTest.java @@ -1,39 +1,228 @@ package io.github.dfa1.vortex.core.model; +import io.github.dfa1.vortex.core.error.VortexException; +import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; + +import java.lang.foreign.MemorySegment; +import java.util.List; +import java.util.stream.Stream; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; class ExtensionIdTest { - @ParameterizedTest - @CsvSource({ - "vortex.date,VORTEX_DATE", - "vortex.time,VORTEX_TIME", - "vortex.timestamp,VORTEX_TIMESTAMP", - "vortex.uuid,VORTEX_UUID" - }) - void parse_knownIds_returnEnumConstant(String wire, ExtensionId expected) { - // Given / When / Then — wire string round-trips to the enum constant - // so the LOOKUP map stays in sync with the enum definition - assertThat(ExtensionId.parse(wire)).contains(expected); + @Nested + class Parse { + + @ParameterizedTest + @EnumSource(ExtensionId.WellKnown.class) + void parse_knownId_returnsMatchingConstant(ExtensionId.WellKnown id) { + // Given the wire string of a well-known constant + // When + ExtensionId result = ExtensionId.parse(id.id()); + // Then the same constant comes back + assertThat(result).isSameAs(id); + } + + @Test + void parse_unknownId_returnsCustomWrappingRawId() { + // Given a wire string no build knows about + String raw = "acme.geopoint"; + // When — parse is total, so a miss is a typed Custom rather than an empty Optional + ExtensionId result = ExtensionId.parse(raw); + // Then + assertThat(result).isEqualTo(new ExtensionId.Custom(raw)); + assertThat(result.id()).isEqualTo(raw); + } + + @ParameterizedTest + @ValueSource(strings = {"", " ", " "}) + void parse_blankId_throwsIllegalArgumentException(String blank) { + // Given / When / Then — blank is not a valid id; parse must not silently wrap it, + // so untrusted-input callers are forced to guard it into their own domain error + assertThatThrownBy(() -> ExtensionId.parse(blank)) + .isInstanceOf(IllegalArgumentException.class); + } + } + + @Nested + class CustomInvariants { + + @Test + void construct_nullId_throwsNullPointerException() { + // Given / When / Then — a Custom must always carry a wire string + assertThatThrownBy(() -> new ExtensionId.Custom(null)) + .isInstanceOf(NullPointerException.class); + } + + @Test + void construct_blankId_throwsIllegalArgumentException() { + // Given / When / Then — blank ids have no wire representation + assertThatThrownBy(() -> new ExtensionId.Custom(" ")) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + void construct_wellKnownId_throwsIllegalArgumentException() { + // Given a wire string that already names a well-known constant + // When / Then — Custom refuses to shadow it and points at the constant to use instead + assertThatThrownBy(() -> new ExtensionId.Custom("vortex.date")) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("VORTEX_DATE"); + } + + @ParameterizedTest + @ValueSource(strings = {"a\nb", "tab\there"}) + void custom_controlCharId_throwsIllegalArgumentException(String id) { + // Given / When / Then — a control char would write a file PostscriptParser crashes on, + // so Custom rejects it just as ColumnName does + assertThatThrownBy(() -> new ExtensionId.Custom(id)) + .isInstanceOf(IllegalArgumentException.class); + } } - @Test - void parse_unknownId_returnsEmpty() { - // Given — open-world extension id; library doesn't recognize it - // When / Then — non-throwing miss so the registry can route to passthrough - assertThat(ExtensionId.parse("acme.geopoint")).isEmpty(); + @Nested + class Properties { + + @ParameterizedTest + @EnumSource(ExtensionId.WellKnown.class) + void id_isNonBlankString(ExtensionId.WellKnown id) { + // Given / When / Then + assertThat(id.id()).isNotBlank(); + } + + @ParameterizedTest + @EnumSource(ExtensionId.WellKnown.class) + void toString_equalsId(ExtensionId.WellKnown id) { + // Given / When / Then + assertThat(id).hasToString(id.id()); + } } - @Test - void id_matchesEnumStringRepresentation() { - // Given / When / Then — id() is the wire-format string used in toString() - // so debug output and switch-case keys stay aligned - for (ExtensionId id : ExtensionId.values()) { - assertThat(id.id()).isEqualTo(id.toString()); + /// Rust's `deserialize_metadata` + `validate_dtype` rules for the four spec extensions + /// (vortex-array/src/extension/datetime/*.rs and encodings/uuid/src/vtable.rs at 0.88.0). + @Nested + class Validate { + + static Stream valid() { + return Stream.of( + Arguments.of("date days I32", ext(ExtensionId.VORTEX_DATE, DType.I32, 4)), + Arguments.of("date millis I64", ext(ExtensionId.VORTEX_DATE, DType.I64, 2)), + Arguments.of("time seconds I32", ext(ExtensionId.VORTEX_TIME, DType.I32, 3)), + Arguments.of("time millis I32", ext(ExtensionId.VORTEX_TIME, DType.I32, 2)), + Arguments.of("time micros I64", ext(ExtensionId.VORTEX_TIME, DType.I64, 1)), + Arguments.of("time nanos I64", ext(ExtensionId.VORTEX_TIME, DType.I64, 0)), + Arguments.of("timestamp no tz", ext(ExtensionId.VORTEX_TIMESTAMP, DType.I64, 2, 0, 0)), + Arguments.of("timestamp tz", ext(ExtensionId.VORTEX_TIMESTAMP, DType.I64, 2, 3, 0, 'U', 'T', 'C')), + Arguments.of("timestamp days", ext(ExtensionId.VORTEX_TIMESTAMP, DType.I64, 4, 0, 0)), + Arguments.of("uuid no version", ext(ExtensionId.VORTEX_UUID, UUID_STORAGE)), + Arguments.of("uuid v4", ext(ExtensionId.VORTEX_UUID, UUID_STORAGE, 4)), + Arguments.of("uuid max 0x0f", ext(ExtensionId.VORTEX_UUID, UUID_STORAGE, 0x0f)), + Arguments.of("uuid max 0xff", ext(ExtensionId.VORTEX_UUID, UUID_STORAGE, 0xff)), + Arguments.of("date nullable storage", ext(ExtensionId.VORTEX_DATE, DType.I32.asNullable(), 4)), + Arguments.of("json", ext(ExtensionId.VORTEX_JSON, DType.UTF8)), + Arguments.of("wkb", ext(ExtensionId.VORTEX_ST_WKB, DType.BINARY)), + Arguments.of("point xy", ext(ExtensionId.VORTEX_ST_POINT, coords("x", "y"))), + Arguments.of("point xyzm", ext(ExtensionId.VORTEX_ST_POINT, coords("x", "y", "z", "m"))), + Arguments.of("linestring", ext(ExtensionId.VORTEX_ST_LINESTRING, lists(coords("x", "y"), 1))), + Arguments.of("multipoint", ext(ExtensionId.VORTEX_ST_MULTIPOINT, lists(coords("x", "y", "z"), 1))), + Arguments.of("polygon", ext(ExtensionId.VORTEX_ST_POLYGON, lists(coords("x", "y"), 2))), + Arguments.of("multilinestring", ext(ExtensionId.VORTEX_ST_MULTILINESTRING, lists(coords("x", "y", "m"), 2))), + Arguments.of("multipolygon", ext(ExtensionId.VORTEX_ST_MULTIPOLYGON, lists(coords("x", "y"), 3))), + Arguments.of("box xy", ext(ExtensionId.VORTEX_ST_BOX, coords("xmin", "ymin", "xmax", "ymax")))); + } + + static Stream invalid() { + return Stream.of( + Arguments.of(ext(ExtensionId.VORTEX_DATE, DType.I32), "metadata must not be empty"), + Arguments.of(ext(ExtensionId.VORTEX_DATE, DType.I32, 9), "invalid time unit"), + Arguments.of(ext(ExtensionId.VORTEX_DATE, DType.I64, 4), "for Days must be I32"), + Arguments.of(ext(ExtensionId.VORTEX_DATE, DType.I64, 0), "does not support time unit Nanoseconds"), + Arguments.of(ext(ExtensionId.VORTEX_DATE, DType.UTF8, 4), "must be I32"), + Arguments.of(ext(ExtensionId.VORTEX_TIME, DType.I64, 2), "for Milliseconds must be I32"), + Arguments.of(ext(ExtensionId.VORTEX_TIME, DType.I32, 4), "does not support time unit Days"), + Arguments.of(ext(ExtensionId.VORTEX_TIMESTAMP, DType.I32, 2, 0, 0), "storage dtype must be I64"), + Arguments.of(ext(ExtensionId.VORTEX_TIMESTAMP, DType.I64, 2), "at least 3 bytes"), + Arguments.of(ext(ExtensionId.VORTEX_TIMESTAMP, DType.I64, 2, 5, 0, 'U'), "truncated"), + Arguments.of(ext(ExtensionId.VORTEX_TIMESTAMP, DType.I64, 2, 1, 0, 0xc3), "not valid UTF-8"), + Arguments.of(ext(ExtensionId.VORTEX_UUID, UUID_STORAGE, 4, 0), "0 or 1 bytes"), + Arguments.of(ext(ExtensionId.VORTEX_UUID, UUID_STORAGE, 9), "unknown UUID version"), + Arguments.of(ext(ExtensionId.VORTEX_UUID, new DType.FixedSizeList(DType.U8, 8, false)), "FixedSizeList"), + Arguments.of(ext(ExtensionId.VORTEX_UUID, new DType.FixedSizeList(DType.U8.asNullable(), 16, false)), "FixedSizeList"), + Arguments.of(ext(ExtensionId.VORTEX_UUID, DType.I64), "FixedSizeList"), + Arguments.of(ext(ExtensionId.VORTEX_JSON, DType.BINARY), "storage dtype must be Utf8"), + Arguments.of(ext(ExtensionId.VORTEX_JSON, DType.UTF8, 1), "metadata must be empty"), + Arguments.of(ext(ExtensionId.VORTEX_ST_WKB, DType.UTF8), "storage dtype must be Binary"), + Arguments.of(ext(ExtensionId.VORTEX_ST_POINT, coords("y", "x")), "not a valid GeoArrow dimension"), + Arguments.of(ext(ExtensionId.VORTEX_ST_POINT, new DType.Struct( + List.of(ColumnName.of("x"), ColumnName.of("y")), List.of(DType.F64, DType.F32), false)), + "must be non-nullable F64"), + Arguments.of(ext(ExtensionId.VORTEX_ST_POINT, DType.F64), "must be a Struct"), + Arguments.of(ext(ExtensionId.VORTEX_ST_POLYGON, lists(coords("x", "y"), 1)), "must nest 2 List level(s)"), + Arguments.of(ext(ExtensionId.VORTEX_ST_MULTIPOLYGON, coords("x", "y")), "must nest 3 List level(s)"), + Arguments.of(ext(ExtensionId.VORTEX_ST_BOX, coords("x", "y")), "not a valid GeoArrow dimension")); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("valid") + void validate_validDtype_passes(String name, DType.Extension dtype) { + // Given a dtype Rust accepts + ExtensionId.WellKnown sut = (ExtensionId.WellKnown) ExtensionId.parse(dtype.extensionId()); + + // When / Then + sut.validate(dtype); + } + + @ParameterizedTest + @MethodSource("invalid") + void validate_invalidDtype_throws(DType.Extension dtype, String message) { + // Given a dtype Rust refuses + ExtensionId.WellKnown sut = (ExtensionId.WellKnown) ExtensionId.parse(dtype.extensionId()); + + // When / Then + assertThatThrownBy(() -> sut.validate(dtype)) + .isInstanceOf(VortexException.class) + .hasMessageContaining(message); + } + + @Test + void validate_nullMetadata_isTreatedAsEmpty() { + // Given — Java callers may pass null; it means "no metadata", which uuid allows + DType.Extension dtype = new DType.Extension("vortex.uuid", UUID_STORAGE, null, false); + + // When / Then + ExtensionId.VORTEX_UUID.validate(dtype); + } + + private static final DType UUID_STORAGE = new DType.FixedSizeList(DType.U8, 16, false); + + /// A GeoArrow coordinate (or box) struct of non-nullable F64 fields. + private static DType coords(String... names) { + return new DType.Struct(java.util.Arrays.stream(names).map(ColumnName::of).toList(), + java.util.Arrays.stream(names).map(n -> (DType) DType.F64).toList(), false); + } + + private static DType lists(DType element, int depth) { + DType out = element; + for (int i = 0; i < depth; i++) { + out = new DType.List(out, false); + } + return out; + } + + private static DType.Extension ext(ExtensionId id, DType storage, int... meta) { + byte[] bytes = new byte[meta.length]; + for (int i = 0; i < meta.length; i++) { + bytes[i] = (byte) meta[i]; + } + return new DType.Extension(id.id(), storage, MemorySegment.ofArray(bytes), storage.nullable()); } } } diff --git a/csv/src/main/java/io/github/dfa1/vortex/csv/CsvExporter.java b/csv/src/main/java/io/github/dfa1/vortex/csv/CsvExporter.java index 2ea0bad24..b882280c2 100644 --- a/csv/src/main/java/io/github/dfa1/vortex/csv/CsvExporter.java +++ b/csv/src/main/java/io/github/dfa1/vortex/csv/CsvExporter.java @@ -47,7 +47,7 @@ public static void exportCsv(Path vortexPath, Path csvPath) throws IOException { } public static void exportCsv(Path vortexPath, Path csvPath, ExportOptions options) throws IOException { - try (VortexReader reader = VortexReader.open(vortexPath); + try (VortexReader reader = VortexReader.open(vortexPath, options.readRegistry()); CsvWriter csvWriter = CsvWriter.builder() .fieldSeparator(options.delimiter()) .build(csvPath)) { @@ -70,7 +70,7 @@ public void close() { // do not close the caller-owned writer } }; - try (VortexReader reader = VortexReader.open(vortexPath); + try (VortexReader reader = VortexReader.open(vortexPath, options.readRegistry()); CsvWriter csvWriter = CsvWriter.builder() .fieldSeparator(options.delimiter()) .build(shielded)) { diff --git a/csv/src/main/java/io/github/dfa1/vortex/csv/ExportOptions.java b/csv/src/main/java/io/github/dfa1/vortex/csv/ExportOptions.java index 7b415c129..370fa02a4 100644 --- a/csv/src/main/java/io/github/dfa1/vortex/csv/ExportOptions.java +++ b/csv/src/main/java/io/github/dfa1/vortex/csv/ExportOptions.java @@ -1,6 +1,7 @@ package io.github.dfa1.vortex.csv; import io.github.dfa1.vortex.core.model.ColumnName; +import io.github.dfa1.vortex.reader.ReadRegistry; import java.util.List; @@ -10,12 +11,15 @@ public final class ExportOptions { private final boolean header; private final List columns; private final ProgressListener progressListener; + private final ReadRegistry readRegistry; - private ExportOptions(char delimiter, boolean header, List columns, ProgressListener progressListener) { + private ExportOptions(char delimiter, boolean header, List columns, ProgressListener progressListener, + ReadRegistry readRegistry) { this.delimiter = delimiter; this.header = header; this.columns = columns; this.progressListener = progressListener; + this.readRegistry = readRegistry; } /// The field delimiter. @@ -46,11 +50,19 @@ public ProgressListener progressListener() { return progressListener; } - /// Default options: comma delimiter, header row written, no projection, no progress listener. + /// The registry the Vortex file is opened with. + /// + /// @return the read registry + public ReadRegistry readRegistry() { + return readRegistry; + } + + /// Default options: comma delimiter, header row written, no projection, no progress listener, + /// the built-in [ReadRegistry]. /// /// @return the default options public static ExportOptions defaults() { - return new ExportOptions(',', true, List.of(), null); + return new ExportOptions(',', true, List.of(), null, ReadRegistry.loadAll()); } /// Override whether a header row with column names is written. @@ -58,7 +70,7 @@ public static ExportOptions defaults() { /// @param header `true` to write the header row /// @return a copy of this options with the header flag applied public ExportOptions withHeader(boolean header) { - return new ExportOptions(delimiter, header, columns, progressListener); + return new ExportOptions(delimiter, header, columns, progressListener, readRegistry); } /// Restrict output to specific columns (projection). Empty list = all columns. @@ -66,7 +78,7 @@ public ExportOptions withHeader(boolean header) { /// @param cols the column names to include, in output order /// @return a copy of this options with the projection applied public ExportOptions withColumns(List cols) { - return new ExportOptions(delimiter, header, List.copyOf(cols), progressListener); + return new ExportOptions(delimiter, header, List.copyOf(cols), progressListener, readRegistry); } /// Attach a progress callback invoked periodically with `(rowsDone, rowsTotal)`. @@ -74,7 +86,17 @@ public ExportOptions withColumns(List cols) { /// @param listener the progress callback /// @return a copy of this options with the listener attached public ExportOptions withProgressListener(ProgressListener listener) { - return new ExportOptions(delimiter, header, columns, listener); + return new ExportOptions(delimiter, header, columns, listener, readRegistry); + } + + /// Open the Vortex file with `registry`, e.g. one built with + /// [ReadRegistry.Builder#allowUnknown()] to export columns of an unregistered extension type as + /// their storage. + /// + /// @param registry the read registry + /// @return a copy of this options with the registry applied + public ExportOptions withReadRegistry(ReadRegistry registry) { + return new ExportOptions(delimiter, header, columns, progressListener, java.util.Objects.requireNonNull(registry)); } /// Whether a column projection has been applied via [#withColumns(List)]. diff --git a/docs/compatibility.md b/docs/compatibility.md index f3a4bbdc1..6c216fe32 100644 --- a/docs/compatibility.md +++ b/docs/compatibility.md @@ -278,13 +278,14 @@ metadata. The Rust catalog lives in each subdir below names a canonical extension id and its on-disk shape. Read-side extensions live in `io.github.dfa1.vortex.reader.extension` (write-side encoders -implement `ExtensionEncoder` in the writer module). Each spec extension is a singleton -implementing the `ExtensionDecoder` interface, with typed decode methods on the concrete -impl — grab the singleton directly: +implement `ExtensionEncoder` in the writer module). Both are registered on their registry, the +built-ins by `registerDefaults()`, as Rust registers an `ExtVTable` on the session. Each spec +extension is a singleton with typed decode methods on the concrete impl — grab the singleton +directly: ```java DType.Extension dtype = (DType.Extension) schema.field(ColumnName.of("birthdays")); -List values = DateExtensionDecoder.INSTANCE.decodeAll(chunk.column(ColumnName.of("birthdays"))); +List values = DateExtensionDecoder.INSTANCE.decodeAll(dtype, chunk.column(ColumnName.of("birthdays"))); ``` End-to-end round-trip — write a `List`, read it back: @@ -301,23 +302,45 @@ try (var iter = reader.scan(ScanOptions.all()); } ``` -`Chunk.as(name, Class)` hides the per-extension decode dispatch for the four -spec extensions (`LocalDate` ↔ `vortex.date`, `LocalTime` ↔ `vortex.time`, -`Instant` ↔ `vortex.timestamp`, `UUID` ↔ `vortex.uuid`); the dispatch is closed over the -spec set. Third-party extensions call their own impl's typed methods directly. - -`ExtensionId` is the enum of known spec ids (`VORTEX_DATE`, `VORTEX_TIME`, -`VORTEX_TIMESTAMP`, `VORTEX_UUID`). Unknown wire ids on `DType.Extension` -round-trip verbatim through the raw `String` field — the registry simply -returns `null` for them and callers can read the storage column directly. +`Chunk.as(name, Class)` decodes through the `ExtensionDecoder` registered for the column's id +(`LocalDate` ↔ `vortex.date`, `LocalTime` ↔ `vortex.time`, `Instant` ↔ `vortex.timestamp`, +`UUID` ↔ `vortex.uuid`, plus any registered third-party type). Both sides run the registered +`validate(DType.Extension)` — Rust's `validate_dtype` — on every extension dtype, at +`VortexWriter.create` and at open, and on the dtypes embedded in Variant metadata. The four +built-ins apply Rust's `deserialize_metadata` + `validate_dtype` rules +(`ExtensionId.WellKnown#validate`): `vortex.date` Days → `I32`, Milliseconds → `I64`; +`vortex.time` s/ms → `I32`, μs/ns → `I64`; `vortex.timestamp` metadata ≥ 3 bytes with a +UTF-8 zone and `I64` storage; `vortex.uuid` 0–1 metadata bytes (a known version) over +`FixedSizeList(non-null U8, 16)`. An extension dtype must carry its metadata field, empty when +there is none, as Rust requires; the writer always emits it. + +`ExtensionId` is a sealed `WellKnown`/`Custom` type like `EncodingId`: the spec ids +(`VORTEX_DATE`, `VORTEX_TIME`, `VORTEX_TIMESTAMP`, `VORTEX_UUID`) plus any third-party id. +The ids the vortex-jni session registers by default are `ExtensionId.WellKnown`: besides the +four above, `vortex.json` and the `vortex.st.*` geometries (`wkb`, `point`, `linestring`, +`polygon`, `multipoint`, `multilinestring`, `multipolygon`, `box`). A file with one of them opens +without a decoder and the column reads as its storage — JSON text as `Utf8`, WKB as `Binary`, +geometries as GeoArrow coordinate structs and lists — validated with Rust's storage rules (the +spatial CRS metadata is not decoded). Java has no JSON or geometry domain type, so `Chunk.as` +refuses them. + +As in Rust, opening a file with any other extension id that has no registered decoder fails +(`VortexException`; Rust: "No such DType extension ID"), unless +`ReadRegistry.Builder#allowUnknown()` is set — then the id round-trips verbatim through +`DType.Extension`'s raw `String` field and the column reads as its storage array, Rust's +`ForeignExtDType`. vortex-jni 0.88.0 exposes no `allow_unknown`, so it cannot read a file with a +custom extension id at all. A blank id, or one with a control character, fails the read with +`VortexException` (Rust accepts any string). | Extension id | Impl | Storage | Metadata | Round-trip | |---------------------|------------------------|-------------------------------------------------|-------------------------------------------|------------| -| `vortex.date` | `DateExtension` | Signed integer days since 1970-01-01 | none | ✅ | +| `vortex.date` | `DateExtension` | `I32` days (Arrow `Date32`) or `I64` milliseconds (`Date64`) since 1970-01-01 | 1 byte: `TimeUnit` (Days or Milliseconds) | ✅ (writes Days) | | `vortex.time` | `TimeExtension` | I32 (s/ms) or I64 (μs/ns) since midnight | 1 byte: `TimeUnit` | ✅ | | `vortex.timestamp` | `TimestampExtension` | I64 epoch count in the recorded `TimeUnit` | unit byte + u16 LE tz_len + UTF-8 tz | ✅ | | `vortex.uuid` | `UuidExtension` | `FixedSizeList(Primitive(U8), 16)` | none | ✅ | -| _custom ids_ | _none_ | _whatever the column declares_ | _opaque bytes_ | passthrough | +| `vortex.json` | _none (storage)_ | `Utf8` | empty | storage | +| `vortex.st.*` | _none (storage)_ | `Binary` (wkb) or GeoArrow coordinate `Struct`/`List` | spatial protobuf (not decoded) | storage | +| _custom ids_ | _registered decoder_ | _whatever the column declares_ | _opaque bytes_ | ✅ when registered; storage with `allowUnknown()` | `TimeUnit` (see [`extension/datetime/unit.rs`](https://github.com/vortex-data/vortex/blob/develop/vortex-array/src/extension/datetime/unit.rs)) encodes precision in the first metadata byte: diff --git a/docs/how-to.md b/docs/how-to.md index 87c56ce43..b25ca567f 100644 --- a/docs/how-to.md +++ b/docs/how-to.md @@ -868,10 +868,148 @@ try (VortexReader vf = VortexReader.open(Path.of("data.vortex"), readRegistry)) } ``` -`ExtensionId` (the `vortex.date`/`vortex.time`/`vortex.timestamp`/`vortex.uuid` family) is a -**closed enum**, unlike `EncodingId` — a custom `ExtensionEncoder`/`ExtensionDecoder` pair can only -re-implement one of those four spec-defined ids, not introduce a wire-new extension type from -outside the codebase. +--- + +## Add a custom extension type + +An extension type gives a storage dtype a logical meaning, as `vortex.timestamp` does for `I64`. +`ExtensionId` is a sealed `WellKnown`/`Custom` type like `EncodingId`, so a third party names its +own id, writes domain values through an `ExtensionEncoder` and reads them back through an +`ExtensionDecoder` — the Java side of Rust's `ExtVTable`. Both sides may override `validate` (Rust's +`validate_dtype`), run on every dtype with that id when the writer is created and when a file is +opened. + +The example is a money column with one currency per column: amounts are `I64` minor units +(cents for EUR), and the ISO 4217 code is the extension metadata, as a timestamp keeps its zone. +Never store money as a float. + +**Write side:** + +```java +import io.github.dfa1.vortex.core.error.VortexException; +import io.github.dfa1.vortex.core.model.ExtensionId; +import io.github.dfa1.vortex.writer.ExtensionEncoder; + +import java.lang.foreign.ValueLayout; +import java.math.BigDecimal; +import java.nio.charset.StandardCharsets; +import java.util.Collection; +import java.util.Currency; + +final class MoneyEncoder implements ExtensionEncoder { + static final ExtensionId ID = ExtensionId.parse("acme.money"); + + static DType.Extension dtype(Currency currency) { + byte[] code = currency.getCurrencyCode().getBytes(StandardCharsets.US_ASCII); + return new DType.Extension(ID.id(), DType.I64, MemorySegment.ofArray(code), false); + } + + @Override + public ExtensionId extensionId() { + return ID; + } + + @Override + public DType.Extension dtype(boolean nullable) { + return dtype(Currency.getInstance("EUR")); + } + + @Override + public void validate(DType.Extension dtype) { + if (!dtype.storageDType().equals(DType.I64)) { + throw new VortexException("acme.money storage must be non-nullable I64"); + } + currency(dtype); // rejects a code that is not ISO 4217 + } + + @Override + public Object encodeAll(DType.Extension dtype, Collection values) { + int digits = currency(dtype).getDefaultFractionDigits(); + long[] minorUnits = new long[values.size()]; + int i = 0; + for (Object value : values) { + // throws on an amount finer than the currency's minor unit (e.g. 0.001 EUR) + minorUnits[i++] = ((BigDecimal) value).movePointRight(digits).longValueExact(); + } + return minorUnits; + } + + static Currency currency(DType.Extension dtype) { + return Currency.getInstance(new String(dtype.metadata().toArray(ValueLayout.JAVA_BYTE), + StandardCharsets.US_ASCII)); + } +} + +DType.Struct schema = new DType.Struct(List.of(ColumnName.of("price")), + List.of(MoneyEncoder.dtype(Currency.getInstance("EUR"))), false); +WriteRegistry registry = WriteRegistry.builder().registerDefaults().register(new MoneyEncoder()).build(); + +try (var ch = FileChannel.open(Path.of("prices.vortex"), StandardOpenOption.CREATE, StandardOpenOption.WRITE); + var writer = VortexWriter.create(ch, schema, WriteOptions.defaults(), registry)) { + writer.writeChunk(chunk -> chunk.put(ColumnName.of("price"), + List.of(new BigDecimal("9.99"), new BigDecimal("120.50")))); +} +``` + +**Read side:** + +```java +import io.github.dfa1.vortex.reader.Chunk; +import io.github.dfa1.vortex.reader.array.LongArray; +import io.github.dfa1.vortex.reader.extension.ExtensionDecoder; + +import java.util.ArrayList; + +final class MoneyDecoder implements ExtensionDecoder { + static final ExtensionId ID = ExtensionId.parse("acme.money"); + + @Override + public ExtensionId extensionId() { + return ID; + } + + @Override + public DType.Extension dtype(boolean nullable) { + byte[] code = "EUR".getBytes(StandardCharsets.US_ASCII); + return new DType.Extension(ID.id(), DType.I64, MemorySegment.ofArray(code), false); + } + + @Override + public Class domainType() { + return BigDecimal.class; + } + + @Override + public List decodeAll(DType.Extension dtype, Array storage) { + Currency currency = Currency.getInstance(new String(dtype.metadata().toArray(ValueLayout.JAVA_BYTE), + StandardCharsets.US_ASCII)); + LongArray minorUnits = (LongArray) storage; + List out = new ArrayList<>(); + for (long i = 0; i < minorUnits.length(); i++) { + out.add(BigDecimal.valueOf(minorUnits.getLong(i), currency.getDefaultFractionDigits())); + } + return out; + } +} + +ReadRegistry readRegistry = ReadRegistry.builder().registerDefaults().register(new MoneyDecoder()).build(); + +try (VortexReader vf = VortexReader.open(Path.of("prices.vortex"), readRegistry); + var iter = vf.scan(ScanOptions.all())) { + while (iter.hasNext()) { + try (Chunk chunk = iter.next()) { + List prices = chunk.as(ColumnName.of("price"), BigDecimal.class); // 9.99, 120.50 + } + } +} +``` + +`validate` checks the dtype, not the values: per-value rules belong in `encodeAll`. Filters and zone +maps see the stored `I64`. A reader without the decoder — including the Rust reference — refuses +the file, as Rust does for an extension id its session has no vtable for; with +`ReadRegistry.Builder#allowUnknown()` the column reads as its storage array. Writing a `Collection` +for an extension column needs a registered encoder; the storage array itself (`long[]` here) can +always be written. Auto-routing a `Collection` applies to top-level columns only. --- diff --git a/docs/reference.md b/docs/reference.md index 83f33073f..849bac9b0 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -67,6 +67,7 @@ shortcut returning a nullable copy), `withNullable(boolean)`, `DType.Struct.fiel | Type | Shape | Notes | |------|-------|-------| | `EncodingId` | `sealed interface` — `WellKnown` enum + `Custom` record | Array-encoding identity; total `parse(String)` over non-blank ids; constants re-exported (`EncodingId.VORTEX_PRIMITIVE`, …) | +| `ExtensionId` | `sealed interface` — `WellKnown` enum + `Custom` record | Extension-type identity (`vortex.date`/`time`/`timestamp`/`uuid` well-known); same `parse(String)` and `Custom` policy as `EncodingId`. `WellKnown#validate(DType.Extension)` applies Rust's metadata and storage checks for the four spec types. `DType.Extension` still carries the raw string; the file parser rejects ids `parse` would refuse | | `LayoutId` | `sealed interface` — `WellKnown` enum + `Custom` record | Layout identity (separate namespace from encodings; `vortex.flat` is layout-only); both zoned aliases `vortex.zoned`/`vortex.stats` | | `ColumnName` | `record ColumnName(String value)` | Validated column name: non-blank, no control characters. `ColumnName.violation(String)` is the policy chokepoint shared by builder, writer, and file parser | | `MemorySize` | `record MemorySize(long bytes)` | Validated non-negative byte count; `ofKiB`/`ofMiB`/`ofGiB` factories, `toGiB()` for display. Rejects negative values (`IllegalArgumentException`) — a programmer-error guard, not a `VortexException`-worthy untrusted-input check | @@ -267,7 +268,7 @@ never desync. | `rowCount()` | Rows in this chunk | | `columns()` | `SequencedMap`, schema order, unmodifiable | | ` column(ColumnName name)` | Typed column lookup; throws `VortexException` if absent | -| `as(ColumnName name, Class domainType)` | Extension column → typed `List` | +| `as(ColumnName name, Class domainType)` | Extension column → typed `List`, via the registered `ExtensionDecoder` | | `isClosed()` | Whether `close()` has run | | `close()` | Releases the chunk's arena. Idempotent. | @@ -282,9 +283,10 @@ Immutable after construction. Build via `ReadRegistry.builder()` or the static c | Method | Notes | |-----------------------------|------------------------------------------------------------------------------| | `static builder()` | Returns a fresh `Builder` | -| `static loadAll()` | Immutable registry populated with all built-in decoders | +| `static loadAll()` | Immutable registry populated with all built-in decoders and extension decoders | | `static empty()` | Immutable empty registry (strict mode) | | `hasDecoder(EncodingId)` | Lookup | +| `lookup(ExtensionId)` | Registered `ExtensionDecoder` for the id, or `null` | | `isAllowUnknown()` | Predicate | ### `ReadRegistry.Builder` @@ -292,14 +294,16 @@ Immutable after construction. Build via `ReadRegistry.builder()` or the static c | Method | Notes | |------------------------------|------------------------------------------------------------------------------------------| | `register(EncodingDecoder)` | Add a custom encoding decoder; throws if already registered | -| `registerDefaults()` | Add every built-in `EncodingDecoder` | -| `allowUnknown()` | Switch to passthrough mode — unknown nodes (and their children) decode as `UnknownArray` | +| `register(ExtensionDecoder)` | Add an extension decoder, keyed by `ExtensionDecoder#extensionId()`; throws on duplicate | +| `registerDefaults()` | Add every built-in `EncodingDecoder` and `ExtensionDecoder` | +| `allowUnknown()` | Switch to passthrough mode — unknown nodes (and their children) decode as `UnknownArray`, and an extension id with no decoder reads as its storage array | | `build()` | Produce the immutable `ReadRegistry` | -Register custom encoding decoders programmatically via `register(EncodingDecoder)` — there is no -`ServiceLoader` discovery. Extension decoders -(`io.github.dfa1.vortex.reader.extension.ExtensionDecoder`) are not registry-managed: the built-in -implementations are singletons invoked directly by their `ExtensionId`. +Register custom encoding and extension decoders programmatically via `register(...)` — there is no +`ServiceLoader` discovery. `Chunk.as(name, Class)` decodes an extension column through the +`ExtensionDecoder` registered for its id. Opening a file runs each registered decoder's +`validate(DType.Extension)` on its dtypes, and refuses an extension id with no decoder unless +`allowUnknown()` is set — Rust's behavior for an id its session has no vtable for. ### `WriteRegistry` (`io.github.dfa1.vortex.writer.WriteRegistry`) @@ -327,12 +331,12 @@ convenience factories. `EncodingId` is a sealed `WellKnown`/`Custom` type (see [Identity types](#identity-types-iogithubdfa1vortexcoremodel)), so `register(EncodingEncoder)` accepts a genuinely third-party encoding — `encodingId()` can return `new EncodingId.Custom("acme.myencoding")`, and `VortexWriter.create(channel, schema, options, -registry)` will pick it whenever its `accepts(DType)` matches. `ExtensionId`, by contrast, is a -**closed enum** with only the four spec-defined constants (`VORTEX_DATE`/`VORTEX_TIME`/ -`VORTEX_TIMESTAMP`/`VORTEX_UUID`) — `register(ExtensionEncoder)` can only re-implement one of -those four, not introduce a wire-new extension type; adding a genuinely new extension type means -adding an `ExtensionId` constant in `core` itself (see [CLAUDE.md § Adding an extension -type](../CLAUDE.md#adding-an-extension-type)). +registry)` will pick it whenever its `accepts(DType)` matches. `ExtensionId` has the same shape, so +`register(ExtensionEncoder)` accepts a third-party extension type: `VortexWriter.create(channel, +schema, options, registry)` runs its `validate(DType.Extension)` on the schema and routes a +`Collection` for a top-level column with that id through `encodeAll`. A `Collection` for an +extension id with no registered encoder fails with `VortexException`; the storage array itself can +always be written. See [how-to.md § Add a custom extension type](how-to.md#add-a-custom-extension-type). --- @@ -650,6 +654,11 @@ even if the command fails, and the exit status is unchanged. The clock starts wh JVM's start-up (roughly half of a short command such as `count`: 33 ms of work in a 60 ms process on a 127 MB file). For the interactive `tui` and `view` it includes the time you spend in them. +Any subcommand also takes `--allow-unknown`, anywhere on the command line: a file with an extension type the CLI has +no decoder for opens, and those columns read as their storage (`ReadRegistry.Builder#allowUnknown()`). Without it +such a file is refused, as the Rust reference refuses it. `vortex.json` and the `vortex.st.*` geometry types need no +flag: they are well-known and always read as their storage. + ### `filter` expression syntax ``` diff --git a/inspector/src/main/java/io/github/dfa1/vortex/inspect/HtmlReport.java b/inspector/src/main/java/io/github/dfa1/vortex/inspect/HtmlReport.java index 2af3ea656..283f15a60 100644 --- a/inspector/src/main/java/io/github/dfa1/vortex/inspect/HtmlReport.java +++ b/inspector/src/main/java/io/github/dfa1/vortex/inspect/HtmlReport.java @@ -6,10 +6,10 @@ import io.github.dfa1.vortex.core.model.PType; import io.github.dfa1.vortex.reader.ArrayStats; import io.github.dfa1.vortex.reader.SegmentSpec; +import io.github.dfa1.vortex.reader.extension.DateExtensionDecoder; import io.github.dfa1.vortex.reader.extension.ExtensionStorage; import io.github.dfa1.vortex.reader.layout.Layout; -import java.time.LocalDate; import java.util.ArrayList; import java.util.Arrays; import java.util.Comparator; @@ -583,10 +583,11 @@ private static Object readable(Object raw, DType type) { return raw; } try { - return switch (ExtensionId.parse(ext.extensionId()).orElse(null)) { - case VORTEX_TIMESTAMP -> ExtensionStorage.instantFromRaw(epoch, ExtensionStorage.readUnit(ext)); - case VORTEX_DATE -> LocalDate.ofEpochDay(epoch); - case null, default -> raw; + return switch (ExtensionId.parse(ext.extensionId())) { + case ExtensionId.WellKnown.VORTEX_TIMESTAMP -> + ExtensionStorage.instantFromRaw(epoch, ExtensionStorage.readUnit(ext)); + case ExtensionId.WellKnown.VORTEX_DATE -> DateExtensionDecoder.INSTANCE.date(ext, epoch); + default -> raw; }; } catch (RuntimeException e) { return raw; diff --git a/integration/src/test/java/io/github/dfa1/vortex/integration/CustomExtensionIntegrationTest.java b/integration/src/test/java/io/github/dfa1/vortex/integration/CustomExtensionIntegrationTest.java new file mode 100644 index 000000000..1cf5deb9a --- /dev/null +++ b/integration/src/test/java/io/github/dfa1/vortex/integration/CustomExtensionIntegrationTest.java @@ -0,0 +1,492 @@ +package io.github.dfa1.vortex.integration; + +import dev.vortex.api.DataSource; +import dev.vortex.api.Session; +import dev.vortex.jni.NativeLoader; +import io.github.dfa1.vortex.core.error.VortexException; +import io.github.dfa1.vortex.core.model.ColumnName; +import io.github.dfa1.vortex.core.model.DType; +import io.github.dfa1.vortex.core.model.ExtensionId; +import io.github.dfa1.vortex.core.model.PType; +import io.github.dfa1.vortex.reader.Chunk; +import io.github.dfa1.vortex.reader.ReadRegistry; +import io.github.dfa1.vortex.reader.RowFilter; +import io.github.dfa1.vortex.reader.ScanOptions; +import io.github.dfa1.vortex.reader.VortexReader; +import io.github.dfa1.vortex.reader.array.Array; +import io.github.dfa1.vortex.reader.array.FixedSizeListArray; +import io.github.dfa1.vortex.reader.array.FloatArray; +import io.github.dfa1.vortex.reader.array.LongArray; +import io.github.dfa1.vortex.reader.array.MaskedArray; +import io.github.dfa1.vortex.reader.extension.ExtensionDecoder; +import io.github.dfa1.vortex.writer.ExtensionEncoder; +import io.github.dfa1.vortex.writer.VortexWriter; +import io.github.dfa1.vortex.writer.WriteOptions; +import io.github.dfa1.vortex.writer.WriteRegistry; +import io.github.dfa1.vortex.writer.encode.FixedSizeListData; +import io.github.dfa1.vortex.writer.encode.NullableData; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.IOException; +import java.lang.foreign.MemorySegment; +import java.lang.foreign.ValueLayout; +import java.nio.channels.FileChannel; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.SequencedMap; +import java.util.stream.LongStream; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/// Acceptance test for a third-party extension type (#530), end to end through the public API: +/// a `example.vector` type is defined outside the library, registered on both registries, written +/// from domain values and read back as domain values — the Java counterpart of Rust's +/// "adding a new logical type" (an `ExtVTable` registered on the session). +/// +/// The type is a fixed-dimension `float[]` embedding over `FixedSizeList` storage, with +/// the dimension also recorded in the metadata (u32 LE), so the test covers non-primitive storage, +/// metadata round-trip, nulls and the `validate` hook on both sides. Reads go through a full scan, a +/// filter on a sibling column and a limit: a full scan never prunes or slices, and a sliced +/// fixed-size list must keep each vector's floats with its row. +class CustomExtensionIntegrationTest { + + static { + NativeLoader.loadJni(); + } + + private static final Session SESSION = Session.create(); + private static final ExtensionId VECTOR = ExtensionId.parse("example.vector"); + private static final ColumnName ID = ColumnName.of("id"); + private static final ColumnName EMBEDDING = ColumnName.of("embedding"); + private static final int DIM = 3; + private static final int CHUNKS = 3; + private static final int CHUNK_ROWS = 10_000; + private static final int ROWS = CHUNKS * CHUNK_ROWS; + + private static final WriteRegistry WRITE_REGISTRY = WriteRegistry.builder() + .registerDefaults() + .register(new VectorEncoder()) + .build(); + private static final ReadRegistry READ_REGISTRY = ReadRegistry.builder() + .registerDefaults() + .register(new VectorDecoder()) + .build(); + + @Nested + class RoundTrip { + + @Test + void fullScan_decodesEveryVectorAndNull(@TempDir Path tmp) throws IOException { + // Given — three batches of nullable vectors written from float[] values + Path file = tmp.resolve("vectors.vortex"); + writeVectors(file, vectorDtype(DIM, true)); + + // When + SequencedMap result = readVectors(file, ScanOptions.all()); + + // Then + assertThat(result.keySet()).containsExactlyElementsOf(ids(0, ROWS)); + assertEachVectorMatchesItsId(result); + } + + @Test + void schema_keepsIdStorageAndMetadata(@TempDir Path tmp) throws IOException { + // Given + Path file = tmp.resolve("vectors.vortex"); + DType.Extension written = vectorDtype(DIM, true); + writeVectors(file, written); + + try (VortexReader reader = VortexReader.open(file, READ_REGISTRY)) { + // When — the metadata is a view of the mapped file, so compare before closing + DType result = ((DType.Struct) reader.dtype()).fieldTypes().get(1); + + // Then — Extension.equals compares the metadata bytes, so the dimension survived too + assertThat(result).isEqualTo(written); + } + } + + @Test + void filteredScan_keepsEachVectorWithItsRow(@TempDir Path tmp) throws IOException { + // Given — a zone-map filter on the sibling id column. It skips only whole chunks that + // cannot match, so the result is a superset of the matching rows; every vector must + // still sit next to its own id. + Path file = tmp.resolve("vectors.vortex"); + writeVectors(file, vectorDtype(DIM, true)); + long threshold = CHUNK_ROWS + CHUNK_ROWS / 2; + + // When + SequencedMap result = + readVectors(file, ScanOptions.all().withFilter(RowFilter.gte(ID, threshold))); + + // Then + assertThat(result.keySet()).containsAll(ids(threshold, ROWS)); + assertEachVectorMatchesItsId(result); + } + + @Test + void limitedScan_returnsTheFirstRows(@TempDir Path tmp) throws IOException { + // Given + Path file = tmp.resolve("vectors.vortex"); + writeVectors(file, vectorDtype(DIM, true)); + int limit = CHUNK_ROWS + 7; + + // When + SequencedMap result = readVectors(file, ScanOptions.limit(limit)); + + // Then + assertThat(result.keySet()).containsExactlyElementsOf(ids(0, limit)); + assertEachVectorMatchesItsId(result); + } + + @Test + void storageArray_withoutEncoder_isWrittenAsIs(@TempDir Path tmp) throws IOException { + // Given — no encoder registered: the caller passes the storage shape itself, as for any + // extension id the writer does not know + Path file = tmp.resolve("raw.vortex"); + DType.Struct schema = schema(vectorDtype(2, false)); + float[] flat = {1f, 2f, 3f, 4f}; + try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE); + var sut = VortexWriter.create(ch, schema, WriteOptions.defaults())) { + // When + sut.writeChunk(Map.of(ID, new long[]{0L, 1L}, EMBEDDING, new FixedSizeListData(flat, 2))); + } + + // Then + assertThat(readVectors(file, ScanOptions.all()).values()) + .containsExactly(new float[]{1f, 2f}, new float[]{3f, 4f}); + } + } + + @Nested + class UnregisteredOnRead { + + @Test + void open_withoutDecoder_throws(@TempDir Path tmp) throws IOException { + // Given — a reader with only the built-ins. Rust refuses an extension id its session + // has no vtable for ("No such DType extension ID") unless unknown ids are allowed. + Path file = tmp.resolve("vectors.vortex"); + writeVectors(file, vectorDtype(DIM, false)); + ReadRegistry builtInsOnly = ReadRegistry.loadAll(); + + // When / Then + assertThatThrownBy(() -> VortexReader.open(file, builtInsOnly)) + .isInstanceOf(VortexException.class) + .hasMessageContaining("no extension decoder registered for example.vector"); + } + + @Test + void allowUnknown_readsColumnAsStorageArray(@TempDir Path tmp) throws IOException { + // Given — with unknown ids allowed the column reads as its fixed-size-list storage, + // Rust's ForeignExtDType + Path file = tmp.resolve("vectors.vortex"); + writeVectors(file, vectorDtype(DIM, false)); + ReadRegistry allowUnknown = ReadRegistry.builder().registerDefaults().allowUnknown().build(); + + // When + float firstElementOfRow1; + try (VortexReader reader = VortexReader.open(file, allowUnknown); + var iter = reader.scan(ScanOptions.limit(2)); + Chunk chunk = iter.next()) { + FixedSizeListArray storage = (FixedSizeListArray) chunk.column(EMBEDDING); + firstElementOfRow1 = ((FloatArray) storage.elements()).getFloat(DIM); + } + + // Then + assertThat(firstElementOfRow1).isEqualTo(vector(1)[0]); + } + + @Test + void as_withoutDecoder_throws(@TempDir Path tmp) throws IOException { + // Given + Path file = tmp.resolve("vectors.vortex"); + writeVectors(file, vectorDtype(DIM, false)); + ReadRegistry allowUnknown = ReadRegistry.builder().registerDefaults().allowUnknown().build(); + + try (VortexReader reader = VortexReader.open(file, allowUnknown); + var iter = reader.scan(ScanOptions.limit(2)); + Chunk chunk = iter.next()) { + // When / Then + assertThatThrownBy(() -> chunk.as(EMBEDDING, float[].class)) + .isInstanceOf(VortexException.class) + .hasMessageContaining("no extension decoder registered for example.vector"); + } + } + + @Test + void as_wrongDomainType_throws(@TempDir Path tmp) throws IOException { + // Given + Path file = tmp.resolve("vectors.vortex"); + writeVectors(file, vectorDtype(DIM, false)); + + try (VortexReader reader = VortexReader.open(file, READ_REGISTRY); + var iter = reader.scan(ScanOptions.limit(2)); + Chunk chunk = iter.next()) { + // When / Then + assertThatThrownBy(() -> chunk.as(EMBEDDING, double[].class)) + .isInstanceOf(VortexException.class) + .hasMessageContaining("decodes to float[], not double[]"); + } + } + + @Test + void vortexJni_withoutVtable_refusesTheFile(@TempDir Path tmp) throws IOException { + // Given — the Rust reference, which has no example.vector vtable and whose JNI session + // exposes no allow_unknown: the behavior the strict default above mirrors + Path file = tmp.resolve("vectors.vortex"); + writeVectors(file, vectorDtype(DIM, false)); + String uri = file.toAbsolutePath().toUri().toString(); + + // When / Then + assertThatThrownBy(() -> DataSource.open(SESSION, uri).scan(dev.vortex.api.ScanOptions.of())) + .hasMessageContaining("No such DType extension ID: example.vector"); + } + } + + @Nested + class Validation { + + @Test + void create_rejectsDtypeTheEncoderRefuses(@TempDir Path tmp) throws IOException { + // Given — metadata says 4 dimensions, storage holds 3 + DType.Extension invalid = new DType.Extension(VECTOR.id(), + new DType.FixedSizeList(DType.F32, DIM, false), dimensionMetadata(4), false); + Path file = tmp.resolve("invalid.vortex"); + + try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE)) { + // When / Then + assertThatThrownBy(() -> VortexWriter.create(ch, schema(invalid), WriteOptions.defaults(), WRITE_REGISTRY)) + .isInstanceOf(VortexException.class) + .hasMessageContaining("example.vector metadata dimension 4 != storage size 3"); + } + } + + @Test + void open_rejectsDtypeTheDecoderRefuses(@TempDir Path tmp) throws IOException { + // Given — a file from a writer that had no encoder, so nothing validated the dtype + DType.Extension invalid = new DType.Extension(VECTOR.id(), + new DType.FixedSizeList(DType.F32, 2, false), dimensionMetadata(5), false); + Path file = tmp.resolve("invalid.vortex"); + try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE); + var writer = VortexWriter.create(ch, schema(invalid), WriteOptions.defaults())) { + writer.writeChunk(Map.of(ID, new long[]{0L}, EMBEDDING, new FixedSizeListData(new float[]{1f, 2f}, 1))); + } + + // When / Then — Rust validates an extension dtype when it deserializes it + assertThatThrownBy(() -> VortexReader.open(file, READ_REGISTRY)) + .isInstanceOf(VortexException.class) + .hasMessageContaining("example.vector metadata dimension 5 != storage size 2"); + } + + @Test + void writeChunk_nullInNonNullableColumn_throws(@TempDir Path tmp) throws IOException { + // Given + Path file = tmp.resolve("nonnull.vortex"); + try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE); + var sut = VortexWriter.create(ch, schema(vectorDtype(DIM, false)), WriteOptions.defaults(), WRITE_REGISTRY)) { + // When / Then + assertThatThrownBy(() -> sut.writeChunk(Map.of(ID, new long[]{0L}, + EMBEDDING, Arrays.asList((float[]) null)))) + .isInstanceOf(VortexException.class) + .hasMessageContaining("null element in non-nullable example.vector column"); + } + } + } + + // ── the custom extension type ──────────────────────────────────────────── + + /// The `example.vector` dtype: `FixedSizeList` storage, dimension also in metadata. + static DType.Extension vectorDtype(int dim, boolean nullable) { + return new DType.Extension(VECTOR.id(), new DType.FixedSizeList(DType.F32, dim, nullable), + dimensionMetadata(dim), nullable); + } + + static MemorySegment dimensionMetadata(int dim) { + MemorySegment metadata = MemorySegment.ofArray(new byte[4]); + metadata.set(ValueLayout.JAVA_INT_UNALIGNED.withOrder(java.nio.ByteOrder.LITTLE_ENDIAN), 0, dim); + return metadata; + } + + /// Shared by encoder and decoder, as one Rust `ExtVTable::validate_dtype` serves both directions. + static void validateVector(DType.Extension dtype) { + if (!(dtype.storageDType() instanceof DType.FixedSizeList list) + || !(list.elementType() instanceof DType.Primitive p && p.ptype() == PType.F32 && !p.nullable())) { + throw new VortexException("example.vector storage must be FixedSizeList, got " + + dtype.storageDType()); + } + MemorySegment metadata = dtype.metadata(); + if (metadata == null || metadata.byteSize() != 4) { + throw new VortexException("example.vector metadata must be a u32 dimension"); + } + int dim = metadata.get(ValueLayout.JAVA_INT_UNALIGNED.withOrder(java.nio.ByteOrder.LITTLE_ENDIAN), 0); + if (dim != list.fixedSize()) { + throw new VortexException("example.vector metadata dimension " + dim + + " != storage size " + list.fixedSize()); + } + } + + static final class VectorEncoder implements ExtensionEncoder { + + @Override + public ExtensionId extensionId() { + return VECTOR; + } + + @Override + public DType.Extension dtype(boolean nullable) { + return vectorDtype(DIM, nullable); + } + + @Override + public void validate(DType.Extension dtype) { + validateVector(dtype); + } + + @Override + public Object encodeAll(DType.Extension dtype, Collection values) { + int dim = ((DType.FixedSizeList) dtype.storageDType()).fixedSize(); + int n = values.size(); + float[] flat = new float[n * dim]; + boolean[] validity = new boolean[n]; + boolean anyNull = false; + int row = 0; + for (Object value : values) { + if (value == null) { + anyNull = true; + } else { + float[] vector = (float[]) value; + if (vector.length != dim) { + throw new VortexException("row " + row + " has " + vector.length + " dimensions, not " + dim); + } + System.arraycopy(vector, 0, flat, row * dim, dim); + validity[row] = true; + } + row++; + } + FixedSizeListData storage = new FixedSizeListData(flat, n); + if (!anyNull) { + return storage; + } + if (!dtype.nullable()) { + throw new VortexException("null element in non-nullable example.vector column"); + } + return new NullableData(storage, validity); + } + } + + static final class VectorDecoder implements ExtensionDecoder { + + @Override + public ExtensionId extensionId() { + return VECTOR; + } + + @Override + public DType.Extension dtype(boolean nullable) { + return vectorDtype(DIM, nullable); + } + + @Override + public Class domainType() { + return float[].class; + } + + @Override + public void validate(DType.Extension dtype) { + validateVector(dtype); + } + + @Override + public List decodeAll(DType.Extension dtype, Array storage) { + MaskedArray.Unwrapped unwrapped = MaskedArray.unwrap(storage); + FixedSizeListArray lists = (FixedSizeListArray) unwrapped.inner(); + FloatArray elements = (FloatArray) MaskedArray.innerOrSelf(lists.elements()); + int dim = lists.fixedSize(); + int n = Math.toIntExact(lists.length()); + List out = new ArrayList<>(n); + for (int row = 0; row < n; row++) { + if (unwrapped.validity() != null && !unwrapped.validity().getBoolean(row)) { + out.add(null); + continue; + } + float[] vector = new float[dim]; + for (int k = 0; k < dim; k++) { + vector[k] = elements.getFloat((long) row * dim + k); + } + out.add(vector); + } + return out; + } + } + + // ── helpers ────────────────────────────────────────────────────────────── + + private static DType.Struct schema(DType.Extension vectorDtype) { + return new DType.Struct(List.of(ID, EMBEDDING), List.of(DType.I64, vectorDtype), false); + } + + /// Every seventh row is null in a nullable column; values differ per row and per dimension, with + /// negatives, so a slice off by one row or one element never reads back equal. + private static float[] vector(long row) { + if (row % 7 == 3) { + return null; + } + float[] vector = new float[DIM]; + for (int k = 0; k < DIM; k++) { + vector[k] = (k % 2 == 0 ? 1 : -1) * (row + k * 0.25f); + } + return vector; + } + + private void writeVectors(Path file, DType.Extension vectorDtype) throws IOException { + try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE); + var writer = VortexWriter.create(ch, schema(vectorDtype), WriteOptions.defaults(), WRITE_REGISTRY)) { + for (int chunk = 0; chunk < CHUNKS; chunk++) { + long[] ids = new long[CHUNK_ROWS]; + List vectors = new ArrayList<>(CHUNK_ROWS); + for (int i = 0; i < CHUNK_ROWS; i++) { + long row = (long) chunk * CHUNK_ROWS + i; + ids[i] = row; + float[] vector = vector(row); + vectors.add(vector == null && !vectorDtype.nullable() ? new float[DIM] : vector); + } + writer.writeChunk(Map.of(ID, ids, EMBEDDING, vectors)); + } + } + } + + /// Reads `id → vector` in scan order; the id column is read alongside so a scan that misaligns + /// the vector column with its sibling is caught. + private SequencedMap readVectors(Path file, ScanOptions options) throws IOException { + SequencedMap out = new LinkedHashMap<>(); + try (VortexReader reader = VortexReader.open(file, READ_REGISTRY); + var iter = reader.scan(options)) { + while (iter.hasNext()) { + try (Chunk chunk = iter.next()) { + List vectors = chunk.as(EMBEDDING, float[].class); + LongArray ids = (LongArray) chunk.column(ID); + for (int i = 0; i < vectors.size(); i++) { + out.put(ids.getLong(i), vectors.get(i)); + } + } + } + } + return out; + } + + private static List ids(long fromInclusive, long toExclusive) { + return LongStream.range(fromInclusive, toExclusive).boxed().toList(); + } + + private static void assertEachVectorMatchesItsId(Map rows) { + rows.forEach((id, vector) -> assertThat(vector).as("row %d", id).isEqualTo(vector(id))); + } +} diff --git a/integration/src/test/java/io/github/dfa1/vortex/integration/DateInteropIntegrationTest.java b/integration/src/test/java/io/github/dfa1/vortex/integration/DateInteropIntegrationTest.java new file mode 100644 index 000000000..5d22af281 --- /dev/null +++ b/integration/src/test/java/io/github/dfa1/vortex/integration/DateInteropIntegrationTest.java @@ -0,0 +1,178 @@ +package io.github.dfa1.vortex.integration; + +import dev.vortex.api.Session; +import dev.vortex.api.VortexWriter; +import dev.vortex.arrow.ArrowAllocation; +import dev.vortex.jni.NativeLoader; +import io.github.dfa1.vortex.core.model.ColumnName; +import io.github.dfa1.vortex.core.model.DType; +import io.github.dfa1.vortex.core.model.TimeUnit; +import io.github.dfa1.vortex.reader.Chunk; +import io.github.dfa1.vortex.reader.ReadRegistry; +import io.github.dfa1.vortex.reader.RowFilter; +import io.github.dfa1.vortex.reader.ScanOptions; +import io.github.dfa1.vortex.reader.VortexReader; +import io.github.dfa1.vortex.reader.array.LongArray; +import org.apache.arrow.c.ArrowArray; +import org.apache.arrow.c.ArrowSchema; +import org.apache.arrow.c.Data; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.DateDayVector; +import org.apache.arrow.vector.DateMilliVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.DateUnit; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.EnumSource; + +import java.io.IOException; +import java.lang.foreign.ValueLayout; +import java.nio.file.Path; +import java.time.LocalDate; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.SequencedMap; + +import static org.assertj.core.api.Assertions.assertThat; + +/// `vortex.date` written by vortex-jni, read by Java. Rust stores an Arrow `Date32` as +/// `vortex.date(Days)` over `I32` and a `Date64` as `vortex.date(Milliseconds)` over `I64`; the +/// Java decoder used to read every storage value as days, so a `Date64` column came back as dates +/// tens of thousands of years away. Both units are read through a full scan and a filter on a +/// sibling column, with nulls and dates before the epoch. +class DateInteropIntegrationTest { + + private static final Session SESSION = Session.create(); + private static final BufferAllocator ALLOCATOR = ArrowAllocation.rootAllocator(); + private static final ColumnName ID = ColumnName.of("id"); + private static final ColumnName DATE = ColumnName.of("d"); + private static final int CHUNKS = 3; + private static final int CHUNK_ROWS = 1_000; + private static final int ROWS = CHUNKS * CHUNK_ROWS; + private static final long MILLIS_PER_DAY = 86_400_000L; + + static { + NativeLoader.loadJni(); + } + + @ParameterizedTest + @EnumSource(DateUnit.class) + void rustWrites_javaReads_fullScan(DateUnit unit, @TempDir Path tmp) throws IOException { + // Given + Path file = tmp.resolve("dates.vortex"); + writeWithJni(file, unit); + + // When + SequencedMap result = readJava(file, ScanOptions.all()); + + // Then + assertThat(result).hasSize(ROWS); + result.forEach((id, date) -> assertThat(date).as("row %d", id).isEqualTo(expected(id))); + } + + @ParameterizedTest + @EnumSource(DateUnit.class) + void rustWrites_javaReads_filtered(DateUnit unit, @TempDir Path tmp) throws IOException { + // Given — a zone-map filter keeps a superset of the matching rows; each date must still + // sit next to its own id + Path file = tmp.resolve("dates.vortex"); + writeWithJni(file, unit); + long threshold = CHUNK_ROWS + CHUNK_ROWS / 2; + + // When + SequencedMap result = readJava(file, ScanOptions.all().withFilter(RowFilter.gte(ID, threshold))); + + // Then + assertThat(result.keySet()).contains(threshold, (long) ROWS - 1); + result.forEach((id, date) -> assertThat(date).as("row %d", id).isEqualTo(expected(id))); + } + + @ParameterizedTest + @EnumSource(DateUnit.class) + void rustWrites_dtypeRecordsTheUnit(DateUnit unit, @TempDir Path tmp) throws IOException { + // Given + Path file = tmp.resolve("dates.vortex"); + writeWithJni(file, unit); + + try (VortexReader reader = VortexReader.open(file, ReadRegistry.loadAll())) { + // When + DType.Extension result = (DType.Extension) ((DType.Struct) reader.dtype()).fieldTypes().get(1); + + // Then — the unit byte and storage width Rust chose, which the decoder must honor + TimeUnit expectedUnit = unit == DateUnit.DAY ? TimeUnit.Days : TimeUnit.Milliseconds; + assertThat(TimeUnit.fromTag(result.metadata().get(ValueLayout.JAVA_BYTE, 0))).isEqualTo(expectedUnit); + assertThat(result.storageDType()).isEqualTo(unit == DateUnit.DAY ? DType.I32.asNullable() : DType.I64.asNullable()); + } + } + + /// Every ninth row is null; the rest step 37 days from 1915, so a quarter of them precede the + /// epoch, where flooring a millisecond count must not round toward zero. + private static LocalDate expected(long row) { + return row % 9 == 4 ? null : LocalDate.ofEpochDay(-20_000 + row * 37); + } + + private static void writeWithJni(Path file, DateUnit unit) throws IOException { + Schema schema = new Schema(List.of( + Field.notNullable("id", new ArrowType.Int(64, true)), + new Field("d", FieldType.nullable(new ArrowType.Date(unit)), null))); + try (VortexWriter writer = VortexWriter.builder(SESSION, file.toAbsolutePath().toUri().toString(), schema, + ALLOCATOR).build()) { + for (int chunk = 0; chunk < CHUNKS; chunk++) { + try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, ALLOCATOR)) { + BigIntVector ids = (BigIntVector) root.getVector("id"); + ids.allocateNew(CHUNK_ROWS); + root.getVector("d").setInitialCapacity(CHUNK_ROWS); + root.getVector("d").allocateNew(); + for (int i = 0; i < CHUNK_ROWS; i++) { + long row = (long) chunk * CHUNK_ROWS + i; + ids.setSafe(i, row); + LocalDate date = expected(row); + if (unit == DateUnit.DAY) { + DateDayVector days = (DateDayVector) root.getVector("d"); + if (date == null) { + days.setNull(i); + } else { + days.setSafe(i, Math.toIntExact(date.toEpochDay())); + } + } else { + DateMilliVector millis = (DateMilliVector) root.getVector("d"); + if (date == null) { + millis.setNull(i); + } else { + millis.setSafe(i, date.toEpochDay() * MILLIS_PER_DAY); + } + } + } + root.setRowCount(CHUNK_ROWS); + try (ArrowArray array = ArrowArray.allocateNew(ALLOCATOR); + ArrowSchema exported = ArrowSchema.allocateNew(ALLOCATOR)) { + Data.exportVectorSchemaRoot(ALLOCATOR, root, null, array, exported); + writer.writeBatch(array.memoryAddress(), exported.memoryAddress()); + } + } + } + } + } + + private static SequencedMap readJava(Path file, ScanOptions options) throws IOException { + SequencedMap out = new LinkedHashMap<>(); + try (VortexReader reader = VortexReader.open(file, ReadRegistry.loadAll()); + var iter = reader.scan(options)) { + while (iter.hasNext()) { + try (Chunk chunk = iter.next()) { + List dates = chunk.as(DATE, LocalDate.class); + LongArray ids = (LongArray) chunk.column(ID); + for (int i = 0; i < dates.size(); i++) { + out.put(ids.getLong(i), dates.get(i)); + } + } + } + } + return out; + } +} diff --git a/integration/src/test/java/io/github/dfa1/vortex/integration/JavaWritesRustReadsIntegrationTest.java b/integration/src/test/java/io/github/dfa1/vortex/integration/JavaWritesRustReadsIntegrationTest.java index 5ce644474..a27cf3b5e 100644 --- a/integration/src/test/java/io/github/dfa1/vortex/integration/JavaWritesRustReadsIntegrationTest.java +++ b/integration/src/test/java/io/github/dfa1/vortex/integration/JavaWritesRustReadsIntegrationTest.java @@ -61,6 +61,7 @@ import java.io.IOException; import java.io.UncheckedIOException; +import java.lang.foreign.MemorySegment; import java.nio.channels.FileChannel; import java.nio.file.Files; import java.nio.file.Path; @@ -79,6 +80,7 @@ import java.util.stream.Stream; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; /// Cross-compatibility: Java writer → Rust (JNI) reader. class JavaWritesRustReadsIntegrationTest { @@ -2286,6 +2288,79 @@ private static boolean hasDictLayout(io.github.dfa1.vortex.inspect.InspectorTree return false; } + @Test + void javaWriter_rustReader_extensionWithNullMetadata(@TempDir Path tmp) throws IOException { + // Given — an extension dtype with null metadata, here vortex.json (Utf8 storage, no + // metadata), which the vortex-jni session registers. The writer used to omit the metadata + // field for null, and Rust refuses an absent field ("failed to parse extension metadata") + // for the whole file, even when only a sibling column is projected. + Path file = tmp.resolve("java_json_null_metadata.vtx"); + DType.Extension json = new DType.Extension("vortex.json", DType.UTF8, null, false); + DType.Struct schema = new DType.Struct(List.of(ColumnName.of("id"), ColumnName.of("j")), + List.of(DType.I64, json), false); + try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE); + var sut = VortexWriter.create(ch, schema, WriteOptions.defaults())) { + // When + sut.writeChunk(Map.of(ColumnName.of("id"), new long[]{7L, 8L}, + ColumnName.of("j"), new String[]{"{\"a\":1}", "[]"})); + } + + // Then — vortex-jni opens the file and reads the sibling column + Scan scan = DataSource.open(SESSION, file.toAbsolutePath().toUri().toString()) + .scan(ScanOptions.builder() + .projection(Expression.select(new String[]{"id"}, Expression.root())) + .build()); + var ids = new ArrayList(); + while (scan.hasNext()) { + try (ArrowReader reader = scan.next().scanArrow(ALLOCATOR)) { + while (reader.loadNextBatch()) { + BigIntVector vec = (BigIntVector) reader.getVectorSchemaRoot().getVector("id"); + for (int i = 0; i < vec.getValueCount(); i++) { + ids.add(vec.get(i)); + } + } + } + } + assertThat(ids).containsExactly(7L, 8L); + } + + @Test + void invalidBuiltInExtensionDtype_refusedByJavaAndRust(@TempDir Path tmp) throws IOException { + // Given — vortex.timestamp over I32 storage, which Rust's validate_dtype refuses ("Timestamp + // storage dtype must be i64"). Java used to write and read it. + DType.Extension invalid = new DType.Extension("vortex.timestamp", DType.I32, + MemorySegment.ofArray(new byte[]{2, 0, 0}), false); + DType.Struct schema = new DType.Struct(List.of(ColumnName.of("ts")), List.of(invalid), false); + Path file = tmp.resolve("timestamp_i32.vtx"); + + // When / Then — the default writer refuses it at create + try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE)) { + assertThatThrownBy(() -> VortexWriter.create(ch, schema, WriteOptions.defaults())) + .isInstanceOf(io.github.dfa1.vortex.core.error.VortexException.class) + .hasMessageContaining("storage dtype must be I64"); + } + + // And — such a file, from another writer, is refused on read by both Java and vortex-jni. + // No Java write path emits it any more, so write the column under a custom id of the same + // length and rename it in place: a FlatBuffer string is stored verbatim, once. + DType.Struct disguised = new DType.Struct(List.of(ColumnName.of("ts")), List.of(new DType.Extension( + "vortex.timestamq", DType.I32, MemorySegment.ofArray(new byte[]{2, 0, 0}), false)), false); + try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE, + StandardOpenOption.TRUNCATE_EXISTING); + var writer = VortexWriter.create(ch, disguised, WriteOptions.defaults())) { + writer.writeChunk(Map.of(ColumnName.of("ts"), new int[]{1, 2})); + } + String bytes = Files.readString(file, java.nio.charset.StandardCharsets.ISO_8859_1); + assertThat(bytes.indexOf("vortex.timestamq")).isEqualTo(bytes.lastIndexOf("vortex.timestamq")).isNotNegative(); + Files.writeString(file, bytes.replace("vortex.timestamq", "vortex.timestamp"), java.nio.charset.StandardCharsets.ISO_8859_1); + assertThatThrownBy(() -> io.github.dfa1.vortex.reader.VortexReader.open(file)) + .isInstanceOf(io.github.dfa1.vortex.core.error.VortexException.class) + .hasMessageContaining("storage dtype must be I64"); + String uri = file.toAbsolutePath().toUri().toString(); + assertThatThrownBy(() -> DataSource.open(SESSION, uri).scan(ScanOptions.of())) + .hasMessageContaining("Timestamp storage dtype must be i64"); + } + @Test void javaWriter_rustReader_nullable_date(@TempDir Path tmp) throws IOException { // Given — nullable vortex.date column; row 1 is null. diff --git a/integration/src/test/java/io/github/dfa1/vortex/integration/WellKnownExtensionsInteropIntegrationTest.java b/integration/src/test/java/io/github/dfa1/vortex/integration/WellKnownExtensionsInteropIntegrationTest.java new file mode 100644 index 000000000..2d2710f5b --- /dev/null +++ b/integration/src/test/java/io/github/dfa1/vortex/integration/WellKnownExtensionsInteropIntegrationTest.java @@ -0,0 +1,168 @@ +package io.github.dfa1.vortex.integration; + +import dev.vortex.api.DataSource; +import dev.vortex.api.Partition; +import dev.vortex.api.Scan; +import dev.vortex.api.ScanOptions; +import dev.vortex.api.Session; +import dev.vortex.arrow.ArrowAllocation; +import dev.vortex.jni.NativeLoader; +import io.github.dfa1.vortex.core.model.ColumnName; +import io.github.dfa1.vortex.core.model.DType; +import io.github.dfa1.vortex.reader.Chunk; +import io.github.dfa1.vortex.reader.ReadRegistry; +import io.github.dfa1.vortex.reader.VortexReader; +import io.github.dfa1.vortex.reader.array.DoubleArray; +import io.github.dfa1.vortex.reader.array.LongArray; +import io.github.dfa1.vortex.reader.array.StructArray; +import io.github.dfa1.vortex.reader.array.VarBinArray; +import io.github.dfa1.vortex.writer.VortexWriter; +import io.github.dfa1.vortex.writer.WriteOptions; +import io.github.dfa1.vortex.writer.encode.StructData; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.complex.StructVector; +import org.apache.arrow.vector.ipc.ArrowReader; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.IOException; +import java.lang.foreign.MemorySegment; +import java.nio.channels.FileChannel; +import java.nio.charset.StandardCharsets; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +/// The extension ids the vortex-jni session registers by default beyond the temporal types and +/// uuid — `vortex.json` and the `vortex.st.*` geometries — open in Java without a decoder and read +/// as their storage. Before, a strict Java reader refused files vortex-jni reads. Java writes one +/// column of each storage shape (`Utf8`, `Binary`, a coordinate `Struct`) and both readers must +/// see the same values. +class WellKnownExtensionsInteropIntegrationTest { + + private static final Session SESSION = Session.create(); + private static final BufferAllocator ALLOCATOR = ArrowAllocation.rootAllocator(); + private static final MemorySegment NO_METADATA = MemorySegment.ofArray(new byte[0]); + private static final DType POINT_STORAGE = new DType.Struct( + List.of(ColumnName.of("x"), ColumnName.of("y")), List.of(DType.F64, DType.F64), false); + private static final DType.Struct SCHEMA = new DType.Struct( + List.of(ColumnName.of("id"), ColumnName.of("j"), ColumnName.of("g"), ColumnName.of("p")), + List.of(DType.I64, + new DType.Extension("vortex.json", DType.UTF8, NO_METADATA, false), + new DType.Extension("vortex.st.wkb", DType.BINARY, NO_METADATA, false), + new DType.Extension("vortex.st.point", POINT_STORAGE, NO_METADATA, false)), + false); + private static final int BATCHES = 2; + private static final int BATCH_ROWS = 500; + + static { + NativeLoader.loadJni(); + } + + @Test + void javaWrites_bothReadersSeeTheStorage(@TempDir Path tmp) throws IOException { + // Given + Path file = tmp.resolve("well_known.vortex"); + writeWithJava(file); + + // When + List java = readJava(file); + List jni = readJni(file); + + // Then + assertThat(java).hasSize(BATCHES * BATCH_ROWS); + for (Row row : java) { + assertThat(row).as("row %d", row.id()).isEqualTo(expected(row.id())); + } + assertThat(jni).containsExactlyElementsOf(java); + } + + /// One row as plain values; `wkb` is compared as a string so record equality works. + record Row(long id, String json, String wkb, double x, double y) { + } + + private static Row expected(long id) { + // a two-byte WKB-like prefix plus the row id: distinct per row, not valid UTF-8 at byte 0 + String wkb = "\u0001" + id; + return new Row(id, "{\"id\":" + id + ",\"tags\":[\"a\",\"é\"]}", wkb, id * 0.5, -id * 0.25); + } + + private static void writeWithJava(Path file) throws IOException { + try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE); + var writer = VortexWriter.create(ch, SCHEMA, WriteOptions.defaults())) { + for (int batch = 0; batch < BATCHES; batch++) { + long[] ids = new long[BATCH_ROWS]; + String[] json = new String[BATCH_ROWS]; + byte[][] wkb = new byte[BATCH_ROWS][]; + double[] xs = new double[BATCH_ROWS]; + double[] ys = new double[BATCH_ROWS]; + for (int i = 0; i < BATCH_ROWS; i++) { + long id = (long) batch * BATCH_ROWS + i; + Row row = expected(id); + ids[i] = id; + json[i] = row.json(); + wkb[i] = row.wkb().getBytes(StandardCharsets.ISO_8859_1); + xs[i] = row.x(); + ys[i] = row.y(); + } + writer.writeChunk(Map.of(ColumnName.of("id"), ids, ColumnName.of("j"), json, + ColumnName.of("g"), wkb, ColumnName.of("p"), new StructData(List.of(xs, ys)))); + } + } + } + + private static List readJava(Path file) throws IOException { + List rows = new ArrayList<>(); + // the default registry: no decoder for these ids, and no allowUnknown + try (VortexReader reader = VortexReader.open(file, ReadRegistry.loadAll()); + var iter = reader.scan(io.github.dfa1.vortex.reader.ScanOptions.all())) { + while (iter.hasNext()) { + try (Chunk chunk = iter.next()) { + LongArray ids = (LongArray) chunk.column(ColumnName.of("id")); + VarBinArray json = (VarBinArray) chunk.column(ColumnName.of("j")); + VarBinArray wkb = (VarBinArray) chunk.column(ColumnName.of("g")); + StructArray points = (StructArray) chunk.column(ColumnName.of("p")); + DoubleArray xs = (DoubleArray) points.field(0); + DoubleArray ys = (DoubleArray) points.field(1); + for (long i = 0; i < chunk.rowCount(); i++) { + rows.add(new Row(ids.getLong(i), json.getString(i), + new String(wkb.getBytes(i), StandardCharsets.ISO_8859_1), + xs.getDouble(i), ys.getDouble(i))); + } + } + } + } + return rows; + } + + private static List readJni(Path file) throws IOException { + List rows = new ArrayList<>(); + Scan scan = DataSource.open(SESSION, file.toAbsolutePath().toUri().toString()).scan(ScanOptions.of()); + while (scan.hasNext()) { + Partition partition = scan.next(); + try (ArrowReader reader = partition.scanArrow(ALLOCATOR)) { + while (reader.loadNextBatch()) { + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + BigIntVector ids = (BigIntVector) root.getVector("id"); + FieldVector json = root.getVector("j"); + FieldVector wkb = root.getVector("g"); + StructVector points = (StructVector) root.getVector("p"); + for (int i = 0; i < root.getRowCount(); i++) { + rows.add(new Row(ids.get(i), json.getObject(i).toString(), + new String((byte[]) wkb.getObject(i), StandardCharsets.ISO_8859_1), + (Double) points.getChild("x").getObject(i), + (Double) points.getChild("y").getObject(i))); + } + } + } + } + return rows; + } +} diff --git a/jdbc/src/main/java/io/github/dfa1/vortex/jdbc/JdbcImporter.java b/jdbc/src/main/java/io/github/dfa1/vortex/jdbc/JdbcImporter.java index aa47e5776..e3fe6f86d 100644 --- a/jdbc/src/main/java/io/github/dfa1/vortex/jdbc/JdbcImporter.java +++ b/jdbc/src/main/java/io/github/dfa1/vortex/jdbc/JdbcImporter.java @@ -218,8 +218,9 @@ private static boolean fillCell(Object buffer, int rowIdx, ResultSet rs, int col @SuppressWarnings("unchecked") private static void fillExtensionCell(List buffer, ResultSet rs, int colIdx, DType.Extension ext) throws SQLException { - ExtensionId id = ExtensionId.parse(ext.extensionId()) - .orElseThrow(() -> new UnsupportedOperationException("unsupported extension: " + ext.extensionId())); + if (!(ExtensionId.parse(ext.extensionId()) instanceof ExtensionId.WellKnown id)) { + throw new UnsupportedOperationException("unsupported extension: " + ext.extensionId()); + } // SQL NULL → null in the buffer. Nullable extension columns round-trip through the // writer's ExtEncoding → MaskedEncoding → primitive layout (validity child preserved); // NOT NULL columns reject any null element with VortexException during encode. @@ -240,6 +241,8 @@ private static void fillExtensionCell(List buffer, ResultSet rs, int col Object raw = rs.getObject(colIdx); buffer.add(rs.wasNull() ? null : toUuid(raw)); } + // SqlTypeToDType never maps a SQL type to the JSON or geometry ids + default -> throw new UnsupportedOperationException("unsupported extension: " + ext.extensionId()); } } diff --git a/jdbc/src/test/java/io/github/dfa1/vortex/jdbc/JdbcImporterTest.java b/jdbc/src/test/java/io/github/dfa1/vortex/jdbc/JdbcImporterTest.java index 99016a4a6..94137d5a6 100644 --- a/jdbc/src/test/java/io/github/dfa1/vortex/jdbc/JdbcImporterTest.java +++ b/jdbc/src/test/java/io/github/dfa1/vortex/jdbc/JdbcImporterTest.java @@ -133,7 +133,7 @@ void roundTripsTemporalSqlTypesViaExtensions(@TempDir Path tmp) throws Exception assertThat(chunk.rowCount()).isEqualTo(2); assertThat(io.github.dfa1.vortex.reader.extension.DateExtensionDecoder.INSTANCE - .decodeAll(chunk.column(ColumnName.of("D")))) + .decodeAll((DType.Extension) schema.fieldTypes().get(1), chunk.column(ColumnName.of("D")))) .containsExactly( java.time.LocalDate.of(1996, 2, 12), java.time.LocalDate.of(2026, 6, 10)); @@ -202,9 +202,10 @@ void roundTripsNullableExtensionColumns(@TempDir Path tmp) throws Exception { assertThat(dMasked.isValid(0)).isTrue(); assertThat(dMasked.isValid(1)).isFalse(); assertThat(dMasked.isValid(2)).isTrue(); - assertThat(io.github.dfa1.vortex.reader.extension.DateExtensionDecoder.INSTANCE.decode(dMasked, 0)) + DType.Extension dDtype = (DType.Extension) schema.fieldTypes().get(1); + assertThat(io.github.dfa1.vortex.reader.extension.DateExtensionDecoder.INSTANCE.decode(dDtype, dMasked, 0)) .isEqualTo(java.time.LocalDate.of(1996, 2, 12)); - assertThat(io.github.dfa1.vortex.reader.extension.DateExtensionDecoder.INSTANCE.decode(dMasked, 2)) + assertThat(io.github.dfa1.vortex.reader.extension.DateExtensionDecoder.INSTANCE.decode(dDtype, dMasked, 2)) .isEqualTo(java.time.LocalDate.of(2026, 6, 10)); DType.Extension tDtype = (DType.Extension) schema.fieldTypes().get(2); @@ -230,7 +231,7 @@ void roundTripsNullableExtensionColumns(@TempDir Path tmp) throws Exception { .isEqualTo(u1); // And — decodeAll preserves nulls in the returned list at row 1 - assertThat(io.github.dfa1.vortex.reader.extension.DateExtensionDecoder.INSTANCE.decodeAll(dMasked)) + assertThat(io.github.dfa1.vortex.reader.extension.DateExtensionDecoder.INSTANCE.decodeAll(dDtype, dMasked)) .containsExactly( java.time.LocalDate.of(1996, 2, 12), null, diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/Chunk.java b/reader/src/main/java/io/github/dfa1/vortex/reader/Chunk.java index 800e9a1a4..0470446ba 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/Chunk.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/Chunk.java @@ -5,19 +5,12 @@ import io.github.dfa1.vortex.core.error.VortexException; import io.github.dfa1.vortex.reader.array.Array; import io.github.dfa1.vortex.core.model.ExtensionId; -import io.github.dfa1.vortex.reader.extension.DateExtensionDecoder; -import io.github.dfa1.vortex.reader.extension.TimeExtensionDecoder; -import io.github.dfa1.vortex.reader.extension.TimestampExtensionDecoder; -import io.github.dfa1.vortex.reader.extension.UuidExtensionDecoder; +import io.github.dfa1.vortex.reader.extension.ExtensionDecoder; import java.lang.foreign.Arena; -import java.time.Instant; -import java.time.LocalDate; -import java.time.LocalTime; import java.util.List; import java.util.Objects; import java.util.SequencedMap; -import java.util.UUID; import java.util.function.Consumer; /// One decoded row group returned by [ScanIterator#next()]. @@ -49,14 +42,16 @@ public final class Chunk implements AutoCloseable { private final long rowCount; private final SequencedMap columns; private final Arena arena; + private final ReadRegistry registry; private final Consumer onClose; private boolean closed; Chunk(long rowCount, SequencedMap columns, - Arena arena, Consumer onClose) { + Arena arena, ReadRegistry registry, Consumer onClose) { this.rowCount = rowCount; this.columns = Objects.requireNonNull(columns); this.arena = Objects.requireNonNull(arena); + this.registry = Objects.requireNonNull(registry); this.onClose = Objects.requireNonNull(onClose); } @@ -109,25 +104,21 @@ public T column(ColumnName name) { return (T) col.array(); } - /// Decodes an extension column into a typed `List` of domain values. - /// - /// Hides extension-decode boilerplate for the four spec extensions: + /// Decodes an extension column into a typed `List` of domain values, through the + /// [io.github.dfa1.vortex.reader.extension.ExtensionDecoder] registered for the column's + /// extension id on the [ReadRegistry] the file was opened with: /// /// ```java /// List dates = chunk.as(ColumnName.of("birthdays"), LocalDate.class); /// List ts = chunk.as(ColumnName.of("events"), Instant.class); /// ``` /// - /// Third-party extensions call the extension's own typed methods directly — - /// this accessor is closed over the spec set. - /// /// @param name column name as declared in the file's schema - /// @param domainType the spec extension's Java domain type (e.g. [LocalDate] for - /// `vortex.date`, [Instant] for `vortex.timestamp`) + /// @param domainType the decoder's domain type (e.g. [java.time.LocalDate] for `vortex.date`) /// @param domain element type /// @return decoded values in row order - /// @throws VortexException if `name` isn't present, isn't an extension column, - /// or the requested `domainType` doesn't match the column's extension id + /// @throws VortexException if `name` isn't present, isn't an extension column, has no + /// registered decoder, or the decoder's domain type isn't `domainType` @SuppressWarnings("unchecked") public List as(ColumnName name, Class domainType) { Column col = columns.get(name); @@ -137,35 +128,16 @@ public List as(ColumnName name, Class domainType) { if (!(col.dtype() instanceof DType.Extension ext)) { throw new VortexException("not an extension column: " + name); } - ExtensionId id = ExtensionId.parse(ext.extensionId()) - .orElseThrow(() -> new VortexException("not a spec extension id: " + ext.extensionId())); - Array storage = col.array(); - Object result = switch (id) { - case VORTEX_DATE -> { - requireDomainType(name, domainType, LocalDate.class); - yield DateExtensionDecoder.INSTANCE.decodeAll(storage); - } - case VORTEX_TIME -> { - requireDomainType(name, domainType, LocalTime.class); - yield TimeExtensionDecoder.INSTANCE.decodeAll(ext, storage); - } - case VORTEX_TIMESTAMP -> { - requireDomainType(name, domainType, Instant.class); - yield TimestampExtensionDecoder.INSTANCE.decodeAll(ext, storage); - } - case VORTEX_UUID -> { - requireDomainType(name, domainType, UUID.class); - yield UuidExtensionDecoder.INSTANCE.decodeAll(storage); - } - }; - return (List) result; - } - - private static void requireDomainType(ColumnName name, Class requested, Class expected) { - if (!requested.equals(expected)) { - throw new VortexException("column '" + name + "' decodes to " + expected.getSimpleName() - + ", not " + requested.getSimpleName()); + ExtensionDecoder decoder = registry.lookup(ExtensionId.parse(ext.extensionId())); + if (decoder == null) { + throw new VortexException("no extension decoder registered for " + ext.extensionId() + + " (column '" + name + "')"); + } + if (!domainType.equals(decoder.domainType())) { + throw new VortexException("column '" + name + "' decodes to " + decoder.domainType().getSimpleName() + + ", not " + domainType.getSimpleName()); } + return (List) decoder.decodeAll(ext, col.array()); } /// Returns whether this chunk has already been closed. diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/PostscriptParser.java b/reader/src/main/java/io/github/dfa1/vortex/reader/PostscriptParser.java index 008ac74da..7e48af281 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/PostscriptParser.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/PostscriptParser.java @@ -3,6 +3,7 @@ import io.github.dfa1.vortex.core.model.DType; import io.github.dfa1.vortex.core.io.IoBounds; import io.github.dfa1.vortex.core.model.EncodingId; +import io.github.dfa1.vortex.core.model.ExtensionId; import io.github.dfa1.vortex.core.model.LayoutId; import io.github.dfa1.vortex.core.model.MemorySize; import io.github.dfa1.vortex.core.model.PType; @@ -153,6 +154,23 @@ private static EncodingId parseEncodingSpec(String raw, int index) { return EncodingId.parse(raw); } + /// Parses an extension dtype's wire id at the schema boundary, so every [DType.Extension] read + /// from a file carries an id [ExtensionId#parse] accepts. A missing, blank or control-character + /// id is untrusted-input corruption that would otherwise surface as an [IllegalArgumentException]. + /// + /// @param raw the wire string from the extension dtype, or `null` when the field is absent + /// @return the typed [ExtensionId] + private static ExtensionId parseExtensionId(String raw) { + if (raw == null) { + throw new VortexException("extension dtype missing id"); + } + try { + return ExtensionId.parse(raw); + } catch (IllegalArgumentException ex) { + throw new VortexException("invalid extension id in file schema: " + ex.getMessage()); + } + } + /// Parses one layout-spec wire string to its [LayoutId] at the footer boundary. A blank id /// is untrusted-input corruption; [LayoutId#parse] would reject it with an /// [IllegalArgumentException], so surface it as a [VortexException] with the spec index. @@ -321,10 +339,15 @@ private static DType convertDType(io.github.dfa1.vortex.core.fbs.FbsDType fbs, i case FbsType.FbsExtension -> { var e = fbs.type(new FbsExtension()); DType storage = convertDType(e.storageDtype(), depth + 1); + // Rust requires the field, empty when there is no metadata + MemorySegment metadata = e.metadataAsSegment(); + if (metadata == null) { + throw new VortexException("extension dtype missing metadata"); + } yield new DType.Extension( - e.id(), + parseExtensionId(e.id()).id(), storage, - e.metadataAsSegment(), + metadata, storage.nullable()); } case FbsType.FbsVariant -> new DType.Variant(fbs.type(new FbsVariant()).nullable()); diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/ReadRegistry.java b/reader/src/main/java/io/github/dfa1/vortex/reader/ReadRegistry.java index 88776ca36..fc73a86c9 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/ReadRegistry.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/ReadRegistry.java @@ -5,7 +5,9 @@ import io.github.dfa1.vortex.reader.array.UnknownArray; import io.github.dfa1.vortex.core.model.Edition; import io.github.dfa1.vortex.core.model.Editions; +import io.github.dfa1.vortex.core.model.DType; import io.github.dfa1.vortex.core.model.EncodingId; +import io.github.dfa1.vortex.core.model.ExtensionId; import io.github.dfa1.vortex.reader.decode.AlpEncodingDecoder; import io.github.dfa1.vortex.reader.decode.AlpRdEncodingDecoder; import io.github.dfa1.vortex.reader.decode.ArrayNode; @@ -47,15 +49,22 @@ import io.github.dfa1.vortex.reader.decode.ZigZagEncodingDecoder; import io.github.dfa1.vortex.reader.decode.ZstdBuffersEncodingDecoder; import io.github.dfa1.vortex.reader.decode.ZstdEncodingDecoder; +import io.github.dfa1.vortex.reader.extension.DateExtensionDecoder; +import io.github.dfa1.vortex.reader.extension.ExtensionDecoder; +import io.github.dfa1.vortex.reader.extension.TimeExtensionDecoder; +import io.github.dfa1.vortex.reader.extension.TimestampExtensionDecoder; +import io.github.dfa1.vortex.reader.extension.UuidExtensionDecoder; import java.lang.foreign.MemorySegment; import java.util.Collections; import java.util.Comparator; +import java.util.HashMap; import java.util.Map; import java.util.Optional; import java.util.TreeMap; -/// Read-side registry: maps [EncodingId] to [EncodingDecoder] implementations. +/// Read-side registry: maps [EncodingId] to [EncodingDecoder] implementations, +/// and [ExtensionId] to [ExtensionDecoder] implementations. /// /// Instances are immutable after construction. Build one via [#builder()] or /// via the [#loadAll()] and [#empty()] convenience factories. @@ -66,16 +75,19 @@ public final class ReadRegistry { // that string (EncodingId is not Comparable; a Custom key must not throw); order is not // load-bearing, but stable ordering keeps the registries consistent. private final Map decoders; + private final Map> extensions; private final boolean allowUnknown; - private ReadRegistry(Map decoders, boolean allowUnknown) { + private ReadRegistry(Map decoders, + Map> extensions, boolean allowUnknown) { var sorted = new TreeMap(Comparator.comparing(EncodingId::id)); sorted.putAll(decoders); this.decoders = Collections.unmodifiableMap(sorted); + this.extensions = Map.copyOf(extensions); this.allowUnknown = allowUnknown; } - /// Loads all built-in [EncodingDecoder] implementations. + /// Loads all built-in [EncodingDecoder] and [ExtensionDecoder] implementations. /// /// @return an immutable [ReadRegistry] populated with all built-in decoders public static ReadRegistry loadAll() { @@ -103,10 +115,11 @@ public Map decoderMap() { return decoders; } - /// Returns whether passthrough decode for unknown encoding ids is enabled. + /// Returns whether passthrough decode for unknown encoding and extension ids is enabled. /// /// @return `true` if unknown encodings are silently wrapped as - /// [io.github.dfa1.vortex.reader.array.UnknownArray] + /// [io.github.dfa1.vortex.reader.array.UnknownArray] and unknown extension ids read as + /// their storage array public boolean isAllowUnknown() { return allowUnknown; } @@ -119,6 +132,85 @@ public boolean hasDecoder(EncodingId encodingId) { return decoders.containsKey(encodingId); } + /// Returns the registered extension decoder for `extensionId`, or `null` if not registered. + /// A file with an unregistered extension id opens only with [Builder#allowUnknown()], and the + /// column then reads as its storage array (Rust's foreign extension). + /// + /// @param extensionId the extension id to look up + /// @return the registered [ExtensionDecoder], or `null` + public ExtensionDecoder lookup(ExtensionId extensionId) { + return extensions.get(extensionId); + } + + /// Checks every extension dtype in `dtype`, nested ones included, as Rust does when it + /// deserializes one: a registered decoder validates it; a [ExtensionId.WellKnown] id without a + /// decoder is validated by [ExtensionId.WellKnown#validate(DType.Extension)] and reads as its + /// storage; any other id is refused unless [Builder#allowUnknown()] is set, when it reads as its + /// storage (Rust's `ForeignExtDType`). + /// Run on a file's schema at open, and on the dtypes embedded in Variant metadata at decode. + /// + /// @param dtype a file's schema, or `null` when the file carries no dtype blob + /// @throws VortexException if a registered decoder rejects its dtype, or an extension id has no + /// decoder and unknown ids are not allowed + public void validateExtensions(DType dtype) { + switch (dtype) { + case null -> { + // a file without a dtype blob parses with no schema: nothing to check + } + case DType.Extension ext -> { + ExtensionId id = parseId(ext.extensionId()); + ExtensionDecoder decoder = extensions.get(id); + if (decoder != null) { + validate(decoder, ext); + } else if (id instanceof ExtensionId.WellKnown wellKnown) { + // the reference session knows it: no decoder needed, the column reads as storage + wellKnown.validate(ext); + } else if (!allowUnknown) { + throw new VortexException("no extension decoder registered for " + ext.extensionId() + + "; register one or enable ReadRegistry.Builder#allowUnknown()"); + } + validateExtensions(ext.storageDType()); + } + case DType.Struct struct -> struct.fieldTypes().forEach(this::validateExtensions); + case DType.List list -> validateExtensions(list.elementType()); + case DType.FixedSizeList list -> validateExtensions(list.elementType()); + case DType.Map map -> { + validateExtensions(map.keyType()); + validateExtensions(map.valueType()); + } + case DType.Union union -> union.variantTypes().forEach(this::validateExtensions); + case DType.Null _, DType.Bool _, DType.Primitive _, DType.Decimal _, DType.Utf8 _, + DType.Binary _, DType.Variant _ -> { + // no nested dtype + } + } + } + + /// Parses an extension id read from a file; an id [ExtensionId#parse] refuses is corrupt input. + private static ExtensionId parseId(String raw) { + if (raw == null) { + throw new VortexException("extension dtype missing id"); + } + try { + return ExtensionId.parse(raw); + } catch (IllegalArgumentException e) { + throw new VortexException("invalid extension id: " + e.getMessage()); + } + } + + /// Runs a decoder's [ExtensionDecoder#validate(DType.Extension)] on a dtype from an untrusted + /// file. A third-party check may trip over malformed metadata with any runtime exception; + /// reading a file must only ever fail with [VortexException]. + private static void validate(ExtensionDecoder decoder, DType.Extension ext) { + try { + decoder.validate(ext); + } catch (VortexException e) { + throw e; + } catch (RuntimeException e) { + throw new VortexException("extension " + ext.extensionId() + " rejected its dtype", e); + } + } + /// Decodes the array described by `ctx`. /// /// @param ctx the decode context @@ -189,6 +281,7 @@ private static UnknownArray decodeUnknown(DecodeContext ctx, ArrayNode node) { public static final class Builder { private final Map decoders = new TreeMap<>(Comparator.comparing(EncodingId::id)); + private final Map> extensions = new HashMap<>(); private boolean allowUnknown = false; private Builder() { @@ -207,11 +300,28 @@ public Builder register(EncodingDecoder decoder) { return this; } - /// Registers all built-in [EncodingDecoder] implementations. + /// Registers an extension decoder. + /// + /// @param extension the [ExtensionDecoder] to register + /// @return this builder, for chaining + /// @throws VortexException if an extension with the same id is already registered + public Builder register(ExtensionDecoder extension) { + ExtensionDecoder old = extensions.put(extension.extensionId(), extension); + if (old != null) { + throw new VortexException("extension %s already registered".formatted(extension.extensionId())); + } + return this; + } + + /// Registers all built-in [EncodingDecoder] and [ExtensionDecoder] implementations. /// /// @return this builder, for chaining - /// @throws VortexException if a built-in decoder collides with one already registered + /// @throws VortexException if a built-in entry collides with one already registered public Builder registerDefaults() { + register(DateExtensionDecoder.INSTANCE) + .register(TimeExtensionDecoder.INSTANCE) + .register(TimestampExtensionDecoder.INSTANCE) + .register(UuidExtensionDecoder.INSTANCE); return register(new AlpEncodingDecoder()) .register(new AlpRdEncodingDecoder()) .register(new BitpackedEncodingDecoder()) @@ -252,10 +362,11 @@ public Builder registerDefaults() { .register(new ZstdEncodingDecoder()); } - /// Enable passthrough decode for unknown encoding ids. + /// Enable passthrough decode for unknown encoding and extension ids. /// /// Default is strict: unknown ids throw [VortexException]. When enabled, unknown - /// nodes are wrapped as [io.github.dfa1.vortex.reader.array.UnknownArray]. + /// encoding nodes are wrapped as [io.github.dfa1.vortex.reader.array.UnknownArray], and + /// an extension id with no registered [ExtensionDecoder] reads as its storage array. /// Mirrors Rust `VortexSession::allow_unknown()`. /// /// @return this builder, for chaining @@ -268,7 +379,7 @@ public Builder allowUnknown() { /// /// @return the immutable registry public ReadRegistry build() { - return new ReadRegistry(decoders, allowUnknown); + return new ReadRegistry(decoders, extensions, allowUnknown); } } } diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/ScanIterator.java b/reader/src/main/java/io/github/dfa1/vortex/reader/ScanIterator.java index 1700058d6..b21b09584 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/ScanIterator.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/ScanIterator.java @@ -350,7 +350,7 @@ public Chunk next() { columns = limitedColumns(columns, chunkRows); } rowsReturned += chunkRows; - Chunk chunk = new Chunk(chunkRows, columns, arena, this::onChunkClosed); + Chunk chunk = new Chunk(chunkRows, columns, arena, file.registry(), this::onChunkClosed); openChunk = chunk; return chunk; } catch (RuntimeException e) { @@ -395,7 +395,7 @@ Chunk decodeChunkAt(int chunkIndex) { Arena arena = Arena.ofConfined(); try { SequencedMap columns = buildSelfContainedColumnMap(spec, arena); - return new Chunk(spec.rowCount(), columns, arena, _ -> { }); + return new Chunk(spec.rowCount(), columns, arena, file.registry(), _ -> { }); } catch (RuntimeException e) { arena.close(); throw e; diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/VortexHttpReader.java b/reader/src/main/java/io/github/dfa1/vortex/reader/VortexHttpReader.java index 09f296e72..865de7f7a 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/VortexHttpReader.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/VortexHttpReader.java @@ -147,6 +147,7 @@ public static VortexHttpReader open(URI uri, ReadRegistry registry, LayoutRegist // Range request from them. The local-file path runs this inside PostscriptParser.parse; // the HTTP path calls parseBlobs directly and must validate here too. PostscriptParser.validateSegmentSpecs(parsed.footer().segmentSpecs(), fileSize); + registry.validateExtensions(parsed.dtype()); return new VortexHttpReader( uri, client, fileSize, trailer.version(), diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/VortexReader.java b/reader/src/main/java/io/github/dfa1/vortex/reader/VortexReader.java index 7f7cec0c4..704a58b76 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/VortexReader.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/VortexReader.java @@ -107,6 +107,7 @@ private static VortexReader parse( } catch (RuntimeException e) { throw new VortexException("malformed postscript", e); } + registry.validateExtensions(parsed.dtype()); return new VortexReader( arena, seg, size, trailer.version(), diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/decode/ConstantEncodingDecoder.java b/reader/src/main/java/io/github/dfa1/vortex/reader/decode/ConstantEncodingDecoder.java index 81a8a0e5e..3a505da08 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/decode/ConstantEncodingDecoder.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/decode/ConstantEncodingDecoder.java @@ -62,7 +62,7 @@ private static Array arrayFromScalar(DecodeContext ctx, ProtoScalarValue scalar, if (inner == null || inner.value() == null) { throw new VortexException(EncodingId.VORTEX_CONSTANT, "constant variant missing variant_value"); } - DType innerDtype = VariantEncodingDecoder.dtypeFromProto(inner.dtype()); + DType innerDtype = VariantEncodingDecoder.dtypeFromProto(inner.dtype(), ctx.registry()); return arrayFromScalar(ctx, inner.value(), innerDtype, n); } if (dtype instanceof DType.Utf8 || dtype instanceof DType.Binary) { diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/decode/ParquetVariantEncodingDecoder.java b/reader/src/main/java/io/github/dfa1/vortex/reader/decode/ParquetVariantEncodingDecoder.java index 8639f28aa..fb53fd333 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/decode/ParquetVariantEncodingDecoder.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/decode/ParquetVariantEncodingDecoder.java @@ -39,7 +39,7 @@ public Array decode(DecodeContext ctx) { } ProtoParquetVariantMetadata meta = parseMetadata(ctx.metadata()); DType typedValueDtype = meta.typed_value_dtype() == null - ? null : VariantEncodingDecoder.dtypeFromProto(meta.typed_value_dtype()); + ? null : VariantEncodingDecoder.dtypeFromProto(meta.typed_value_dtype(), ctx.registry()); if (!meta.has_value() && typedValueDtype == null) { throw new VortexException(EncodingId.VORTEX_PARQUET_VARIANT, "at least one of value or typed_value must be present"); diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/decode/VariantEncodingDecoder.java b/reader/src/main/java/io/github/dfa1/vortex/reader/decode/VariantEncodingDecoder.java index f25d8e3a9..4747e81dd 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/decode/VariantEncodingDecoder.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/decode/VariantEncodingDecoder.java @@ -7,6 +7,7 @@ import io.github.dfa1.vortex.reader.array.VariantArray; import io.github.dfa1.vortex.core.model.EncodingId; import io.github.dfa1.vortex.core.proto.ProtoVariantMetadata; +import io.github.dfa1.vortex.reader.ReadRegistry; import java.io.IOException; import java.lang.foreign.MemorySegment; @@ -23,7 +24,7 @@ public EncodingId encodingId() { @Override public Array decode(DecodeContext ctx) { - DType shreddedDtype = parseShreddedDtype(ctx.metadata()); + DType shreddedDtype = parseShreddedDtype(ctx.metadata(), ctx.registry()); int numChildren = ctx.node().children().length; if (numChildren < 1 || numChildren > 2) { @@ -41,7 +42,7 @@ public Array decode(DecodeContext ctx) { return new VariantArray(ctx.dtype(), ctx.rowCount(), coreStorage, shredded); } - private static DType parseShreddedDtype(MemorySegment rawMeta) { + private static DType parseShreddedDtype(MemorySegment rawMeta, ReadRegistry registry) { if (rawMeta == null || rawMeta.byteSize() == 0) { return null; } @@ -51,12 +52,26 @@ private static DType parseShreddedDtype(MemorySegment rawMeta) { if (meta.shredded_dtype() == null) { return null; } - return dtypeFromProto(meta.shredded_dtype()); + return dtypeFromProto(meta.shredded_dtype(), registry); } catch (IOException e) { throw new VortexException(EncodingId.VORTEX_VARIANT, "invalid metadata", e); } } + /// Converts a dtype embedded in array metadata and checks its extension dtypes against + /// `registry`, as Rust's `DType::from_proto` resolves each extension id on the session. + /// + /// @param proto the embedded dtype + /// @param registry the registry of the file being read + /// @return the converted dtype + /// @throws VortexException if an extension id is unregistered (without allow-unknown) or its + /// decoder rejects the dtype + static DType dtypeFromProto(io.github.dfa1.vortex.core.proto.ProtoDType proto, ReadRegistry registry) { + DType dtype = dtypeFromProto(proto); + registry.validateExtensions(dtype); + return dtype; + } + static DType dtypeFromProto(io.github.dfa1.vortex.core.proto.ProtoDType proto) { if (proto.null_() != null) { return new DType.Null(true); @@ -105,11 +120,13 @@ static DType dtypeFromProto(io.github.dfa1.vortex.core.proto.ProtoDType proto) { proto.fixed_size_list().nullable()); } if (proto.extension() != null) { + // an extension's nullability is its storage's, as the file-schema parser has it + DType storage = dtypeFromProto(proto.extension().storage_dtype()); return new DType.Extension( proto.extension().id(), - dtypeFromProto(proto.extension().storage_dtype()), + storage, MemorySegment.ofArray(proto.extension().metadata() != null ? proto.extension().metadata() : new byte[0]).asReadOnly(), - false); + storage.nullable()); } if (proto.variant() != null) { return new DType.Variant(proto.variant().nullable()); diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/extension/DateExtensionDecoder.java b/reader/src/main/java/io/github/dfa1/vortex/reader/extension/DateExtensionDecoder.java index ca6c724a9..b4ace8f98 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/extension/DateExtensionDecoder.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/extension/DateExtensionDecoder.java @@ -1,5 +1,6 @@ package io.github.dfa1.vortex.reader.extension; +import io.github.dfa1.vortex.core.error.VortexException; import io.github.dfa1.vortex.core.model.DType; import io.github.dfa1.vortex.core.io.IoBounds; import io.github.dfa1.vortex.core.model.PType; @@ -12,18 +13,31 @@ import java.util.ArrayList; import java.util.List; -/// `vortex.date` — days since the Unix epoch, signed integer storage. -/// Per Arrow's canonical Date type. -public final class DateExtensionDecoder implements ExtensionDecoder { +/// `vortex.date` — days (`I32`) or milliseconds (`I64`) since the Unix epoch, per the time unit in +/// the metadata. Per Arrow's canonical `Date32`/`Date64` types. +public final class DateExtensionDecoder implements ExtensionDecoder { /// Singleton instance. public static final DateExtensionDecoder INSTANCE = new DateExtensionDecoder(); + private static final long MILLIS_PER_DAY = 86_400_000L; + @Override public ExtensionId extensionId() { return ExtensionId.VORTEX_DATE; } + /// Applies Rust's `validate_dtype` rules for `vortex.date`. + @Override + public void validate(DType.Extension dtype) { + ExtensionId.VORTEX_DATE.validate(dtype); + } + + @Override + public Class domainType() { + return LocalDate.class; + } + @Override public DType.Extension dtype(boolean nullable) { // Rust vortex.date metadata: 1 byte = TimeUnit tag (Days = 4), required by Rust reader. @@ -35,34 +49,60 @@ public DType.Extension dtype(boolean nullable) { nullable); } - /// Decodes the date cell at row `i`. + /// Decodes the date cell at row `i`, in the column's time unit. /// + /// @param dtype declared extension dtype, carrying the time unit /// @param storage signed-integer storage (Byte/Short/Int/Long, possibly Masked) /// @param i row index, `0 <= i < storage.length()` /// @return decoded date /// @throws io.github.dfa1.vortex.core.error.VortexException if storage isn't an integer primitive - public LocalDate decode(Array storage, long i) { + /// or the unit is not Days or Milliseconds + public LocalDate decode(DType.Extension dtype, Array storage, long i) { ExtensionStorage.checkBounds(i, storage.length()); - return LocalDate.ofEpochDay(ExtensionStorage.epochInteger(storage, i)); + return date(ExtensionStorage.readUnit(dtype), ExtensionStorage.epochInteger(storage, i)); + } + + /// Converts one raw storage value to a date, in the column's time unit. + /// + /// @param dtype declared extension dtype, carrying the time unit + /// @param raw storage value: days, or milliseconds, since the Unix epoch + /// @return the date + /// @throws io.github.dfa1.vortex.core.error.VortexException if the unit is not Days or Milliseconds + public LocalDate date(DType.Extension dtype, long raw) { + return date(ExtensionStorage.readUnit(dtype), raw); } /// Decodes every row of `storage` into a list of dates. [MaskedArray] /// storage yields `null` at invalid positions instead of throwing. /// + /// @param dtype declared extension dtype, carrying the time unit /// @param storage signed-integer storage array (optionally wrapped in `MaskedArray`) /// @return list of decoded dates in row order; `null` entries mark invalid rows - public List decodeAll(Array storage) { + @Override + public List decodeAll(DType.Extension dtype, Array storage) { + TimeUnit unit = ExtensionStorage.readUnit(dtype); int n = IoBounds.toIntSize(storage.length()); List out = new ArrayList<>(n); if (storage instanceof MaskedArray masked) { for (int i = 0; i < n; i++) { - out.add(masked.isValid(i) ? decode(masked.inner(), i) : null); + out.add(masked.isValid(i) ? date(unit, ExtensionStorage.epochInteger(masked.inner(), i)) : null); } return out; } for (int i = 0; i < n; i++) { - out.add(decode(storage, i)); + out.add(date(unit, ExtensionStorage.epochInteger(storage, i))); } return out; } + + /// Rust stores an Arrow `Date32` as Days and a `Date64` as Milliseconds. A millisecond count + /// floors to its day, so an instant before the epoch lands on the earlier date. + private static LocalDate date(TimeUnit unit, long raw) { + return switch (unit) { + case Days -> LocalDate.ofEpochDay(raw); + case Milliseconds -> LocalDate.ofEpochDay(Math.floorDiv(raw, MILLIS_PER_DAY)); + case Nanoseconds, Microseconds, Seconds -> + throw new VortexException("vortex.date does not support time unit " + unit); + }; + } } diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/extension/ExtensionDecoder.java b/reader/src/main/java/io/github/dfa1/vortex/reader/extension/ExtensionDecoder.java index ed4d42148..5b1bccb88 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/extension/ExtensionDecoder.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/extension/ExtensionDecoder.java @@ -2,19 +2,49 @@ import io.github.dfa1.vortex.core.model.DType; import io.github.dfa1.vortex.core.model.ExtensionId; +import io.github.dfa1.vortex.reader.array.Array; -/// Read-side contract for a Vortex extension type. +import java.util.List; + +/// Read-side contract for a Vortex extension type, the Java side of Rust's `ExtVTable`. +/// +/// Implementations pair an identity ([ExtensionId]) with the matching [DType.Extension] dtype and +/// decode a column's storage array into domain values. Register one on +/// [io.github.dfa1.vortex.reader.ReadRegistry.Builder#register(ExtensionDecoder)] and +/// [io.github.dfa1.vortex.reader.Chunk#as(io.github.dfa1.vortex.core.model.ColumnName, Class)] +/// dispatches to it. A file with an extension id that has no registered decoder fails to open +/// unless [io.github.dfa1.vortex.reader.ReadRegistry.Builder#allowUnknown()] is set; the column +/// then reads as its storage array. /// -/// Implementations pair a spec identity ([ExtensionId]) with the matching -/// [DType.Extension] dtype. Typed decode methods live on each concrete -/// implementation — they are not on this interface, so read callers get typed return -/// values without casting through `Object`. -public interface ExtensionDecoder { +/// @param the domain type one row decodes to (e.g. [java.time.LocalDate] for `vortex.date`) +public interface ExtensionDecoder { - /// @return the spec identity of this extension + /// @return the identity of this extension ExtensionId extensionId(); /// @param nullable whether the column allows nulls /// @return matching [DType.Extension] (storage dtype + default metadata) DType.Extension dtype(boolean nullable); + + /// @return the domain type one row decodes to + Class domainType(); + + /// Decodes every row of `storage` into domain values. + /// + /// @param dtype declared extension dtype (carries metadata such as a time unit) + /// @param storage the column's storage array; nullable columns arrive as a + /// [io.github.dfa1.vortex.reader.array.MaskedArray] + /// @return decoded values in row order, `null` at null rows + List decodeAll(DType.Extension dtype, Array storage); + + /// Rejects a dtype this extension cannot represent — Rust's `validate_dtype`. Called for every + /// extension dtype with this id in a file's schema when the file is opened. + /// + /// The default accepts every dtype. + /// + /// @param dtype the extension dtype declared by the file + /// @throws io.github.dfa1.vortex.core.error.VortexException if the storage dtype or metadata is + /// invalid for this extension + default void validate(DType.Extension dtype) { + } } diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/extension/TimeExtensionDecoder.java b/reader/src/main/java/io/github/dfa1/vortex/reader/extension/TimeExtensionDecoder.java index d6ed8787d..b1bbcc883 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/extension/TimeExtensionDecoder.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/extension/TimeExtensionDecoder.java @@ -14,7 +14,7 @@ import java.util.List; /// `vortex.time` — sub-day count in the [TimeUnit] recorded in the metadata byte. -public final class TimeExtensionDecoder implements ExtensionDecoder { +public final class TimeExtensionDecoder implements ExtensionDecoder { /// Singleton instance. public static final TimeExtensionDecoder INSTANCE = new TimeExtensionDecoder(); @@ -24,6 +24,17 @@ public ExtensionId extensionId() { return ExtensionId.VORTEX_TIME; } + /// Applies Rust's `validate_dtype` rules for `vortex.time`. + @Override + public void validate(DType.Extension dtype) { + ExtensionId.VORTEX_TIME.validate(dtype); + } + + @Override + public Class domainType() { + return LocalTime.class; + } + /// Returns the default dtype using [TimeUnit#Milliseconds] over I32 storage. /// Use [#dtype(TimeUnit, boolean)] for non-default units. @Override @@ -65,6 +76,7 @@ public LocalTime decode(DType.Extension ext, Array storage, long i) { /// @param ext declared extension dtype carrying the unit /// @param storage signed-integer storage array (optionally wrapped in `MaskedArray`) /// @return list of decoded times in row order; `null` entries mark invalid rows + @Override public List decodeAll(DType.Extension ext, Array storage) { int n = IoBounds.toIntSize(storage.length()); List out = new ArrayList<>(n); diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/extension/TimestampExtensionDecoder.java b/reader/src/main/java/io/github/dfa1/vortex/reader/extension/TimestampExtensionDecoder.java index d8968dde4..a6a107bb4 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/extension/TimestampExtensionDecoder.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/extension/TimestampExtensionDecoder.java @@ -18,7 +18,7 @@ import java.util.Optional; /// `vortex.timestamp` — I64 epoch count plus optional IANA timezone. -public final class TimestampExtensionDecoder implements ExtensionDecoder { +public final class TimestampExtensionDecoder implements ExtensionDecoder { /// Singleton instance. public static final TimestampExtensionDecoder INSTANCE = new TimestampExtensionDecoder(); @@ -28,6 +28,17 @@ public ExtensionId extensionId() { return ExtensionId.VORTEX_TIMESTAMP; } + /// Applies Rust's `validate_dtype` rules for `vortex.timestamp`. + @Override + public void validate(DType.Extension dtype) { + ExtensionId.VORTEX_TIMESTAMP.validate(dtype); + } + + @Override + public Class domainType() { + return Instant.class; + } + /// Returns the default dtype using milliseconds resolution and no timezone. /// Use [#dtype(TimeUnit, ZoneId, boolean)] for non-default settings. @Override @@ -88,6 +99,7 @@ public Optional timezone(DType.Extension ext) { /// @param ext declared extension dtype carrying the unit /// @param storage signed-integer storage array (optionally wrapped in `MaskedArray`) /// @return list of decoded instants in row order; `null` entries mark invalid rows + @Override public List decodeAll(DType.Extension ext, Array storage) { int n = IoBounds.toIntSize(storage.length()); List out = new ArrayList<>(n); diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/extension/UuidExtensionDecoder.java b/reader/src/main/java/io/github/dfa1/vortex/reader/extension/UuidExtensionDecoder.java index 348efbd8a..98c7d404e 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/extension/UuidExtensionDecoder.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/extension/UuidExtensionDecoder.java @@ -14,7 +14,7 @@ import java.util.UUID; /// `vortex.uuid` — 16-byte UUID stored as `FixedSizeList(Primitive(U8), 16)`. -public final class UuidExtensionDecoder implements ExtensionDecoder { +public final class UuidExtensionDecoder implements ExtensionDecoder { /// Singleton instance. public static final UuidExtensionDecoder INSTANCE = new UuidExtensionDecoder(); @@ -24,6 +24,22 @@ public ExtensionId extensionId() { return ExtensionId.VORTEX_UUID; } + /// Applies Rust's `validate_dtype` rules for `vortex.uuid`. + @Override + public void validate(DType.Extension dtype) { + ExtensionId.VORTEX_UUID.validate(dtype); + } + + @Override + public Class domainType() { + return UUID.class; + } + + @Override + public List decodeAll(DType.Extension dtype, Array storage) { + return decodeAll(storage); + } + @Override public DType.Extension dtype(boolean nullable) { DType.Primitive u8 = DType.U8; diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/layout/LayoutRegistry.java b/reader/src/main/java/io/github/dfa1/vortex/reader/layout/LayoutRegistry.java index 188be0037..ddf249424 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/layout/LayoutRegistry.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/layout/LayoutRegistry.java @@ -16,7 +16,7 @@ /// Instances are immutable after construction. Build one via [#builder()], or take the four /// built-ins directly via [#defaults()]. /// -/// Layouts register programmatically only (mirroring `ExtensionDecoder`). Unknown layout ids fail +/// Layouts register programmatically only (mirroring encodings and extensions). Unknown layout ids fail /// loudly: there is no allow-unknown mode, matching the Rust reference. public final class LayoutRegistry { diff --git a/reader/src/test/java/io/github/dfa1/vortex/reader/ComputeFilteredAggregateTest.java b/reader/src/test/java/io/github/dfa1/vortex/reader/ComputeFilteredAggregateTest.java index 6082acfcd..40b0187ca 100644 --- a/reader/src/test/java/io/github/dfa1/vortex/reader/ComputeFilteredAggregateTest.java +++ b/reader/src/test/java/io/github/dfa1/vortex/reader/ComputeFilteredAggregateTest.java @@ -625,7 +625,7 @@ private static Array floatArray(float... values) { private static Chunk chunk(long rows, Map columns) { SequencedMap typed = new LinkedHashMap<>(); columns.forEach((name, array) -> typed.put(ColumnName.of(name), new Chunk.Column(array, array.dtype()))); - return new Chunk(rows, typed, ARENA, c -> { /* no-op: the shared arena outlives the chunk */ }); + return new Chunk(rows, typed, ARENA, ReadRegistry.loadAll(), c -> { /* no-op: the shared arena outlives the chunk */ }); } /// Builds a long-domain column from a reference, optionally narrowing it to an `i32` array, and diff --git a/reader/src/test/java/io/github/dfa1/vortex/reader/PostscriptParserDTypeGuardsTest.java b/reader/src/test/java/io/github/dfa1/vortex/reader/PostscriptParserDTypeGuardsTest.java index f967cf120..9e62e6437 100644 --- a/reader/src/test/java/io/github/dfa1/vortex/reader/PostscriptParserDTypeGuardsTest.java +++ b/reader/src/test/java/io/github/dfa1/vortex/reader/PostscriptParserDTypeGuardsTest.java @@ -6,6 +6,7 @@ import io.github.dfa1.vortex.core.fbs.FbsArraySpec; import io.github.dfa1.vortex.core.fbs.FbsBuilder; import io.github.dfa1.vortex.core.fbs.FbsDType; +import io.github.dfa1.vortex.core.fbs.FbsExtension; import io.github.dfa1.vortex.core.fbs.FbsFooter; import io.github.dfa1.vortex.core.fbs.FbsLayout; import io.github.dfa1.vortex.core.fbs.FbsLayoutSpec; @@ -111,6 +112,62 @@ void convertFooter_blankArraySpec_throwsVortexException() { .hasMessageContaining("blank encoding id at array spec index 0"); } + @Test + void convertDType_customExtensionId_isAccepted() { + // Given — an extension id this build has no constant for (Rust reads it as a foreign + // extension over its storage) + MemorySegment dtype = extensionDType("acme.geopoint"); + + // When + PostscriptParser.ParsedFile result = PostscriptParser.parseBlobs(minimalFooter(), flatLayout(), dtype); + + // Then + assertThat(result.dtype()).isInstanceOfSatisfying(DType.Extension.class, ext -> + assertThat(ext.extensionId()).isEqualTo("acme.geopoint")); + } + + @org.junit.jupiter.params.ParameterizedTest + @org.junit.jupiter.params.provider.ValueSource(strings = {"", "a\nb"}) + void convertDType_invalidExtensionId_throwsVortexException(String id) { + // Given — an id ExtensionId.parse rejects with IllegalArgumentException; validating at the + // schema boundary lets every later ExtensionId.parse of a read dtype run unguarded + MemorySegment dtype = extensionDType(id); + + // When / Then + MemorySegment footer = minimalFooter(); + MemorySegment layout = flatLayout(); + assertThatThrownBy(() -> PostscriptParser.parseBlobs(footer, layout, dtype)) + .isInstanceOf(VortexException.class) + .hasMessageContaining("invalid extension id in file schema"); + } + + @Test + void convertDType_extensionMissingId_throwsVortexException() { + // Given — FlatBuffers fields are optional on the wire; a crafted extension omits its id + MemorySegment dtype = extensionDType(null); + + // When / Then + MemorySegment footer = minimalFooter(); + MemorySegment layout = flatLayout(); + assertThatThrownBy(() -> PostscriptParser.parseBlobs(footer, layout, dtype)) + .isInstanceOf(VortexException.class) + .hasMessageContaining("extension dtype missing id"); + } + + @Test + void convertDType_extensionMissingMetadata_throwsVortexException() { + // Given — no metadata field at all, which older vortex-java wrote for null metadata. Rust + // refuses it ("failed to parse extension metadata from flatbuffer"); empty metadata is fine. + MemorySegment dtype = extensionDType("acme.geopoint", false); + + // When / Then + MemorySegment footer = minimalFooter(); + MemorySegment layout = flatLayout(); + assertThatThrownBy(() -> PostscriptParser.parseBlobs(footer, layout, dtype)) + .isInstanceOf(VortexException.class) + .hasMessageContaining("extension dtype missing metadata"); + } + @Test void convertDType_mapMissingKeyType_throwsVortexException() { // Given — a map dtype table with no key_type field. FlatBuffers makes every field @@ -238,6 +295,32 @@ private static MemorySegment mapDType(boolean withKeyType, boolean withValueType } + /// Builds an extension dtype blob over non-nullable I64 storage with empty metadata; a `null` + /// id omits the field. + private static MemorySegment extensionDType(String id) { + return extensionDType(id, true); + } + + private static MemorySegment extensionDType(String id, boolean withMetadata) { + var fbb = new FbsBuilder(256); + int storage = FbsDType.createFbsDType(fbb, FbsType.FbsPrimitive, + FbsPrimitive.createFbsPrimitive(fbb, FbsPType.I64, false)); + int idOff = id == null ? 0 : fbb.createString(id); + int metaOff = withMetadata ? FbsExtension.createMetadataVector(fbb, new byte[0]) : 0; + FbsExtension.startFbsExtension(fbb); + if (id != null) { + FbsExtension.addId(fbb, idOff); + } + FbsExtension.addStorageDtype(fbb, storage); + if (withMetadata) { + FbsExtension.addMetadata(fbb, metaOff); + } + int ext = FbsExtension.endFbsExtension(fbb); + int root = FbsDType.createFbsDType(fbb, FbsType.FbsExtension, ext); + fbb.finish(root); + return fbb.dataSegment(); + } + private static MemorySegment footerWithArraySpec(String arraySpec) { var fbb = new FbsBuilder(256); int asv = FbsFooter.createArraySpecsVector(fbb, new int[]{ diff --git a/reader/src/test/java/io/github/dfa1/vortex/reader/ReadRegistryTest.java b/reader/src/test/java/io/github/dfa1/vortex/reader/ReadRegistryTest.java index dc7400ddc..73a06328e 100644 --- a/reader/src/test/java/io/github/dfa1/vortex/reader/ReadRegistryTest.java +++ b/reader/src/test/java/io/github/dfa1/vortex/reader/ReadRegistryTest.java @@ -4,7 +4,12 @@ import io.github.dfa1.vortex.reader.array.Array; import io.github.dfa1.vortex.reader.array.UnknownArray; import io.github.dfa1.vortex.core.testing.DTypes; +import io.github.dfa1.vortex.core.model.ColumnName; +import io.github.dfa1.vortex.core.model.DType; import io.github.dfa1.vortex.core.model.EncodingId; +import io.github.dfa1.vortex.core.model.ExtensionId; +import io.github.dfa1.vortex.reader.extension.DateExtensionDecoder; +import io.github.dfa1.vortex.reader.extension.ExtensionDecoder; import io.github.dfa1.vortex.reader.decode.ArrayNode; import io.github.dfa1.vortex.reader.decode.DecodeContext; import io.github.dfa1.vortex.reader.decode.EncodingDecoder; @@ -12,6 +17,7 @@ import java.lang.foreign.Arena; import java.lang.foreign.MemorySegment; +import java.util.List; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -196,9 +202,120 @@ void builder_retainsAllRegisteredDecoders() { assertThat(result.hasDecoder(EncodingId.VORTEX_CONSTANT)).isFalse(); } + @Test + void registerExtension_lookupFindsCustomAndBuiltIn() { + // Given — a Custom id next to the built-ins registerDefaults() adds + ExtensionDecoder custom = extension("acme.geopoint"); + + // When + ReadRegistry result = ReadRegistry.builder().registerDefaults().register(custom).build(); + + // Then + assertThat(result.lookup(ExtensionId.parse("acme.geopoint"))).isSameAs(custom); + assertThat(result.lookup(ExtensionId.VORTEX_DATE)).isSameAs(DateExtensionDecoder.INSTANCE); + } + + @Test + void registerExtension_duplicateId_throws() { + // Given + ReadRegistry.Builder sut = ReadRegistry.builder().register(extension("acme.geopoint")); + ExtensionDecoder duplicate = extension("acme.geopoint"); + + // When / Then + assertThatThrownBy(() -> sut.register(duplicate)) + .isInstanceOf(VortexException.class) + .hasMessageContaining("already registered"); + } + + @Test + void validateExtensions_unregisteredNestedId_throwsUnlessAllowUnknown() { + // Given — the unknown id sits inside a list inside a struct: Rust checks every extension + // dtype it deserializes, not only top-level columns + DType schema = new DType.Struct(List.of(ColumnName.of("c")), + List.of(new DType.List(new DType.Extension("acme.geopoint", DTypes.I32, null, false), false)), false); + ReadRegistry strict = ReadRegistry.loadAll(); + ReadRegistry lenient = ReadRegistry.builder().registerDefaults().allowUnknown().build(); + + // When / Then + assertThatThrownBy(() -> strict.validateExtensions(schema)) + .isInstanceOf(VortexException.class) + .hasMessageContaining("no extension decoder registered for acme.geopoint"); + lenient.validateExtensions(schema); + } + + @Test + void validateExtensions_registeredDecoder_runsItsValidate() { + // Given — a decoder whose validate refuses every dtype + ExtensionDecoder refusing = extension("acme.geopoint"); + org.mockito.BDDMockito.willThrow(new VortexException("bad geopoint")).given(refusing).validate(org.mockito.ArgumentMatchers.any()); + ReadRegistry sut = ReadRegistry.builder().register(refusing).build(); + DType schema = new DType.Struct(List.of(ColumnName.of("c")), + List.of(new DType.Extension("acme.geopoint", DTypes.I32, null, false)), false); + + // When / Then + assertThatThrownBy(() -> sut.validateExtensions(schema)) + .isInstanceOf(VortexException.class) + .hasMessage("bad geopoint"); + } + + @Test + void validateExtensions_noDtype_passes() { + // Given — a file without a dtype blob parses with a null schema (PostscriptParserBlobBoundsTest); + // the walk's switch used to throw NullPointerException on it + ReadRegistry sut = ReadRegistry.loadAll(); + + // When / Then + sut.validateExtensions(null); + } + + @Test + void validateExtensions_validateThrowsRuntimeException_wrapsInVortexException() { + // Given — a third-party validate that trips over malformed metadata from the file + ExtensionDecoder fragile = extension("acme.geopoint"); + IndexOutOfBoundsException cause = new IndexOutOfBoundsException("metadata too short"); + org.mockito.BDDMockito.willThrow(cause).given(fragile).validate(org.mockito.ArgumentMatchers.any()); + ReadRegistry sut = ReadRegistry.builder().register(fragile).build(); + DType schema = new DType.Extension("acme.geopoint", DTypes.I32, null, false); + + // When / Then — reading a file only ever fails with VortexException + assertThatThrownBy(() -> sut.validateExtensions(schema)) + .isInstanceOf(VortexException.class) + .hasMessageContaining("acme.geopoint") + .hasCause(cause); + } + + @Test + void validateExtensions_wellKnownIdWithoutDecoder_passesStrictMode() { + // Given — vortex.json is in the vortex-jni default session, so a strict reader must open it + // (reading the column as its Utf8 storage) though Java registers no decoder for it + DType schema = new DType.Extension("vortex.json", DType.UTF8, MemorySegment.ofArray(new byte[0]), false); + ReadRegistry sut = ReadRegistry.loadAll(); + + // When / Then + sut.validateExtensions(schema); + } + + @Test + void validateExtensions_wellKnownIdWithoutDecoder_stillValidated() { + // Given — no decoder, but Rust's rules for the id still apply + DType schema = new DType.Extension("vortex.json", DType.I32, MemorySegment.ofArray(new byte[0]), false); + ReadRegistry sut = ReadRegistry.loadAll(); + + // When / Then + assertThatThrownBy(() -> sut.validateExtensions(schema)) + .isInstanceOf(VortexException.class) + .hasMessageContaining("storage dtype must be Utf8"); + } + private static EncodingDecoder decoder(EncodingId id) { EncodingDecoder decoder = mock(EncodingDecoder.class); given(decoder.encodingId()).willReturn(id); return decoder; } + + private static ExtensionDecoder extension(String id) { + ExtensionDecoder decoder = mock(ExtensionDecoder.class); + given(decoder.extensionId()).willReturn(ExtensionId.parse(id)); + return decoder; + } } diff --git a/reader/src/test/java/io/github/dfa1/vortex/reader/decode/VariantEncodingDecoderTest.java b/reader/src/test/java/io/github/dfa1/vortex/reader/decode/VariantEncodingDecoderTest.java index 4dbc0cd6a..38742c675 100644 --- a/reader/src/test/java/io/github/dfa1/vortex/reader/decode/VariantEncodingDecoderTest.java +++ b/reader/src/test/java/io/github/dfa1/vortex/reader/decode/VariantEncodingDecoderTest.java @@ -1,6 +1,7 @@ package io.github.dfa1.vortex.reader.decode; import io.github.dfa1.vortex.core.testing.TestSegments; +import io.github.dfa1.vortex.core.error.VortexException; import io.github.dfa1.vortex.reader.ReadRegistry; import io.github.dfa1.vortex.core.model.ColumnName; @@ -287,6 +288,61 @@ void extension_nullMetadata_becomesEmptyBuffer() { assertThat(result.metadata().byteSize()).isZero(); } + @Test + void extension_nullableStorage_extensionIsNullable() { + // Given — the conversion hardcoded nullable=false; Rust's extension nullability is its storage's + var proto = io.github.dfa1.vortex.core.proto.ProtoDType.ofExtension(new io.github.dfa1.vortex.core.proto.ProtoExtension( + "ip.address", prim(io.github.dfa1.vortex.core.proto.ProtoPType.I32, true), new byte[0])); + + // When + DType result = VariantEncodingDecoder.dtypeFromProto(proto); + + // Then + assertThat(result.nullable()).isTrue(); + } + + @Test + void extension_unregisteredId_throwsUnlessAllowUnknown() { + // Given — Rust's DType::from_proto refuses an unregistered id ("Unregistered extension + // type ID") unless the session allows unknown ids + var proto = io.github.dfa1.vortex.core.proto.ProtoDType.ofExtension(new io.github.dfa1.vortex.core.proto.ProtoExtension( + "ip.address", prim(io.github.dfa1.vortex.core.proto.ProtoPType.I32, false), new byte[0])); + ReadRegistry strict = ReadRegistry.loadAll(); + ReadRegistry lenient = ReadRegistry.builder().registerDefaults().allowUnknown().build(); + + // When / Then + assertThatThrownBy(() -> VariantEncodingDecoder.dtypeFromProto(proto, strict)) + .isInstanceOf(VortexException.class) + .hasMessageContaining("no extension decoder registered for ip.address"); + assertThat(VariantEncodingDecoder.dtypeFromProto(proto, lenient)).isInstanceOf(DType.Extension.class); + } + + @Test + void extension_invalidBuiltIn_throws() { + // Given — vortex.timestamp over I32 storage, which Rust's validate_dtype refuses + var proto = io.github.dfa1.vortex.core.proto.ProtoDType.ofExtension(new io.github.dfa1.vortex.core.proto.ProtoExtension( + "vortex.timestamp", prim(io.github.dfa1.vortex.core.proto.ProtoPType.I32, false), new byte[]{2, 0, 0})); + ReadRegistry registry = ReadRegistry.loadAll(); + + // When / Then + assertThatThrownBy(() -> VariantEncodingDecoder.dtypeFromProto(proto, registry)) + .isInstanceOf(VortexException.class) + .hasMessageContaining("storage dtype must be I64"); + } + + @Test + void extension_controlCharacterId_throwsVortexException() { + // Given — the id is untrusted file content; ExtensionId.parse would throw IllegalArgumentException + var proto = io.github.dfa1.vortex.core.proto.ProtoDType.ofExtension(new io.github.dfa1.vortex.core.proto.ProtoExtension( + "a\nb", prim(io.github.dfa1.vortex.core.proto.ProtoPType.I32, false), new byte[0])); + ReadRegistry registry = ReadRegistry.loadAll(); + + // When / Then + assertThatThrownBy(() -> VariantEncodingDecoder.dtypeFromProto(proto, registry)) + .isInstanceOf(VortexException.class) + .hasMessageContaining("invalid extension id"); + } + @Test void variant() { // Given / When diff --git a/reader/src/test/java/io/github/dfa1/vortex/reader/extension/DateExtensionDecoderTest.java b/reader/src/test/java/io/github/dfa1/vortex/reader/extension/DateExtensionDecoderTest.java index 458a1208b..3d20ac8e5 100644 --- a/reader/src/test/java/io/github/dfa1/vortex/reader/extension/DateExtensionDecoderTest.java +++ b/reader/src/test/java/io/github/dfa1/vortex/reader/extension/DateExtensionDecoderTest.java @@ -1,14 +1,20 @@ package io.github.dfa1.vortex.reader.extension; +import io.github.dfa1.vortex.core.error.VortexException; import io.github.dfa1.vortex.core.model.DType; import io.github.dfa1.vortex.core.model.ExtensionId; +import io.github.dfa1.vortex.core.model.TimeUnit; import io.github.dfa1.vortex.reader.array.BoolArray; import io.github.dfa1.vortex.reader.array.IntArray; +import io.github.dfa1.vortex.reader.array.LongArray; import io.github.dfa1.vortex.reader.array.MaskedArray; import io.github.dfa1.vortex.reader.array.MaterializedBoolArray; import io.github.dfa1.vortex.reader.array.MaterializedIntArray; +import io.github.dfa1.vortex.reader.array.MaterializedLongArray; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; import java.lang.foreign.Arena; import java.lang.foreign.MemorySegment; @@ -19,9 +25,13 @@ import static io.github.dfa1.vortex.reader.extension.ExtensionTestSupport.I32; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; class DateExtensionDecoderTest { + private static final DType.Extension DAYS = DateExtensionDecoder.INSTANCE.dtype(false); + private static final DType.Extension MILLIS = dateDtype(TimeUnit.Milliseconds, DType.I64); + private final DateExtensionDecoder sut = DateExtensionDecoder.INSTANCE; @Test @@ -53,7 +63,7 @@ void decode_tpchSample() { IntArray storage = new MaterializedIntArray(I32, 1, buf); // When / Then - assertThat(sut.decode(storage, 0)).isEqualTo(LocalDate.of(1996, 2, 12)); + assertThat(sut.decode(DAYS, storage, 0)).isEqualTo(LocalDate.of(1996, 2, 12)); } } @@ -66,7 +76,7 @@ void decode_negativeDays_returnsPreEpoch() { IntArray storage = new MaterializedIntArray(I32, 1, buf); // When / Then - assertThat(sut.decode(storage, 0)).isEqualTo(LocalDate.of(1969, 12, 31)); + assertThat(sut.decode(DAYS, storage, 0)).isEqualTo(LocalDate.of(1969, 12, 31)); } } @@ -86,7 +96,7 @@ void decodeAll_maskedArray_yieldsNullAtInvalidPositions() { BoolArray validity = new MaterializedBoolArray(DType.BOOL, 3, validityBuf); // When - List out = sut.decodeAll(new MaskedArray(inner, validity)); + List out = sut.decodeAll(DAYS, new MaskedArray(inner, validity)); // Then assertThat(out).containsExactly( @@ -95,4 +105,64 @@ void decodeAll_maskedArray_yieldsNullAtInvalidPositions() { LocalDate.of(2026, 1, 3)); } } + + @ParameterizedTest + @CsvSource({ + // Rust stores an Arrow Date64 as vortex.date(Milliseconds) over I64. The decoder used + // to read the count as days, so these came back ~86 million times too far from 1970. + "0, 1970-01-01", + "86400000, 1970-01-02", + "824083200000, 1996-02-12", + // before the epoch a millisecond count floors to the earlier day, not toward zero + "-1, 1969-12-31", + "-86400000, 1969-12-31", + "-86400001, 1969-12-30" + }) + void decode_milliseconds_floorsToTheDay(long millis, LocalDate expected) { + // Given + try (Arena arena = Arena.ofConfined()) { + MemorySegment buf = arena.allocate(8); + buf.set(ValueLayout.JAVA_LONG_UNALIGNED, 0, millis); + LongArray storage = new MaterializedLongArray(DType.I64, 1, buf); + + // When + LocalDate result = sut.decode(MILLIS, storage, 0); + + // Then + assertThat(result).isEqualTo(expected); + } + } + + @Test + void decodeAll_milliseconds_usesTheUnit() { + // Given — the same unit must reach decodeAll, the path Chunk.as takes + try (Arena arena = Arena.ofConfined()) { + MemorySegment buf = arena.allocate(16); + buf.set(ValueLayout.JAVA_LONG_UNALIGNED, 0, 0L); + buf.set(ValueLayout.JAVA_LONG_UNALIGNED, 8, 824_083_200_000L); + LongArray storage = new MaterializedLongArray(DType.I64, 2, buf); + + // When + List result = sut.decodeAll(MILLIS, storage); + + // Then + assertThat(result).containsExactly(LocalDate.EPOCH, LocalDate.of(1996, 2, 12)); + } + } + + @Test + void date_unsupportedUnit_throws() { + // Given — Rust's Date supports only Days and Milliseconds + DType.Extension seconds = dateDtype(TimeUnit.Seconds, DType.I64); + + // When / Then + assertThatThrownBy(() -> sut.date(seconds, 1L)) + .isInstanceOf(VortexException.class) + .hasMessageContaining("does not support time unit Seconds"); + } + + private static DType.Extension dateDtype(TimeUnit unit, DType storage) { + return new DType.Extension("vortex.date", storage, + MemorySegment.ofArray(new byte[]{(byte) unit.ordinal()}), false); + } } diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/DTypeFbsSerializer.java b/writer/src/main/java/io/github/dfa1/vortex/writer/DTypeFbsSerializer.java index a80f9dfab..d89cf8ace 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/DTypeFbsSerializer.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/DTypeFbsSerializer.java @@ -108,11 +108,10 @@ static int serializeDType(FbsBuilder fbb, DType dtype) { case DType.Extension e -> { int idOff = fbb.createString(e.extensionId()); int storageDtypeOff = serializeDType(fbb, e.storageDType()); - int metaOff = 0; - if (e.metadata() != null) { - byte[] metaBytes = e.metadata().toArray(ValueLayout.JAVA_BYTE); - metaOff = FbsExtension.createMetadataVector(fbb, metaBytes); - } + // Always present, empty when there is none: Rust refuses an extension dtype whose + // metadata field is absent ("failed to parse extension metadata"), failing the file. + byte[] metaBytes = e.metadata() == null ? new byte[0] : e.metadata().toArray(ValueLayout.JAVA_BYTE); + int metaOff = FbsExtension.createMetadataVector(fbb, metaBytes); int inner = FbsExtension.createFbsExtension(fbb, idOff, storageDtypeOff, metaOff); yield FbsDType.createFbsDType(fbb, FbsType.FbsExtension, inner); } diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/ExtensionEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/ExtensionEncoder.java index f38e62988..ea0902684 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/ExtensionEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/ExtensionEncoder.java @@ -5,11 +5,14 @@ import java.util.Collection; -/// Write-side contract for a Vortex extension type. +/// Write-side contract for a Vortex extension type, the Java side of Rust's `ExtVTable`. /// -/// Implementations pair a spec identity ([ExtensionId]) with the matching +/// Implementations pair an identity ([ExtensionId]) with the matching /// [DType.Extension] dtype and a polymorphic [#encodeAll] entry point used -/// by the writer's auto-routing path. +/// by the writer's auto-routing path. Register one on +/// [WriteRegistry.Builder#register(ExtensionEncoder)] and pass the registry to +/// [VortexWriter#create(java.nio.channels.WritableByteChannel, DType.Struct, WriteOptions, WriteRegistry)]. +/// An extension id with no registered encoder can still be written from its storage array. public interface ExtensionEncoder { /// @return the spec identity of this extension @@ -28,4 +31,15 @@ public interface ExtensionEncoder { /// @param values domain-typed values to encode /// @return packed storage array Object encodeAll(DType.Extension dtype, Collection values); + + /// Rejects a dtype this extension cannot represent — Rust's `validate_dtype`. Called for every + /// extension dtype with this id in the schema when the writer is created. + /// + /// The default accepts every dtype. + /// + /// @param dtype the extension dtype declared by the schema + /// @throws io.github.dfa1.vortex.core.error.VortexException if the storage dtype or metadata is + /// invalid for this extension + default void validate(DType.Extension dtype) { + } } diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java b/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java index f2ac8ab90..6a6b2255a 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java @@ -31,11 +31,7 @@ import io.github.dfa1.vortex.writer.encode.BoolEncodingEncoder; import io.github.dfa1.vortex.writer.encode.CascadingCompressor; import io.github.dfa1.vortex.writer.encode.ConstantEncodingEncoder; -import io.github.dfa1.vortex.writer.encode.DateExtensionEncoder; import io.github.dfa1.vortex.writer.encode.DateTimePartsEncodingEncoder; -import io.github.dfa1.vortex.writer.encode.TimeExtensionEncoder; -import io.github.dfa1.vortex.writer.encode.TimestampExtensionEncoder; -import io.github.dfa1.vortex.writer.encode.UuidExtensionEncoder; import io.github.dfa1.vortex.writer.encode.DecimalBytePartsEncodingEncoder; import io.github.dfa1.vortex.writer.encode.DecimalEncodingEncoder; import io.github.dfa1.vortex.writer.encode.DeltaEncodingEncoder; @@ -123,6 +119,8 @@ public final class VortexWriter implements Closeable { new DictEncodingEncoder(), new VarBinEncodingEncoder(), new ExtEncodingEncoder(), new FixedSizeListEncodingEncoder(), new ListEncodingEncoder(), new DecimalEncodingEncoder()); + private static final WriteRegistry DEFAULT_EXTENSIONS = WriteRegistry.builder().registerDefaultExtensions().build(); + // Base cascade codec list — no Zstd or Pco. They are appended (before PrimitiveEncoding) when // WriteOptions.compact() is true. See WriteOptions.withCompact(boolean) for the tradeoff. @@ -131,6 +129,8 @@ public final class VortexWriter implements Closeable { private final WriteOptions options; private final List encodings; private final WriteRegistry defaultRegistry; + // Extension encoders for the Collection auto-route: the built-ins, or the caller's registry. + private final WriteRegistry extensionRegistry; private final List cascadeCodecs; private final WriteRegistry cascadeRegistry; private final List segs = new ArrayList<>(); @@ -193,7 +193,8 @@ public final class VortexWriter implements Closeable { private int submittedSegments = 0; private VortexWriter( - WritableByteChannel channel, DType.Struct schema, WriteOptions options, List encodings + WritableByteChannel channel, DType.Struct schema, WriteOptions options, List encodings, + WriteRegistry extensionRegistry ) { // Wire contract, enforced by the reference writer ("StructLayout must have unique field // names"): duplicate-name schemas are constructible via the DType.Struct record (only @@ -213,6 +214,8 @@ private VortexWriter( this.dictRetainedBudget = options.globalDictMaxRetainedBytes().bytes(); this.encodings = encodings; this.defaultRegistry = buildRegistry(encodings); + this.extensionRegistry = extensionRegistry; + validateExtensions(schema); // A custom encoder list is the whole candidate set, cascade included: emitting a built-in // the caller left out would produce a file their own read registry may not decode. this.cascadeCodecs = encodings == DEFAULT_CODECS ? buildCascadeCodecs(options) : List.copyOf(encodings); @@ -287,16 +290,43 @@ private static Set editionExcluded( return Set.copyOf(excluded); } - /// Builds a [WriteRegistry] from the given encoder list plus all built-in extension encoders. + /// Runs [ExtensionEncoder#validate(DType.Extension)] on every extension dtype in `dtype`, nested + /// ones included, that has a registered encoder — Rust validates each extension dtype it builds. + private void validateExtensions(DType dtype) { + switch (dtype) { + case DType.Extension ext -> { + io.github.dfa1.vortex.core.model.ExtensionId id = + io.github.dfa1.vortex.core.model.ExtensionId.parse(ext.extensionId()); + ExtensionEncoder encoder = extensionRegistry.lookup(id); + if (encoder != null) { + encoder.validate(ext); + } else if (id instanceof io.github.dfa1.vortex.core.model.ExtensionId.WellKnown wellKnown) { + // written from its storage array; still held to Rust's rules for the id + wellKnown.validate(ext); + } + validateExtensions(ext.storageDType()); + } + case DType.Struct struct -> struct.fieldTypes().forEach(this::validateExtensions); + case DType.List list -> validateExtensions(list.elementType()); + case DType.FixedSizeList list -> validateExtensions(list.elementType()); + case DType.Map map -> { + validateExtensions(map.keyType()); + validateExtensions(map.valueType()); + } + case DType.Union union -> union.variantTypes().forEach(this::validateExtensions); + case DType.Null _, DType.Bool _, DType.Primitive _, DType.Decimal _, DType.Utf8 _, + DType.Binary _, DType.Variant _ -> { + // no nested dtype + } + } + } + + /// Builds a [WriteRegistry] from the given encoder list. private static WriteRegistry buildRegistry(List encoders) { WriteRegistry.Builder b = WriteRegistry.builder(); for (EncodingEncoder e : encoders) { b.register(e); } - b.register(DateExtensionEncoder.INSTANCE) - .register(TimeExtensionEncoder.INSTANCE) - .register(TimestampExtensionEncoder.INSTANCE) - .register(UuidExtensionEncoder.INSTANCE); return b.build(); } @@ -378,7 +408,7 @@ private static List buildCascadeCodecs(WriteOptions options) { public static VortexWriter create( WritableByteChannel channel, DType.Struct schema, WriteOptions options ) { - return new VortexWriter(channel, schema, options, DEFAULT_CODECS); + return new VortexWriter(channel, schema, options, DEFAULT_CODECS, DEFAULT_EXTENSIONS); } /// Creates a [VortexWriter] with a custom encoder list. @@ -395,7 +425,7 @@ public static VortexWriter create( ) { // Custom encoding list: disable global dict — using DEFAULT_CODECS for values/codes behind the scenes // would violate the user's expectation that only their encoding list is used. - return new VortexWriter(channel, schema, options.withGlobalDict(false), encodings); + return new VortexWriter(channel, schema, options.withGlobalDict(false), encodings, DEFAULT_EXTENSIONS); } /// Creates a [VortexWriter] with a custom [WriteRegistry]. @@ -409,7 +439,7 @@ public static VortexWriter create( WritableByteChannel channel, DType.Struct schema, WriteOptions options, WriteRegistry registry ) { return new VortexWriter(channel, schema, options.withGlobalDict(false), - List.copyOf(registry.encoderMap().values())); + List.copyOf(registry.encoderMap().values()), registry); } /// Counts rows for the length-consistency check in [#writeChunk]. Accepts the @@ -555,14 +585,16 @@ public void writeChunk(Map columns) throws IOException { // ExtEncoding wraps the storage child below — matches Rust's nested layout // (ExtEncoding → PrimitiveEncoding) and lets Registry skip its unwrap path. if (colDtype instanceof DType.Extension extDtype && data instanceof java.util.Collection coll) { - ExtensionEncoder impl = - io.github.dfa1.vortex.core.model.ExtensionId.parse(extDtype.extensionId()) - .map(defaultRegistry::lookup) - .orElse(null); - if (impl != null) { - data = impl.encodeAll(extDtype, coll); - adapted.put(colName, data); + ExtensionEncoder impl = extensionRegistry.lookup( + io.github.dfa1.vortex.core.model.ExtensionId.parse(extDtype.extensionId())); + if (impl == null) { + // Without an encoder the Collection would reach the encoders unconverted and + // fail far from the cause; raw storage arrays still pass through. + throw new VortexException("column '" + colName + "': no extension encoder registered for " + + extDtype.extensionId() + "; register one or pass the storage array"); } + data = impl.encodeAll(extDtype, coll); + adapted.put(colName, data); } // Sizing a string batch reads every string and was the caller thread's largest cost diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/WriteRegistry.java b/writer/src/main/java/io/github/dfa1/vortex/writer/WriteRegistry.java index eaaed0fe9..c76aff932 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/WriteRegistry.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/WriteRegistry.java @@ -68,8 +68,8 @@ public final class WriteRegistry { private WriteRegistry(Map encoders, Map extensions) { - // Order by encoding name so it is stable regardless of enum declaration order. (The id - // enums sort by ordinal naturally — Enum.compareTo is final — so a Comparator is required.) + // Order by id string so it is stable regardless of declaration order. (The id types are not + // Comparable — a WellKnown enum would sort by ordinal, a Custom record not at all.) // VortexWriter.create(.., WriteRegistry) selects the first encoder whose accepts() matches // the dtype, so iteration order is significant: a HashMap would make selection vary across // runs and platforms. Sorting by id makes the order a pure function of the registered set, @@ -128,7 +128,8 @@ public static final class Builder { // EncodingId is not Comparable (a Custom id would throw on natural ordering); order by // wire string, matching the constructor's sortedByName. private final Map encoders = new TreeMap<>(Comparator.comparing(EncodingId::id)); - private final Map extensions = new TreeMap<>(); + // Same as encoders: ExtensionId is not Comparable, so a Custom key would throw on natural ordering. + private final Map extensions = new TreeMap<>(Comparator.comparing(ExtensionId::id)); private Builder() { } @@ -202,11 +203,16 @@ public Builder registerDefaults() { .register(new VarBinViewEncodingEncoder()) .register(new ZigZagEncodingEncoder()) .register(new ZstdEncodingEncoder()); - register(DateExtensionEncoder.INSTANCE) + return registerDefaultExtensions(); + } + + /// Registers the built-in [ExtensionEncoder]s only — also the writer's extension set when + /// it is created without a [WriteRegistry]. + Builder registerDefaultExtensions() { + return register(DateExtensionEncoder.INSTANCE) .register(TimeExtensionEncoder.INSTANCE) .register(TimestampExtensionEncoder.INSTANCE) .register(UuidExtensionEncoder.INSTANCE); - return this; } /// Builds an immutable [WriteRegistry]. diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/DateExtensionEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/DateExtensionEncoder.java index 02c434e11..66c669e56 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/DateExtensionEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/DateExtensionEncoder.java @@ -23,6 +23,12 @@ public ExtensionId extensionId() { return ExtensionId.VORTEX_DATE; } + /// Applies Rust's `validate_dtype` rules for `vortex.date`. + @Override + public void validate(DType.Extension dtype) { + ExtensionId.VORTEX_DATE.validate(dtype); + } + @Override public DType.Extension dtype(boolean nullable) { // Rust vortex.date metadata: 1 byte = TimeUnit tag (Days = 4), required by Rust reader. diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/TimeExtensionEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/TimeExtensionEncoder.java index 44f00c3d8..50c561aa1 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/TimeExtensionEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/TimeExtensionEncoder.java @@ -22,6 +22,12 @@ public ExtensionId extensionId() { return ExtensionId.VORTEX_TIME; } + /// Applies Rust's `validate_dtype` rules for `vortex.time`. + @Override + public void validate(DType.Extension dtype) { + ExtensionId.VORTEX_TIME.validate(dtype); + } + /// Returns the default dtype using [TimeUnit#Milliseconds] over I32 storage. /// Use [#dtype(TimeUnit, boolean)] for non-default units. @Override diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/TimestampExtensionEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/TimestampExtensionEncoder.java index 0f90cbff8..1f05e8591 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/TimestampExtensionEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/TimestampExtensionEncoder.java @@ -23,6 +23,12 @@ public ExtensionId extensionId() { return ExtensionId.VORTEX_TIMESTAMP; } + /// Applies Rust's `validate_dtype` rules for `vortex.timestamp`. + @Override + public void validate(DType.Extension dtype) { + ExtensionId.VORTEX_TIMESTAMP.validate(dtype); + } + /// Returns the default dtype using milliseconds resolution and no timezone. /// Use [#dtype(TimeUnit, ZoneId, boolean)] for non-default settings. @Override diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/UuidExtensionEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/UuidExtensionEncoder.java index 6f4f8e328..ff14ebf6b 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/UuidExtensionEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/UuidExtensionEncoder.java @@ -21,6 +21,12 @@ public ExtensionId extensionId() { return ExtensionId.VORTEX_UUID; } + /// Applies Rust's `validate_dtype` rules for `vortex.uuid`. + @Override + public void validate(DType.Extension dtype) { + ExtensionId.VORTEX_UUID.validate(dtype); + } + @Override public DType.Extension dtype(boolean nullable) { DType.Primitive u8 = DType.U8; diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/VortexWriterTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/VortexWriterTest.java index 4005464f4..2c24b0759 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/VortexWriterTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/VortexWriterTest.java @@ -178,6 +178,72 @@ void writeChunk_autoroutesExtensionCollectionViaSpecExtension(@TempDir Path tmp) } } + @Test + void create_withRegistry_usesItsExtensionEncoders(@TempDir Path tmp) throws IOException { + // Given — a registry whose vortex.date encoder takes epoch-day Integers instead of + // LocalDates. create(WriteRegistry) used to rebuild the extension set from the four + // built-ins, silently dropping this one: the built-in then cast an Integer to LocalDate. + var dateSchema = new DType.Struct( + List.of(ColumnName.of("d")), + List.of(io.github.dfa1.vortex.writer.encode.DateExtensionEncoder.INSTANCE.dtype(false)), + false); + ExtensionEncoder epochDays = new ExtensionEncoder() { + @Override + public io.github.dfa1.vortex.core.model.ExtensionId extensionId() { + return io.github.dfa1.vortex.core.model.ExtensionId.VORTEX_DATE; + } + + @Override + public DType.Extension dtype(boolean nullable) { + return io.github.dfa1.vortex.writer.encode.DateExtensionEncoder.INSTANCE.dtype(nullable); + } + + @Override + public Object encodeAll(DType.Extension dtype, java.util.Collection values) { + return values.stream().mapToInt(v -> (Integer) v).toArray(); + } + }; + WriteRegistry.Builder builder = WriteRegistry.builder(); + WriteRegistry.loadAll().encoderMap().values().forEach(builder::register); + WriteRegistry registry = builder.register(epochDays).build(); + Path file = tmp.resolve("custom_date_encoder.vtx"); + + // When + try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE); + var sut = VortexWriter.create(ch, dateSchema, WriteOptions.defaults(), registry)) { + sut.writeChunk(Map.of(ColumnName.of("d"), List.of(0, 20_000))); + } + + // Then + try (var vf = VortexReader.open(file, ReadRegistry.loadAll()); + var iter = vf.scan(ScanOptions.all()); + Chunk chunk = iter.next()) { + assertThat(chunk.as(ColumnName.of("d"), java.time.LocalDate.class)) + .containsExactly(java.time.LocalDate.ofEpochDay(0), java.time.LocalDate.ofEpochDay(20_000)); + } + } + + @Test + void writeChunk_extensionCollectionWithoutEncoder_failsFast(@TempDir Path tmp) throws IOException { + // Given — a registry with every encoder but no extension encoders. A Collection used to + // pass through unconverted and fail later with an unrelated encoder error. + var dateSchema = new DType.Struct( + List.of(ColumnName.of("d")), + List.of(io.github.dfa1.vortex.writer.encode.DateExtensionEncoder.INSTANCE.dtype(false)), + false); + WriteRegistry.Builder builder = WriteRegistry.builder(); + WriteRegistry.loadAll().encoderMap().values().forEach(builder::register); + Path file = tmp.resolve("no_ext_encoder.vtx"); + + try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE); + var sut = VortexWriter.create(ch, dateSchema, WriteOptions.defaults(), builder.build())) { + // When / Then + assertThatThrownBy(() -> sut.writeChunk(Map.of(ColumnName.of("d"), List.of(java.time.LocalDate.EPOCH)))) + .isInstanceOf(io.github.dfa1.vortex.core.error.VortexException.class) + .hasMessageContaining("no extension encoder registered for vortex.date"); + } + } + @Test void writeChunk_extensionCollectionColumn_rowCountValidatedAgainstSibling(@TempDir Path tmp) throws IOException { diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/WriteRegistryTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/WriteRegistryTest.java index 24bd9362c..a32927b06 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/WriteRegistryTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/WriteRegistryTest.java @@ -103,6 +103,21 @@ void registerExtension_duplicateId_throws() { .hasMessageContaining("already registered"); } + @Test + void registerExtension_customAndWellKnownIds_bothLookUp() { + // Given — a Custom id next to a WellKnown one. The builder's TreeMap used natural ordering, + // which a Custom record does not have: the second put threw ClassCastException. + ExtensionEncoder custom = extension(new ExtensionId.Custom("acme.geopoint")); + ExtensionEncoder date = extension(ExtensionId.VORTEX_DATE); + + // When + WriteRegistry result = WriteRegistry.builder().register(date).register(custom).build(); + + // Then + assertThat(result.lookup(ExtensionId.parse("acme.geopoint"))).isSameAs(custom); + assertThat(result.lookup(ExtensionId.VORTEX_DATE)).isSameAs(date); + } + @Test void lookup_unregisteredExtension_returnsNull() { // Given — empty registry diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/WriterZoneMapTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/WriterZoneMapTest.java index f69af68cd..0cc59e019 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/WriterZoneMapTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/WriterZoneMapTest.java @@ -15,6 +15,7 @@ import io.github.dfa1.vortex.reader.ScanIterator; import io.github.dfa1.vortex.reader.ScanOptions; import io.github.dfa1.vortex.reader.SegmentSpec; +import io.github.dfa1.vortex.reader.ReadRegistry; import io.github.dfa1.vortex.reader.VortexReader; import io.github.dfa1.vortex.reader.array.DoubleArray; import io.github.dfa1.vortex.reader.array.LongArray; @@ -308,8 +309,10 @@ void extensionColumn_emitsStorageMinMaxZoneMap(@TempDir Path tmp) throws IOExcep sut.writeChunk(Map.of(ColumnName.of("t"), new long[]{20, 21})); } - // When / Then the column is zoned with MAX+MIN+NULL_COUNT and the stats decode per zone - try (VortexReader reader = VortexReader.open(file)) { + // When / Then the column is zoned with MAX+MIN+NULL_COUNT and the stats decode per zone; + // test.ext has no decoder, so the file opens only with unknown ids allowed (Rust parity) + ReadRegistry allowUnknown = ReadRegistry.builder().registerDefaults().allowUnknown().build(); + try (VortexReader reader = VortexReader.open(file, allowUnknown)) { Layout column = reader.layout().children().get(0); assertThat(column.isZoned()).isTrue(); MemorySegment meta = column.metadata();