From 14cb6be4502fce2e2049087fba5e3d3e291d9570 Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Thu, 22 Oct 2026 22:04:49 +0200 Subject: [PATCH 1/4] docs: no ternary that wraps lines or binds a pattern Co-Authored-By: Claude Sonnet 5 --- CLAUDE.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/CLAUDE.md b/CLAUDE.md index b8a5d60d8..42c788a3b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -366,6 +366,9 @@ procedure): `forEachChunk`; use `Optional`, records, sealed types, pattern switches, virtual threads, FFM). New APIs should feel like JDK APIs. - Always braces for `if`/`else`/`for`/`while`, even one-liners (`if (c) { return a; }`). +- **No ternary that wraps or binds.** A `?:` that spans lines, or whose condition binds a pattern + (`x instanceof Foo f && … ? f.a() : x`), is an `if`: assign the default, then override it. + A short one-line ternary (`n == 0 ? 1 : n`) is fine. - **Enums: never call `values()` on a call path.** `values()` allocates a fresh array on every call (64 B for `PType`), which the JIT does not remove. Look a constant up by wire tag, ordinal or number with a `switch` expression: no allocation, and about 4x faster than `values()[i]` when the ordinals repeat, as in real files From c39c95b72dbb61e1ddca8502bb272d96f206a4fa Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Thu, 22 Oct 2026 22:13:04 +0200 Subject: [PATCH 2/4] fix(writer): order Utf8 dict bounds by UTF-8 bytes DictEncodingEncoder picked a string column's min/max with String::compareTo, UTF-16 order, where every other string bound in the writer uses UTF-8 byte order (Utf8Order), as Rust compares strings. A dictionary mixing a supplementary character (a surrogate pair in UTF-16) with U+E000..U+FFFF reported the bounds the wrong way round. The default vortex.zoned layout computes its own bounds, so this reached files only through the encoder stats (the legacy vortex.stats layout). Co-Authored-By: Claude Sonnet 5 --- .../writer/encode/DictEncodingEncoder.java | 16 ++++++++-------- .../writer/encode/DictEncodingEncoderTest.java | 15 +++++++++++++++ 2 files changed, 23 insertions(+), 8 deletions(-) diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoder.java index 534f855d4..5a76f260f 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoder.java @@ -217,10 +217,10 @@ private static CascadeStep encodeUtf8Cascade(String[] strings, EncodeContext ctx String[] distinct = valueMap.keySet().toArray(new String[0]); - String minStr = valueMap.keySet().stream().min(String::compareTo).orElse(null); - String maxStr = valueMap.keySet().stream().max(String::compareTo).orElse(null); - ProtoScalarValue statsMin = minStr != null ? ProtoScalarValue.ofStringValue(minStr) : null; - ProtoScalarValue statsMax = maxStr != null ? ProtoScalarValue.ofStringValue(maxStr) : null; + // UTF-8 byte order like every other string bound (Rust compares bytes), not String's UTF-16 order + MinMax stats = VarBinEncodingEncoder.minMaxStats(valueMap.keySet().toArray(new String[0])); + ProtoScalarValue statsMin = stats != null ? stats.min() : null; + ProtoScalarValue statsMax = stats != null ? stats.max() : null; return new CascadeStep(partialRoot, List.of(), List.of(new ChildSlot(new DType.Primitive(codePType, false), codesArr, 0, CODES_EXCLUDED), @@ -356,10 +356,10 @@ private static EncodeResult encodeUtf8(String[] strings, EncodeContext ctx) { new EncodeNode[]{codesNode, valuesNode}, new int[0]); - String minStr = valueMap.keySet().stream().min(String::compareTo).orElse(null); - String maxStr = valueMap.keySet().stream().max(String::compareTo).orElse(null); - ProtoScalarValue statsMin = minStr != null ? ProtoScalarValue.ofStringValue(minStr) : null; - ProtoScalarValue statsMax = maxStr != null ? ProtoScalarValue.ofStringValue(maxStr) : null; + // UTF-8 byte order like every other string bound (Rust compares bytes), not String's UTF-16 order + MinMax stats = VarBinEncodingEncoder.minMaxStats(valueMap.keySet().toArray(new String[0])); + ProtoScalarValue statsMin = stats != null ? stats.min() : null; + ProtoScalarValue statsMax = stats != null ? stats.max() : null; return new EncodeResult(root, List.of(EncodedBuffer.bytes(dictBytesBuf), EncodedBuffer.of(dictOffsetsBuf, PType.I64), EncodedBuffer.of(codesBuf, codePType)), statsMin, statsMax); } diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoderTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoderTest.java index fa9977b24..d8b6a5aea 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoderTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoderTest.java @@ -206,6 +206,21 @@ void encodeCascade_utf8_exposesCodesAndValuesAsOpenChildren() { assertThat((String[]) values.childData()).containsExactly("apple", "banana", "cherry"); } + @Test + void encodeCascade_utf8_statsUseUtf8ByteOrder() { + // Given — U+FF61 (UTF-8 EF BD A1) and U+1F600 (F0 9F 98 80, a surrogate pair in UTF-16): + // UTF-16 order puts the emoji first, UTF-8 order (Rust's, and every other string bound's) + // puts it last. Dict used String::compareTo, so its bounds disagreed with VarBin's. + String[] data = repeat(new String[]{"😀", "｡"}, 100); + + // When + CascadeStep result = ENCODER.encodeCascade(DTypes.UTF8, data, EncodeTestHelper.testCtx()); + + // Then + assertThat(result.statsMin().string_value()).isEqualTo("｡"); + assertThat(result.statsMax().string_value()).isEqualTo("😀"); + } + @Test void encode_primitive_statsCarryMinAndMaxOfLogicalValues() throws IOException { // Given — low-cardinality so dict encoding is a plausible pick; unordered so a broken From 1e29bc23f3a4f076aa0eb99215b46a42a62e47a2 Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Thu, 22 Oct 2026 22:56:17 +0200 Subject: [PATCH 3/4] feat(writer): accept a Utf8 column as UTF-8 bytes writeChunk takes a Utf8 column as byte[][] as well as String[]: a caller that already holds UTF-8 (a CSV parser) skips decoding to String only for every encoder to encode it back. Rows are checked to be well-formed UTF-8 (Utf8Validation, RFC 3629, as Rust validates a Utf8 array). The encoders told Utf8 from Binary by carrier (String[] vs byte[][]); they now ask the dtype, through VarBinEncodingEncoder.textMinMax for the bounds, and Dict and Constant key and compare the bytes directly. The global dictionary keys Utf8 by String either way, so a column mixing both carriers codes each value once. Utf8BytesInputTest writes 5 column shapes x 3 option sets both ways and asserts byte-identical files. Co-Authored-By: Claude Sonnet 5 --- CHANGELOG.md | 3 + CLAUDE.md | 2 +- .../vortex/core/compute/Utf8Validation.java | 61 +++++++++ .../core/compute/Utf8ValidationTest.java | 95 ++++++++++++++ docs/reference.md | 5 +- .../io/github/dfa1/vortex/writer/Chunk.java | 2 +- .../github/dfa1/vortex/writer/ChunkImpl.java | 20 ++- .../dfa1/vortex/writer/DictColumnState.java | 14 +- .../dfa1/vortex/writer/ZoneMapStatCodec.java | 2 +- .../encode/ConstantEncodingEncoder.java | 41 ++++-- .../writer/encode/DictEncodingEncoder.java | 116 ++++++++++++----- .../writer/encode/FsstEncodingEncoder.java | 24 +--- .../writer/encode/MaskedEncodingEncoder.java | 2 +- .../writer/encode/OnPairEncodingEncoder.java | 7 +- .../writer/encode/VarBinEncodingEncoder.java | 78 +++++++++-- .../vortex/writer/encode/ZoneMapStats.java | 12 +- .../vortex/writer/Utf8BytesInputTest.java | 123 ++++++++++++++++++ 17 files changed, 510 insertions(+), 97 deletions(-) create mode 100644 core/src/main/java/io/github/dfa1/vortex/core/compute/Utf8Validation.java create mode 100644 core/src/test/java/io/github/dfa1/vortex/core/compute/Utf8ValidationTest.java create mode 100644 writer/src/test/java/io/github/dfa1/vortex/writer/Utf8BytesInputTest.java diff --git a/CHANGELOG.md b/CHANGELOG.md index e58c24c69..0843e4ec9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- `VortexWriter` accepts a `Utf8` column as its UTF-8 bytes (`byte[][]`), skipping the per-string encoding; malformed UTF-8 is rejected. ([#590](https://github.com/dfa1/vortex-java/pull/590)) + ### Changed - **Breaking:** SQL identifiers are quoted with double quotes, as IntelliJ and DBeaver generate them, instead of back-ticks, so opening a table or sorting its grid works; write `"date"` where you wrote `` `date` ``. ([#588](https://github.com/dfa1/vortex-java/pull/588)) - **Breaking:** the `ScanIterator` constructor and the reader's Pco helpers (`PcoTansDecoder`, `PcoBin`, `LeBitReader`) are no longer public, and the built-in extension codecs are reached only through their `INSTANCE`. ([4ba6ea12](https://github.com/dfa1/vortex-java/commit/4ba6ea12)) diff --git a/CLAUDE.md b/CLAUDE.md index 42c788a3b..9c9b9f434 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -38,7 +38,7 @@ core — everything lives under `io.github.dfa1.vortex.core.*`: EditionId, Edition, EditionFamily, Editions, MemorySize core.io IoBounds, PTypeIO, VortexFormat core.error VortexException - core.compute FastLanes, PrimitiveArrays, Utf8Order + core.compute FastLanes, PrimitiveArrays, Utf8Order, Utf8Validation core.simd SimdOperations (kernels shared by reader+writer, all eleven ptypes), AutoVectorizedSimdOperations (C2-shaped loops, the Rust-parity default), VectorApiSimdOperations (opt-in via --add-modules jdk.incubator.vector; Rust-parity results, diff --git a/core/src/main/java/io/github/dfa1/vortex/core/compute/Utf8Validation.java b/core/src/main/java/io/github/dfa1/vortex/core/compute/Utf8Validation.java new file mode 100644 index 000000000..b35cc0cd1 --- /dev/null +++ b/core/src/main/java/io/github/dfa1/vortex/core/compute/Utf8Validation.java @@ -0,0 +1,61 @@ +package io.github.dfa1.vortex.core.compute; + +/// Checks that bytes are well-formed UTF-8, as a Utf8 array's bytes must be. +/// +/// Rust validates a Utf8 array's bytes when it builds the array, so a writer that takes UTF-8 +/// bytes from its caller must reject what Rust would. Well-formed is RFC 3629: no overlong forms, +/// no surrogates (U+D800..U+DFFF), nothing above U+10FFFF. +public final class Utf8Validation { + + private Utf8Validation() { + } + + /// Whether `bytes[from, to)` is well-formed UTF-8. + /// + /// @param bytes the bytes + /// @param from the first byte, inclusive + /// @param to the last byte, exclusive + /// @return `true` if the range is well-formed UTF-8 + public static boolean isWellFormed(byte[] bytes, int from, int to) { + int i = from; + while (i < to) { + int lead = bytes[i] & 0xFF; + if (lead < 0x80) { + i++; + continue; + } + int length; + if (lead >= 0xC2 && lead <= 0xDF) { + length = 2; + } else if (lead >= 0xE0 && lead <= 0xEF) { + length = 3; + } else if (lead >= 0xF0 && lead <= 0xF4) { + length = 4; + } else { + return false; + } + if (to - i < length || !secondByteFits(lead, bytes[i + 1] & 0xFF)) { + return false; + } + for (int k = 2; k < length; k++) { + if ((bytes[i + k] & 0xC0) != 0x80) { + return false; + } + } + i += length; + } + return true; + } + + /// The second byte's range narrows after the leads that would otherwise allow an overlong + /// form (E0, F0), a surrogate (ED) or a code point above U+10FFFF (F4): RFC 3629 table 3.7. + private static boolean secondByteFits(int lead, int second) { + return switch (lead) { + case 0xE0 -> second >= 0xA0 && second <= 0xBF; + case 0xED -> second >= 0x80 && second <= 0x9F; + case 0xF0 -> second >= 0x90 && second <= 0xBF; + case 0xF4 -> second >= 0x80 && second <= 0x8F; + default -> second >= 0x80 && second <= 0xBF; + }; + } +} diff --git a/core/src/test/java/io/github/dfa1/vortex/core/compute/Utf8ValidationTest.java b/core/src/test/java/io/github/dfa1/vortex/core/compute/Utf8ValidationTest.java new file mode 100644 index 000000000..94328d95a --- /dev/null +++ b/core/src/test/java/io/github/dfa1/vortex/core/compute/Utf8ValidationTest.java @@ -0,0 +1,95 @@ +package io.github.dfa1.vortex.core.compute; + +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.ValueSource; + +import java.nio.ByteBuffer; +import java.nio.charset.CharacterCodingException; +import java.nio.charset.CodingErrorAction; +import java.nio.charset.StandardCharsets; +import java.util.HexFormat; +import java.util.Random; + +import static org.assertj.core.api.Assertions.assertThat; + +class Utf8ValidationTest { + + @ParameterizedTest + @CsvSource({ + "'', true", + "61, true", // a + "c3a9, true", // é + "e282ac, true", // € + "f09f9880, true", // U+1F600 + "f48fbfbf, true", // U+10FFFF, the last code point + "efbfbd, true", // U+FFFD + "c080, false", // overlong NUL + "e08080, false", // overlong 3-byte + "f0808080, false", // overlong 4-byte + "eda080, false", // U+D800, a surrogate + "f4908080, false", // above U+10FFFF + "f5808080, false", // lead byte past F4 + "80, false", // a stray continuation byte + "c3, false", // truncated + "e282, false", // truncated + "c328, false", // second byte is not a continuation + "e28228, false", // third byte is not a continuation + }) + void isWellFormed_namedCases(String hex, boolean expected) { + // Given + byte[] bytes = HexFormat.of().parseHex(hex); + + // When + boolean result = Utf8Validation.isWellFormed(bytes, 0, bytes.length); + + // Then + assertThat(result).isEqualTo(expected); + } + + @ParameterizedTest + @ValueSource(longs = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10}) + void isWellFormed_agreesWithTheJdkDecoder(long seed) { + // Given: short random byte runs biased toward lead and continuation bytes, so most are + // near-misses rather than plain ASCII; the JDK's strict decoder is the oracle + Random random = new Random(seed); + int[] palette = {0x41, 0x80, 0x8f, 0x90, 0x9f, 0xa0, 0xbf, 0xc2, 0xdf, 0xe0, 0xed, 0xef, 0xf0, 0xf4, 0xf5, 0xff}; + for (int run = 0; run < 2_000; run++) { + byte[] bytes = new byte[random.nextInt(6)]; + for (int i = 0; i < bytes.length; i++) { + bytes[i] = (byte) palette[random.nextInt(palette.length)]; + } + + // When + boolean result = Utf8Validation.isWellFormed(bytes, 0, bytes.length); + + // Then + assertThat(result).as(HexFormat.of().formatHex(bytes)).isEqualTo(jdkAccepts(bytes)); + } + } + + @ParameterizedTest + @CsvSource({"0, 1, true", "1, 3, true", "0, 2, false", "2, 3, false"}) + void isWellFormed_checksOnlyTheRange(int from, int to, boolean expected) { + // Given: 'a' then 'é' (c3 a9): a range that cuts the 2-byte character is malformed + byte[] bytes = {0x61, (byte) 0xc3, (byte) 0xa9}; + + // When + boolean result = Utf8Validation.isWellFormed(bytes, from, to); + + // Then + assertThat(result).isEqualTo(expected); + } + + private static boolean jdkAccepts(byte[] bytes) { + try { + StandardCharsets.UTF_8.newDecoder() + .onMalformedInput(CodingErrorAction.REPORT) + .onUnmappableCharacter(CodingErrorAction.REPORT) + .decode(ByteBuffer.wrap(bytes)); + return true; + } catch (CharacterCodingException _) { + return false; + } + } +} diff --git a/docs/reference.md b/docs/reference.md index 80d81b4c4..082af8887 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -170,10 +170,13 @@ Accepted array types per column `DType`: | `Primitive(I64/U64)` | `long[]` | `Long[]` | | `Primitive(F32)` | `float[]` | `Float[]` | | `Primitive(F64)` | `double[]` | `Double[]` | -| `Utf8` | `String[]` | `String[]` (nulls allowed) | +| `Utf8` | `String[]` or UTF-8 `byte[][]` | either (nulls allowed) | | `Bool` | `boolean[]` | `Boolean[]` | | `Decimal(precision, scale)` | `BigDecimal[]` | `BigDecimal[]` (nulls allowed) | +A `Utf8` column given as UTF-8 `byte[][]` skips encoding each string and writes the same file as the +`String[]` would; each row must be well-formed UTF-8 (as Rust requires), or `IllegalArgumentException`. + Decimal values are rescaled to the column scale exactly: a value with more fractional digits than the scale, or more digits than the precision, is rejected with `IllegalArgumentException` rather than rounded. Every precision Rust supports (1–76, i8 … i256 storage) and negative scales work. diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/Chunk.java b/writer/src/main/java/io/github/dfa1/vortex/writer/Chunk.java index 01e01993e..8ce1a97b4 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/Chunk.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/Chunk.java @@ -22,7 +22,7 @@ /// | Primitive(I64/U64) | `long[]` | `Long[]` | /// | Primitive(F32) | `float[]` | `Float[]` | /// | Primitive(F64) | `double[]` | `Double[]` | -/// | Utf8 | `String[]` | `String[]` (nulls allowed) | +/// | Utf8 | `String[]`, or the UTF-8 bytes as `byte[][]` | either, nulls allowed | /// | Bool | `boolean[]` | `Boolean[]` | public interface Chunk { diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/ChunkImpl.java b/writer/src/main/java/io/github/dfa1/vortex/writer/ChunkImpl.java index 978c7eac5..183d5050a 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/ChunkImpl.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/ChunkImpl.java @@ -1,5 +1,6 @@ package io.github.dfa1.vortex.writer; +import io.github.dfa1.vortex.core.compute.Utf8Validation; import io.github.dfa1.vortex.core.model.ColumnName; import io.github.dfa1.vortex.core.model.DType; import io.github.dfa1.vortex.core.model.PType; @@ -120,15 +121,24 @@ private static Object adaptPrimitive(String column, DType.Primitive dtype, Objec }; } + /// Utf8 takes `String[]`, or the strings' UTF-8 bytes as `byte[][]`, which skips encoding each + /// string and is checked to be well-formed UTF-8, as Rust checks a Utf8 array's bytes. private static Object adaptUtf8(String column, DType.Utf8 dtype, Object value) { - if (!(value instanceof String[] arr)) { - throw new IllegalArgumentException( - "column '" + column + "' expects String[] for Utf8; got " + value.getClass().getSimpleName()); + if (value instanceof byte[][] utf8) { + for (int i = 0; i < utf8.length; i++) { + if (utf8[i] != null && !Utf8Validation.isWellFormed(utf8[i], 0, utf8[i].length)) { + throw new IllegalArgumentException( + "column '" + column + "' row " + i + " is not well-formed UTF-8"); + } + } + } else if (!(value instanceof String[])) { + throw new IllegalArgumentException("column '" + column + "' expects String[] or UTF-8 byte[][] for Utf8; got " + + value.getClass().getSimpleName()); } - return adaptVarBin(column, dtype.nullable(), arr); + return adaptVarBin(column, dtype.nullable(), (Object[]) value); } - /// Shared by Utf8 (`String[]`) and Binary (`byte[][]`): a non-nullable column rejects any + /// Shared by Utf8 (`String[]` or UTF-8 `byte[][]`) and Binary (`byte[][]`): a non-nullable column rejects any /// `null` element. A nullable one unifies on the NullableData carrier like the primitive /// paths: the raw array (null elements preserved) plus a derived validity bitmap. The writer /// then routes it through MaskedEncoding (or a nullable-capable encoder such as vortex.zstd). diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/DictColumnState.java b/writer/src/main/java/io/github/dfa1/vortex/writer/DictColumnState.java index b2a5d9f5d..18c233f80 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/DictColumnState.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/DictColumnState.java @@ -11,6 +11,7 @@ import io.github.dfa1.vortex.writer.encode.VarBinEncodingEncoder; import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; import java.util.HashSet; @@ -251,7 +252,7 @@ boolean ingestDictChunk(Object data) { codes[i] = (short) nullCode; continue; } - Object v = varBinKey(strings[i]); + Object v = dictKey(strings[i]); Integer code = valueToCode.get(v); if (code == null) { if (valueToCode.size() == cap) { @@ -273,7 +274,7 @@ boolean ingestDictChunk(Object data) { } if (statsStrings && seenInChunk[code] != chunkStamp) { seenInChunk[code] = chunkStamp; - String s = (String) strings[i]; + String s = (String) v; if (minString == null || Utf8Order.compare(s, minString) < 0) { minString = s; } @@ -608,6 +609,15 @@ private static int dictMaxCardinality(boolean varBin) { return varBin ? GLOBAL_DICT_MAX_CARDINALITY_UTF8 : GLOBAL_DICT_MAX_CARDINALITY; } + // The dictionary key of one Utf8/Binary row. Utf8 keys are always Strings, a UTF-8 row decoded, + // so a column written partly as String[] and partly as byte[][] still codes each value once. + private Object dictKey(Object value) { + if (!binary && value instanceof byte[] utf8) { + return new String(utf8, StandardCharsets.UTF_8); + } + return varBinKey(value); + } + // A byte[] compares by identity, so Binary values are keyed by a ByteBuffer view, whose // equals/hashCode compare content. The array is never mutated after ChunkImpl hands it over. private static Object varBinKey(Object value) { diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/ZoneMapStatCodec.java b/writer/src/main/java/io/github/dfa1/vortex/writer/ZoneMapStatCodec.java index 8515e45ef..3199cb42c 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/ZoneMapStatCodec.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/ZoneMapStatCodec.java @@ -147,7 +147,7 @@ static MinMax columnMinMax(DType dtype, Object data) { PrimitiveEncodingEncoder.minMaxStats(p.ptype(), compactIfNullable(p.ptype(), data)); case DType.Utf8 _ -> { Object values = data instanceof NullableData nd ? nd.values() : data; - yield values instanceof String[] strings ? VarBinEncodingEncoder.minMaxStats(strings) : null; + yield VarBinEncodingEncoder.textMinMax(dtype, values); } default -> null; }; diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ConstantEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ConstantEncodingEncoder.java index ef108b0dc..4a0e25303 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ConstantEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ConstantEncodingEncoder.java @@ -9,6 +9,8 @@ import java.lang.foreign.MemorySegment; import java.lang.reflect.Array; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; import java.util.OptionalLong; /// Write-only encoder for `vortex.constant`. @@ -42,8 +44,7 @@ public Estimate expectedRatio(DType dtype, ArrayAndStats data, EncodeContext ctx } if (dtype instanceof DType.Utf8) { // Strings carry no stats; Rust's constant check is one pass over the array - String[] strings = (String[]) data.data(); - return strings.length > 0 && isConstantStrings(strings) ? Estimate.ALWAYS_USE : Estimate.SKIP; + return isConstantText(data.data()) ? Estimate.ALWAYS_USE : Estimate.SKIP; } ArrayStats stats = data.stats(); if (stats.valueCount() == 0) { @@ -64,7 +65,7 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { return encodeBool((boolean[]) data); } if (dtype instanceof DType.Utf8) { - return encodeUtf8((String[]) data); + return encodeUtf8(data); } if (!(dtype instanceof DType.Primitive p)) { throw new VortexException(EncodingId.VORTEX_CONSTANT, "encode only supports Primitive or Bool dtype, got " + dtype); @@ -84,8 +85,7 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { @Override public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext encodeCtx) { if (dtype instanceof DType.Utf8) { - String[] strings = (String[]) data; - if (strings.length == 0 || !isConstantStrings(strings)) { + if (!isConstantText(data)) { return CascadeStep.notApplicable(); } return CascadeStep.terminal(encode(dtype, data, encodeCtx)); @@ -102,24 +102,41 @@ public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext encodeC return CascadeStep.terminal(encode(dtype, data, encodeCtx)); } - private static EncodeResult encodeUtf8(String[] data) { - if (data.length == 0 || !isConstantStrings(data)) { + /// @param data a `String[]` or the column's UTF-8 `byte[][]` + private static EncodeResult encodeUtf8(Object data) { + if (!isConstantText(data)) { throw new VortexException(EncodingId.VORTEX_CONSTANT, "not a constant array"); } - ProtoScalarValue scalar = ProtoScalarValue.ofStringValue(data[0]); + String value; + if (data instanceof byte[][] utf8) { + value = new String(utf8[0], StandardCharsets.UTF_8); + } else { + value = ((String[]) data)[0]; + } + ProtoScalarValue scalar = ProtoScalarValue.ofStringValue(value); byte[] scalarBytes = scalar.encode(); // Both extremes are the one repeated value, by construction return EncodeResult.simple(EncodingId.VORTEX_CONSTANT, EncodedBuffer.bytes(MemorySegment.ofArray(scalarBytes)), scalar, scalar); } - private static boolean isConstantStrings(String[] data) { - for (int i = 1; i < data.length; i++) { - if (!data[i].equals(data[0])) { + /// Whether a non-empty `String[]` or UTF-8 `byte[][]` repeats one value. + private static boolean isConstantText(Object data) { + if (data instanceof byte[][] utf8) { + for (int i = 1; i < utf8.length; i++) { + if (!Arrays.equals(utf8[i], utf8[0])) { + return false; + } + } + return utf8.length > 0; + } + String[] strings = (String[]) data; + for (int i = 1; i < strings.length; i++) { + if (!strings[i].equals(strings[0])) { return false; } } - return true; + return strings.length > 0; } private static EncodeResult encodeBool(boolean[] data) { diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoder.java index 5a76f260f..db8ebff0e 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoder.java @@ -16,10 +16,13 @@ import java.lang.foreign.MemorySegment; import java.lang.foreign.ValueLayout; import java.lang.reflect.Array; +import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; +import java.util.Collection; import java.util.Set; import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; /// Write-only encoder for `vortex.dict`. public final class DictEncodingEncoder implements EncodingEncoder { @@ -112,7 +115,7 @@ static double integerDictRatio(int bitWidth, long n, long distinct, long average @Override public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { if (dtype instanceof DType.Utf8) { - return encodeUtf8((String[]) data, ctx); + return encodeUtf8(data, ctx); } DictData d = buildDictData(dtype, data); PType codePType = d.codePType(); @@ -145,7 +148,7 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { @Override public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { if (dtype instanceof DType.Utf8) { - return encodeUtf8Cascade((String[]) data, ctx); + return encodeUtf8Cascade(data, ctx); } DictData d = buildDictData(dtype, data); PType codePType = d.codePType(); @@ -173,14 +176,16 @@ public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { /// varbin. [CascadingCompressor] excludes the winning `vortex.dict` from that child's /// competition, so the pool is never wrapped in a second dict the reader cannot unwrap. /// - /// @param strings the column's string values - /// @param ctx encoding context supplying the arena + /// @param data the column's values, a `String[]` or UTF-8 `byte[][]` + /// @param ctx encoding context supplying the arena /// @return a cascade step whose values child (index 1) is left open for compression - private static CascadeStep encodeUtf8Cascade(String[] strings, EncodeContext ctx) { - int n = strings.length; - var valueMap = new LinkedHashMap(); - for (String s : strings) { - valueMap.computeIfAbsent(s, _ -> valueMap.size()); + private static CascadeStep encodeUtf8Cascade(Object data, EncodeContext ctx) { + int n = Array.getLength(data); + Object[] keys = new Object[n]; + var valueMap = new LinkedHashMap(); + for (int i = 0; i < n; i++) { + keys[i] = textKey(data, i); + valueMap.computeIfAbsent(keys[i], _ -> valueMap.size()); // Distinct count only grows as the scan proceeds, so once it alone already exceeds // n/2 the final count is guaranteed to too — the same "dict can't win" rule // DictEncodingEncoder#expectedRatio applies to Primitive via stats, mirrored here @@ -199,7 +204,7 @@ private static CascadeStep encodeUtf8Cascade(String[] strings, EncodeContext ctx // Codes as a typed array (U8->byte[], U16->short[], U32->int[]) so the cascade can bitpack // them, rather than emitting a raw vortex.primitive leaf that pays the full U8/U16/U32 width // per row (#303). The primitive dict path already exposes its codes this way. - Object codesArr = buildCodesArray(strings, valueMap, codePType); + Object codesArr = buildCodesArray(keys, valueMap, codePType); byte[] metaBytes = new ProtoDictMetadata( dictSize, @@ -215,10 +220,10 @@ private static CascadeStep encodeUtf8Cascade(String[] strings, EncodeContext ctx new EncodeNode[]{null, null}, new int[0]); - String[] distinct = valueMap.keySet().toArray(new String[0]); + Object distinct = textPool(valueMap.keySet(), data instanceof byte[][]); // UTF-8 byte order like every other string bound (Rust compares bytes), not String's UTF-16 order - MinMax stats = VarBinEncodingEncoder.minMaxStats(valueMap.keySet().toArray(new String[0])); + MinMax stats = VarBinEncodingEncoder.textMinMax(DType.UTF8, distinct); ProtoScalarValue statsMin = stats != null ? stats.min() : null; ProtoScalarValue statsMax = stats != null ? stats.max() : null; @@ -239,16 +244,17 @@ private static CascadeStep encodeUtf8Cascade(String[] strings, EncodeContext ctx /// /// @param dtype the column's nullable Utf8 dtype; also the values pool's dtype, as Rust /// requires - /// @param strings the row values (any element at an invalid row is ignored) + /// @param rows the row values, a `String[]` or UTF-8 `byte[][]` (any element at an invalid + /// row is ignored) /// @param validity per-row validity /// @param ctx the encode context of the enclosing masked column /// @return the encoded dict, or `null` when more than half the rows are distinct (Dict cannot win) - static EncodeResult encodeNullableUtf8(DType dtype, String[] strings, boolean[] validity, EncodeContext ctx) { - int n = strings.length; - String[] keys = new String[n]; - var valueMap = new LinkedHashMap(); + static EncodeResult encodeNullableUtf8(DType dtype, Object rows, boolean[] validity, EncodeContext ctx) { + int n = validity.length; + Object[] keys = new Object[n]; + var valueMap = new LinkedHashMap(); for (int i = 0; i < n; i++) { - keys[i] = validity[i] ? strings[i] : null; + keys[i] = validity[i] ? textKey(rows, i) : null; valueMap.computeIfAbsent(keys[i], _ -> valueMap.size()); // Same "dict can't win" bail-out as encodeUtf8Cascade. if (valueMap.size() * 2 > n) { @@ -257,10 +263,11 @@ static EncodeResult encodeNullableUtf8(DType dtype, String[] strings, boolean[] } int dictSize = valueMap.size(); PType codePType = codePType(dictSize); - String[] pool = valueMap.keySet().toArray(new String[0]); + Object pool = textPool(valueMap.keySet(), rows instanceof byte[][]); boolean[] poolValidity = new boolean[dictSize]; - for (int i = 0; i < dictSize; i++) { - poolValidity[i] = pool[i] != null; + int p = 0; + for (Object key : valueMap.keySet()) { + poolValidity[p++] = key != null; } EncodeContext childCtx = ctx.withDecrementedDepth(); @@ -288,20 +295,63 @@ static EncodeResult encodeNullableUtf8(DType dtype, String[] strings, boolean[] return new EncodeResult(root, buffers, null, null); } - private static Object buildCodesArray(String[] strings, java.util.Map valueMap, PType codePType) { - int[] codes = new int[strings.length]; + private static Object buildCodesArray(Object[] keys, Map valueMap, PType codePType) { + int[] codes = new int[keys.length]; for (int i = 0; i < codes.length; i++) { - codes[i] = valueMap.get(strings[i]); + codes[i] = valueMap.get(keys[i]); } return PrimitiveArrays.fromIntsArray(codes, codePType, EncodingId.VORTEX_DICT); } - private static EncodeResult encodeUtf8(String[] strings, EncodeContext ctx) { - int n = strings.length; + /// The hash key of a Utf8 row: the string itself, or the row's UTF-8 bytes wrapped so they hash + /// and compare by content. A `null` row is a `null` key. + /// + /// @param data a `String[]` or the column's UTF-8 `byte[][]` + /// @param i the row + /// @return the row's key + private static Object textKey(Object data, int i) { + if (data instanceof byte[][] utf8) { + return utf8[i] == null ? null : ByteBuffer.wrap(utf8[i]); + } + return ((String[]) data)[i]; + } - var valueMap = new LinkedHashMap(); - for (String s : strings) { - valueMap.computeIfAbsent(s, _ -> valueMap.size()); + /// The distinct values in code order, in the shape the rows came in: `String[]`, or `byte[][]` + /// when `bytes`. + /// + /// @param keys the dictionary keys in first-seen order, from [#textKey(Object, int)] + /// @param bytes whether the keys wrap UTF-8 bytes + /// @return the pool, `null` where a key is `null` + private static Object textPool(Collection keys, boolean bytes) { + if (!bytes) { + return keys.toArray(new String[0]); + } + byte[][] pool = new byte[keys.size()][]; + int i = 0; + for (Object key : keys) { + if (key != null) { + pool[i] = ((ByteBuffer) key).array(); + } + i++; + } + return pool; + } + + /// One pool entry's UTF-8 bytes. + private static byte[] utf8Of(Object key) { + if (key instanceof ByteBuffer bytes) { + return bytes.array(); + } + return ((String) key).getBytes(StandardCharsets.UTF_8); + } + + private static EncodeResult encodeUtf8(Object data, EncodeContext ctx) { + int n = Array.getLength(data); + Object[] keys = new Object[n]; + var valueMap = new LinkedHashMap(); + for (int i = 0; i < n; i++) { + keys[i] = textKey(data, i); + valueMap.computeIfAbsent(keys[i], _ -> valueMap.size()); } int dictSize = valueMap.size(); @@ -311,8 +361,8 @@ private static EncodeResult encodeUtf8(String[] strings, EncodeContext ctx) { byte[][] dictByteArrays = new byte[dictSize][]; int j = 0; long totalDictBytes = 0; - for (String s : valueMap.keySet()) { - dictByteArrays[j] = s.getBytes(StandardCharsets.UTF_8); + for (Object key : valueMap.keySet()) { + dictByteArrays[j] = utf8Of(key); totalDictBytes += dictByteArrays[j].length; j++; } @@ -331,7 +381,7 @@ private static EncodeResult encodeUtf8(String[] strings, EncodeContext ctx) { MemorySegment codesBuf = arena.allocate((long) n * codeBytes); for (int i = 0; i < n; i++) { - writeCodeToSeg(codesBuf, codePType, i, valueMap.get(strings[i])); + writeCodeToSeg(codesBuf, codePType, i, valueMap.get(keys[i])); } byte[] metaBytes = new ProtoDictMetadata( @@ -357,7 +407,7 @@ private static EncodeResult encodeUtf8(String[] strings, EncodeContext ctx) { new int[0]); // UTF-8 byte order like every other string bound (Rust compares bytes), not String's UTF-16 order - MinMax stats = VarBinEncodingEncoder.minMaxStats(valueMap.keySet().toArray(new String[0])); + MinMax stats = VarBinEncodingEncoder.textMinMax(DType.UTF8, textPool(valueMap.keySet(), data instanceof byte[][])); ProtoScalarValue statsMin = stats != null ? stats.min() : null; ProtoScalarValue statsMax = stats != null ? stats.max() : null; return new EncodeResult(root, List.of(EncodedBuffer.bytes(dictBytesBuf), EncodedBuffer.of(dictOffsetsBuf, PType.I64), diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/FsstEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/FsstEncodingEncoder.java index 510c28e1a..c011aa6c8 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/FsstEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/FsstEncodingEncoder.java @@ -63,24 +63,7 @@ public boolean accepts(DType dtype) { public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { Arena arena = ctx.arena(); VarBinBytes.Rows rows = VarBinBytes.toContiguous(data); - return toEncodeResult(compress(data, rows, arena), arena, zoneMapStats(data, rows)); - } - - /// Zone-map min/max for the rows being encoded, or `null` when there is nothing useful to - /// record. Utf8 rows go through [VarBinEncodingEncoder#minMaxStats(String[])] so an - /// FSST-compressed column reports the same lexicographic bounds a plain `vortex.varbin` one - /// would - without this, choosing FSST for a column silently costs it zone-map pruning and - /// `MIN`/`MAX` push-down. Binary rows are skipped for the same reason [VarBinEncodingEncoder] - /// skips them: a zone map over arbitrary blobs prunes nothing. - /// - /// @param data the Utf8 (`String[]`) or Binary (`byte[][]`) values being encoded - /// @param rows `data` as [VarBinBytes#toContiguous(Object)] lays it out, or `null` above 2 GB - /// @return the [MinMax] of encoded scalars, or `null` - private static MinMax zoneMapStats(Object data, VarBinBytes.Rows rows) { - if (!(data instanceof String[] strings)) { - return null; - } - return rows == null ? VarBinEncodingEncoder.minMaxStats(strings) : VarBinEncodingEncoder.minMaxStats(strings, rows); + return toEncodeResult(compress(data, rows, arena), arena, VarBinEncodingEncoder.textMinMax(dtype, data, rows)); } /// Lays out a compression product in the terminal `vortex.fsst` wire format: the per-row @@ -94,7 +77,8 @@ private static MinMax zoneMapStats(Object data, VarBinBytes.Rows rows) { /// /// @param c the compression product to lay out /// @param arena arena backing the length/offset buffers - /// @param stats the `{min, max}` zone-map scalars from [#zoneMapStats(Object, VarBinBytes.Rows)], or `null` + /// @param stats the `{min, max}` zone-map scalars (the same bounds a plain `vortex.varbin` would + /// report, so choosing FSST costs no pruning), or `null` /// @return the terminal `vortex.fsst` encode result static EncodeResult toEncodeResult(Fsst c, Arena arena, MinMax stats) { long uncompLenBytes = c.uncompLenPType().byteSize(); @@ -142,7 +126,7 @@ public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { } VarBinBytes.Rows rows = VarBinBytes.toContiguous(data); Fsst c = compress(data, rows, ctx.arena()); - MinMax stats = zoneMapStats(data, rows); + MinMax stats = VarBinEncodingEncoder.textMinMax(dtype, data, rows); Object uncompLens = PrimitiveArrays.fromIntsArray(c.uncompLens(), c.uncompLenPType(), EncodingId.VORTEX_FSST); Object codesOffsets = PrimitiveArrays.fromIntsArray(c.codesOffsets(), c.codesOffPType(), EncodingId.VORTEX_FSST); EncodeNode partialRoot = new EncodeNode( diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/MaskedEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/MaskedEncodingEncoder.java index c25dd5f5f..25aa0ab80 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/MaskedEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/MaskedEncodingEncoder.java @@ -75,7 +75,7 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { // makes null one more (invalid) pool entry. Offer that shape too and keep the smaller. if (ctx.allowedCascading() > 0 && nonNullable instanceof DType.Utf8 && !ctx.excluded().contains(EncodingId.VORTEX_DICT)) { - EncodeResult dict = DictEncodingEncoder.encodeNullableUtf8(dtype, (String[]) values, validity, ctx); + EncodeResult dict = DictEncodingEncoder.encodeNullableUtf8(dtype, values, validity, ctx); if (dict != null && totalBytes(dict) < totalBytes(masked)) { return new EncodeResult(dict.rootNode(), dict.encodedBuffers(), statsMin, statsMax); } diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/OnPairEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/OnPairEncodingEncoder.java index e36504a2d..c7885b4a8 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/OnPairEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/OnPairEncodingEncoder.java @@ -72,7 +72,7 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { } EncodeNode root = new EncodeNode(EncodingId.VORTEX_ONPAIR, MemorySegment.ofArray(c.metadata()), children, new int[]{0}); - MinMax stats = zoneMapStats(data, rows); + MinMax stats = VarBinEncodingEncoder.textMinMax(dtype, data, rows); return new EncodeResult(root, buffers, stats != null ? stats.min() : null, stats != null ? stats.max() : null); } @@ -94,13 +94,10 @@ public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { } EncodeNode partialRoot = new EncodeNode(EncodingId.VORTEX_ONPAIR, MemorySegment.ofArray(c.metadata()), new EncodeNode[4], new int[]{0}); - MinMax stats = zoneMapStats(data, rows); + MinMax stats = VarBinEncodingEncoder.textMinMax(dtype, data, rows); return new CascadeStep(partialRoot, List.of(EncodedBuffer.bytes(c.dictBytes(ctx.arena()))), slots, stats != null ? stats.min() : null, stats != null ? stats.max() : null, true); } - private static MinMax zoneMapStats(Object data, VarBinBytes.Rows rows) { - return data instanceof String[] strings ? VarBinEncodingEncoder.minMaxStats(strings, rows) : null; - } /// The compressed column: read-padded dictionary bytes plus the four integer children in slot /// order (`dict_offsets`, `codes`, `codes_offsets`, `uncompressed_lengths`). diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/VarBinEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/VarBinEncodingEncoder.java index 9af696a98..daccbf8d6 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/VarBinEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/VarBinEncodingEncoder.java @@ -12,6 +12,7 @@ import java.lang.foreign.Arena; import java.lang.foreign.MemorySegment; import java.lang.foreign.ValueLayout; +import java.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.List; @@ -34,12 +35,10 @@ public boolean accepts(DType dtype) { @Override public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { - // Binary (DType.Binary) arrives as raw byte[][] and must round-trip byte-for-byte — - // routing it through the Utf8 String[] path would corrupt any byte sequence that isn't - // valid UTF-8 (e.g. an embedded audio/image blob). Utf8 arrives as String[] and is UTF-8 - // encoded. Either way a null entry (this encoder is the values child of a masked/nullable - // layout, where validity carries nullity) contributes a zero-length slot. - String[] strings = data instanceof String[] s ? s : null; + // Binary arrives as raw byte[][] and round-trips byte-for-byte; Utf8 arrives as String[] + // (UTF-8 encoded here) or as its UTF-8 bytes. Either way a null entry (this encoder is the + // values child of a masked/nullable layout, where validity carries nullity) contributes a + // zero-length slot. VarBinBytes.Rows rows = VarBinBytes.toContiguous(data); if (rows == null) { throw new VortexException(EncodingId.VORTEX_VARBIN, "chunk exceeds 2 GB of row bytes"); @@ -57,9 +56,7 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { byte[] metaBytes = new ProtoVarBinMetadata(io.github.dfa1.vortex.core.proto.ProtoPType.fromValue(PType.I64.ordinal())).encode(); - // Zone-map min/max stats are lexicographic-string-only (minMaxStats(String[])); binary - // blobs (e.g. audio bytes) aren't usefully zone-mapped, so skip stats for that path. - MinMax stats = strings != null ? minMaxStats(strings, rows) : null; + MinMax stats = textMinMax(dtype, data, rows); ProtoScalarValue statsMin = stats != null ? stats.min() : null; ProtoScalarValue statsMax = stats != null ? stats.max() : null; @@ -124,6 +121,69 @@ static MinMax minMaxStats(String[] strings, VarBinBytes.Rows rows) { return min < 0 ? null : encodeMinMax(strings[min], strings[max]); } + /// The zone-map min/max of a string column, or `null` for binary values or when every value is + /// `null`. A `String[]` is always text; a `byte[][]` is text (UTF-8) when `dtype`, or an + /// extension's storage, is Utf8. Binary has no useful ordering: a zone map over arbitrary blobs + /// prunes nothing. + /// + /// @param dtype the values' dtype + /// @param values a `String[]` or `byte[][]`, `null` at a null row + /// @return the [MinMax] of encoded string scalars, or `null` + public static MinMax textMinMax(DType dtype, Object values) { + if (values instanceof String[] strings) { + return minMaxStats(strings); + } + if (values instanceof byte[][] utf8 && isUtf8(dtype)) { + return minMaxStats(utf8); + } + return null; + } + + /// [#textMinMax(DType, Object)] reusing the contiguous rows an encoder already built. + /// + /// @param dtype the values' dtype + /// @param values a `String[]` or `byte[][]`, `null` at a null row + /// @param rows `values` as [VarBinBytes#toContiguous(Object)] lays them out, or `null` + /// @return the [MinMax] of encoded string scalars, or `null` + static MinMax textMinMax(DType dtype, Object values, VarBinBytes.Rows rows) { + if (values instanceof String[] strings && rows != null) { + return minMaxStats(strings, rows); + } + return textMinMax(dtype, values); + } + + /// Whether `dtype` holds UTF-8 text, directly or as an extension's storage. + /// + /// @param dtype the dtype, or `null` + /// @return `true` for Utf8 and Utf8-backed extensions + static boolean isUtf8(DType dtype) { + return dtype instanceof DType.Utf8 + || (dtype instanceof DType.Extension ext && ext.storageDType() instanceof DType.Utf8); + } + + /// [#minMaxStats(String[])] over UTF-8 rows, comparing their bytes: [Arrays#compareUnsigned] + /// orders UTF-8 exactly as [Utf8Order] orders the strings. Only the two winners are decoded. + /// + /// @param utf8 the rows' UTF-8 bytes, `null` at a null row + /// @return the [MinMax] of encoded scalars, or `null` if all entries are `null` + static MinMax minMaxStats(byte[][] utf8) { + int min = -1; + int max = -1; + for (int i = 0; i < utf8.length; i++) { + if (utf8[i] == null) { + continue; + } + if (min < 0 || Arrays.compareUnsigned(utf8[i], utf8[min]) < 0) { + min = i; + } + if (max < 0 || Arrays.compareUnsigned(utf8[i], utf8[max]) > 0) { + max = i; + } + } + return min < 0 ? null : encodeMinMax(new String(utf8[min], StandardCharsets.UTF_8), + new String(utf8[max], StandardCharsets.UTF_8)); + } + /// Encodes an already-computed string min/max as [#minMaxStats(String[])] would. /// /// @param minStr the minimum in UTF-8 order, or `null` when there is no non-null value diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ZoneMapStats.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ZoneMapStats.java index 6f0064d54..3193b838a 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ZoneMapStats.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ZoneMapStats.java @@ -25,7 +25,7 @@ /// values are compacted against the validity mask first. Taking the raw array's minimum here /// would report `0` for a column of positive values, which is worse than reporting nothing. /// - a primitive array (`long[]`, `double[]`, …) matching the dtype's [io.github.dfa1.vortex.core.model.PType] -/// - `String[]`, compared lexicographically, skipping nulls +/// - `String[]`, or a Utf8 column's UTF-8 `byte[][]`, compared in UTF-8 byte order, skipping nulls /// /// Anything else - binary blobs, nested types - has no useful ordering and yields `null`. final class ZoneMapStats { @@ -48,12 +48,12 @@ static MinMax of(DType dtype, Object data) { if (nonNullable instanceof DType.Decimal d && values instanceof BigDecimal[] decimals) { return ofDecimal(d, decimals); } - // A String[] carries a real null at its invalid positions, so the lexicographic - // null-skip is already correct and needs no compaction. - return values instanceof String[] strings ? VarBinEncodingEncoder.minMaxStats(strings) : null; + // A String[]/byte[][] carries a real null at its invalid positions, so the + // lexicographic null-skip is already correct and needs no compaction. + return VarBinEncodingEncoder.textMinMax(nonNullable, values); } - if (data instanceof String[] strings) { - return VarBinEncodingEncoder.minMaxStats(strings); + if (data instanceof String[] || data instanceof byte[][]) { + return VarBinEncodingEncoder.textMinMax(dtype, data); } DType effective = dtype instanceof DType.Extension ext ? ext.storageDType() : dtype; if (effective instanceof DType.Primitive p && data != null && data.getClass().isArray()) { diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/Utf8BytesInputTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/Utf8BytesInputTest.java new file mode 100644 index 000000000..1853244c7 --- /dev/null +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/Utf8BytesInputTest.java @@ -0,0 +1,123 @@ +package io.github.dfa1.vortex.writer; + +import io.github.dfa1.vortex.core.model.ColumnName; +import io.github.dfa1.vortex.core.model.DType; +import org.junit.jupiter.api.Named; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +import java.io.IOException; +import java.nio.channels.FileChannel; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.function.IntFunction; +import java.util.stream.Stream; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.junit.jupiter.api.Named.named; + +/// A Utf8 column given as its UTF-8 bytes must write exactly the file the same strings do: the +/// bytes path skips encoding each string, and every encoder, stat and dictionary has to treat the +/// two carriers alike for that to hold. +class Utf8BytesInputTest { + + private static final ColumnName S = ColumnName.of("s"); + private static final int CHUNKS = 3; + private static final int ROWS = 20_000; + + @TempDir + Path tmp; + + static Stream columns() { + List>> shapes = List.of( + // dict wins: a few values + named("low cardinality", i -> "city-" + (i * 7919 % 8)), + // FSST/Zstd compete: unique values sharing a prefix + named("high cardinality", i -> "note-" + Integer.toHexString(i * 0x9E3779B1)), + // UTF-16 and UTF-8 order disagree: a surrogate pair against U+E000.., plus 2/3-byte + named("non-ASCII", i -> switch (i % 5) { + case 0 -> "café " + (i % 50); + case 1 -> "😀" + (i % 50); + case 2 -> "｡" + (i % 50); + case 3 -> "€" + i; + default -> "plain" + i; + }), + named("constant", _ -> "same"), + named("nullable", i -> i % 11 == 0 ? null : "v" + (i % 40))); + List> options = List.of( + named("defaults", WriteOptions.defaults()), + named("compact", WriteOptions.defaults().withCompact(true)), + named("no global dict", WriteOptions.defaults().withGlobalDict(false))); + List cases = new ArrayList<>(); + for (Named> shape : shapes) { + for (Named option : options) { + cases.add(Arguments.of(shape, option)); + } + } + return cases.stream(); + } + + @ParameterizedTest(name = "{0}, {1}") + @MethodSource("columns") + void utf8Bytes_writeTheSameFileAsStrings(IntFunction value, WriteOptions options) throws IOException { + // Given + List chunks = new ArrayList<>(); + Random random = new Random(42); + for (int c = 0; c < CHUNKS; c++) { + String[] rows = new String[ROWS]; + for (int i = 0; i < ROWS; i++) { + rows[i] = value.apply(random.nextInt(1_000_000)); + } + chunks.add(rows); + } + boolean nullable = chunks.stream().flatMap(Stream::of).anyMatch(s -> s == null); + Path expected = write("strings.vortex", nullable, options, chunks.stream().map(r -> (Object) r).toList()); + + // When + Path result = write("bytes.vortex", nullable, options, chunks.stream().map(Utf8BytesInputTest::utf8).toList()); + + // Then + assertThat(Files.readAllBytes(result)).isEqualTo(Files.readAllBytes(expected)); + } + + @Test + void malformedUtf8_isRejected() throws IOException { + // Given: an overlong encoding of '/', which Rust's Utf8 validation rejects + byte[][] rows = {"ok".getBytes(StandardCharsets.UTF_8), {(byte) 0xc0, (byte) 0xaf}}; + + // When / Then + assertThatThrownBy(() -> write("bad.vortex", false, WriteOptions.defaults(), List.of((Object) rows))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("row 1 is not well-formed UTF-8"); + } + + private Path write(String name, boolean nullable, WriteOptions options, List chunks) throws IOException { + Path file = tmp.resolve(name); + var schema = new DType.Struct(List.of(S), List.of(new DType.Utf8(nullable)), false); + try (var channel = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE); + var sut = VortexWriter.create(channel, schema, options)) { + for (Object chunk : chunks) { + sut.writeChunk(Map.of(S, chunk)); + } + } + return file; + } + + private static Object utf8(String[] rows) { + byte[][] out = new byte[rows.length][]; + for (int i = 0; i < rows.length; i++) { + out[i] = rows[i] == null ? null : rows[i].getBytes(StandardCharsets.UTF_8); + } + return out; + } +} From dbae5123d5c27270ce1dc5712ac8907b4261adbb Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Thu, 22 Oct 2026 22:56:40 +0200 Subject: [PATCH 4/4] perf: CSV import about 1.7x faster Profiled a 2.7 GB / 45M-row CSV import (YourKit, JFR with DebugNonSafepoints) and fixed what showed up. Output is byte-identical in default and compact modes. CSV import: - Parse the next chunk while a background thread converts and writes the previous one: parsing is the import's one sequential step. - CsvByteReader parses bytes into a column-wise ColumnChunk instead of FastCSV's chars and a String per field, and Utf8 columns reach the writer as UTF-8 byte[][]. Same format as FastCSV's defaults: CsvByteReaderTest reads 200 random files (quotes, "", every line ending, empty lines, multi-byte and malformed UTF-8, a BOM, 1-7-byte reads) against FastCSV. The delimiter must now be ASCII. - Numbers parse in place; a plain decimal whose digits fit 2^53 is digits / 10^k, bit-identical to Double.parseDouble (Clinger's fast path, 20 000 random decimals checked). A string column is checked for non-ASCII once per chunk, eight bytes at a time. Writer (CPU the profile pinned on it): - Zstd orders string zone stats by the bytes it already encoded (6.9 s -> 1.2 s of a compact import) instead of comparing Strings again. - Zone min/max kept as primitives: an ascending id boxed a Long per row. - Repartitioner sizes strings with getBytes(UTF_8).length, twice as fast as the hand-written count and exact by definition. - OnPair's first two trie levels are direct-indexed and deeper edges sit in a one-array table: it wins high-cardinality strings by default and its per-byte hashing was 38% of a default import. 2.7 GB CSV, interleaved runs: default 13.13 s -> 7.69 s, compact (loaded machine) 15.65 s -> 10.71 s. Tried and dropped: pooled zstd contexts, regrouping fields at parse time, and parallel block parsing (correct, but slower on a CPU-bound 10-core laptop). Co-Authored-By: Claude Sonnet 5 --- CHANGELOG.md | 1 + .../github/dfa1/vortex/csv/ColumnChunk.java | 111 ++++++ .../github/dfa1/vortex/csv/CsvByteReader.java | 200 +++++++++++ .../github/dfa1/vortex/csv/CsvImporter.java | 327 ++++++++++++++---- .../github/dfa1/vortex/csv/ImportOptions.java | 3 +- .../dfa1/vortex/csv/CsvByteReaderTest.java | 175 ++++++++++ .../dfa1/vortex/csv/CsvImporterTest.java | 187 +++++++++- .../dfa1/vortex/writer/Repartitioner.java | 29 +- .../dfa1/vortex/writer/ZoneAccumulator.java | 52 ++- .../writer/encode/OnPairEncodingEncoder.java | 119 ++++++- .../writer/encode/VarBinEncodingEncoder.java | 23 ++ .../writer/encode/ZstdEncodingEncoder.java | 28 +- .../dfa1/vortex/writer/RepartitionerTest.java | 16 - .../encode/VarBinEncodingEncoderTest.java | 39 ++- 14 files changed, 1147 insertions(+), 163 deletions(-) create mode 100644 csv/src/main/java/io/github/dfa1/vortex/csv/ColumnChunk.java create mode 100644 csv/src/main/java/io/github/dfa1/vortex/csv/CsvByteReader.java create mode 100644 csv/src/test/java/io/github/dfa1/vortex/csv/CsvByteReaderTest.java diff --git a/CHANGELOG.md b/CHANGELOG.md index 0843e4ec9..899c06bde 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed - **Breaking:** SQL identifiers are quoted with double quotes, as IntelliJ and DBeaver generate them, instead of back-ticks, so opening a table or sorting its grid works; write `"date"` where you wrote `` `date` ``. ([#588](https://github.com/dfa1/vortex-java/pull/588)) - **Breaking:** the `ScanIterator` constructor and the reader's Pco helpers (`PcoTansDecoder`, `PcoBin`, `LeBitReader`) are no longer public, and the built-in extension codecs are reached only through their `INSTANCE`. ([4ba6ea12](https://github.com/dfa1/vortex-java/commit/4ba6ea12)) +- CSV import is about 40% faster: it parses bytes instead of `String`s and overlaps parsing with writing; the delimiter must be an ASCII character. ([#590](https://github.com/dfa1/vortex-java/pull/590)) - SQL `ORDER BY ... LIMIT` on a string column is up to 3 times faster: rows that cannot make the top are ranked by their bytes without building a `String`, and a chunk stored in the opposite order is read backward. ([#575](https://github.com/dfa1/vortex-java/pull/575)) ## [0.18.0] — 2026-10-20 diff --git a/csv/src/main/java/io/github/dfa1/vortex/csv/ColumnChunk.java b/csv/src/main/java/io/github/dfa1/vortex/csv/ColumnChunk.java new file mode 100644 index 000000000..e973287e2 --- /dev/null +++ b/csv/src/main/java/io/github/dfa1/vortex/csv/ColumnChunk.java @@ -0,0 +1,111 @@ +package io.github.dfa1.vortex.csv; + +import java.nio.charset.StandardCharsets; +import java.util.Arrays; + +/// One chunk of CSV fields kept column by column as raw bytes: each column is one growing byte +/// buffer plus the end offset of each row's field, so parsing allocates nothing per field and each +/// column's conversion reads one contiguous buffer. +final class ColumnChunk { + + private static final int INITIAL_COLUMN_BYTES = 1 << 12; + + private byte[][] bytes = new byte[0][]; + private int[][] ends = new int[0][]; + private int[] used = new int[0]; + private final int capacity; + private int rows; + + /// @param capacity the most rows the chunk holds + ColumnChunk(int capacity) { + this.capacity = capacity; + } + + /// Appends `src[from, to)` to field `column` of the row being read; a field may arrive in + /// several pieces (a quoted field around each escaped quote). + void append(int column, byte[] src, int from, int to) { + if (column >= used.length) { + addColumns(column + 1); + } + int length = to - from; + byte[] buffer = bytes[column]; + int at = used[column]; + if (buffer.length - at < length) { + buffer = Arrays.copyOf(buffer, Math.max(buffer.length * 2, at + length)); + bytes[column] = buffer; + } + System.arraycopy(src, from, buffer, at, length); + used[column] = at + length; + } + + /// Ends field `column` of the row being read. + void endField(int column) { + if (column >= used.length) { + addColumns(column + 1); + } + ends[column][rows] = used[column]; + } + + /// Ends the row being read. + void endRow() { + rows++; + } + + /// Drops the fields of the row being read, which the input ended in the middle of. + void abortRow() { + for (int c = 0; c < used.length; c++) { + used[c] = rows == 0 ? 0 : ends[c][rows - 1]; + } + } + + /// The number of complete rows. + int rows() { + return rows; + } + + /// Whether the chunk holds its capacity of rows. + boolean isFull() { + return rows == capacity; + } + + /// The number of columns any row has had. + int columns() { + return used.length; + } + + /// Column `column`'s bytes, or `null` for a column no row has. + byte[] bytes(int column) { + return column < bytes.length ? bytes[column] : null; + } + + /// Where row `row`'s field starts in [#bytes(int)]. + int start(int column, int row) { + return row == 0 ? 0 : ends[column][row - 1]; + } + + /// Where row `row`'s field ends in [#bytes(int)]. + int end(int column, int row) { + return ends[column][row]; + } + + /// Row `row`'s field in column `column` as text, `""` for a column the row lacks; malformed + /// UTF-8 decodes to U+FFFD as the JDK decoder does. + String text(int column, int row) { + if (column >= bytes.length) { + return ""; + } + int from = start(column, row); + return new String(bytes[column], from, end(column, row) - from, StandardCharsets.UTF_8); + } + + private void addColumns(int count) { + int old = used.length; + bytes = Arrays.copyOf(bytes, count); + ends = Arrays.copyOf(ends, count); + used = Arrays.copyOf(used, count); + for (int c = old; c < count; c++) { + bytes[c] = new byte[INITIAL_COLUMN_BYTES]; + ends[c] = new int[capacity]; + } + } +} diff --git a/csv/src/main/java/io/github/dfa1/vortex/csv/CsvByteReader.java b/csv/src/main/java/io/github/dfa1/vortex/csv/CsvByteReader.java new file mode 100644 index 000000000..eea8666f1 --- /dev/null +++ b/csv/src/main/java/io/github/dfa1/vortex/csv/CsvByteReader.java @@ -0,0 +1,200 @@ +package io.github.dfa1.vortex.csv; + +import java.io.Closeable; +import java.io.IOException; +import java.io.InputStream; +import java.util.Arrays; + +/// Reads CSV records straight from bytes into a [ColumnChunk], never decoding to characters. +/// +/// The import used FastCSV, which decodes the input to `char`s and builds a `String` per field; +/// the writer then encodes every string back to UTF-8. Parsing bytes skips both, and since the +/// delimiter, quote and line breaks are ASCII, they can never be part of a multi-byte UTF-8 +/// character, so the bytes split exactly where the characters do. +/// +/// The format is FastCSV's default, which the import had: RFC 4180 fields, `"` quotes with `""` +/// for a quote inside one, records ended by `\n`, `\r` or `\r\n`, empty lines skipped, and every +/// record holding as many fields as the first. Anything but a delimiter or line break after a +/// closing quote is an error; a quote left open runs to the end of the input. +final class CsvByteReader implements Closeable { + + private static final int BUFFER_BYTES = 1 << 20; + private static final byte QUOTE = '"'; + private static final byte CR = '\r'; + private static final byte LF = '\n'; + + private final InputStream in; + private final byte delimiter; + private byte[] buf = new byte[BUFFER_BYTES]; + private int pos; + private int limit; + private boolean eof; + private int expectedFields = -1; + private long records; + + /// @param in the CSV bytes + /// @param delimiter the field delimiter, an ASCII character other than a quote or line break + CsvByteReader(InputStream in, char delimiter) { + if (delimiter >= 0x80 || delimiter == QUOTE || delimiter == CR || delimiter == LF) { + throw new IllegalArgumentException("the CSV delimiter must be an ASCII character other than a quote" + + " or line break, got '" + delimiter + "'"); + } + this.in = in; + this.delimiter = (byte) delimiter; + } + + /// Appends the next record's fields to `out` as one row. + /// + /// @param out the chunk receiving the record + /// @return `false` at the end of the input, with nothing appended + /// @throws IOException if reading fails or the record is malformed + boolean readRecord(ColumnChunk out) throws IOException { + while (true) { + int end = parseRecord(out); + if (end >= 0) { + pos = end; + return true; + } + if (end == END_OF_INPUT) { + return false; + } + out.abortRow(); + refill(); + } + } + + private static final int NEED_MORE = -1; + private static final int END_OF_INPUT = -2; + + /// Parses the record starting at `pos`, returning where it ends, [#NEED_MORE] when the buffer + /// ends before the record does (the caller rolls back, refills and parses it again), or + /// [#END_OF_INPUT]. + private int parseRecord(ColumnChunk out) throws IOException { + byte[] b = buf; + int p = pos; + // An empty line, or the \n of a \r\n split across two reads, holds no record. + while (p < limit && (b[p] == LF || b[p] == CR)) { + p++; + } + pos = p; + if (p == limit) { + return eof ? END_OF_INPUT : NEED_MORE; + } + int field = 0; + while (true) { + if (p == limit) { + // a delimiter ended the input: one last, empty field + if (!eof) { + return NEED_MORE; + } + out.endField(field); + return endRecord(out, field + 1, p); + } + if (b[p] == QUOTE) { + p = parseQuoted(out, field, p); + if (p < 0) { + return NEED_MORE; + } + } else { + int start = p; + byte d = delimiter; + while (p < limit) { + byte c = b[p]; + if (c == d || c == LF || c == CR) { + break; + } + p++; + } + if (p == limit && !eof) { + return NEED_MORE; + } + out.append(field, b, start, p); + } + out.endField(field); + field++; + if (p == limit) { + return endRecord(out, field, p); + } + byte c = b[p++]; + if (c != delimiter) { + // a line break: \r\n is one, a lone \n left by a split read is skipped as an empty line + if (c == CR && p < limit && b[p] == LF) { + p++; + } + return endRecord(out, field, p); + } + } + } + + /// Parses the quoted field opening at `p`, returning the position after its closing quote, or + /// -1 when the buffer ends first. + private int parseQuoted(ColumnChunk out, int field, int p) throws IOException { + byte[] b = buf; + int q = p + 1; + int chunkStart = q; + while (true) { + while (q < limit && b[q] != QUOTE) { + q++; + } + if (q == limit) { + if (eof) { + // FastCSV keeps an unclosed quoted field: it runs to the end of the input + out.append(field, b, chunkStart, q); + return q; + } + return -1; + } + out.append(field, b, chunkStart, q); + if (q + 1 == limit && !eof) { + // the next byte decides whether this quote is escaped or closing + return -1; + } + if (q + 1 < limit && b[q + 1] == QUOTE) { + out.append(field, b, q, q + 1); + q += 2; + chunkStart = q; + continue; + } + int after = q + 1; + if (after < limit && b[after] != delimiter && b[after] != LF && b[after] != CR) { + throw new IOException("record " + (records + 1) + " has characters after a closing quote"); + } + return after; + } + } + + private int endRecord(ColumnChunk out, int fields, int end) throws IOException { + records++; + if (expectedFields < 0) { + expectedFields = fields; + } else if (fields != expectedFields) { + throw new IOException("record " + records + " has " + fields + " fields, but the first record had " + + expectedFields); + } + out.endRow(); + return end; + } + + /// Keeps the unparsed tail, growing the buffer when one record fills it, and reads more. + private void refill() throws IOException { + int tail = limit - pos; + if (pos == 0 && limit == buf.length) { + buf = Arrays.copyOf(buf, buf.length * 2); + } else { + System.arraycopy(buf, pos, buf, 0, tail); + } + pos = 0; + limit = tail; + int read = in.read(buf, limit, buf.length - limit); + if (read < 0) { + eof = true; + } else { + limit += read; + } + } + + @Override + public void close() throws IOException { + in.close(); + } +} diff --git a/csv/src/main/java/io/github/dfa1/vortex/csv/CsvImporter.java b/csv/src/main/java/io/github/dfa1/vortex/csv/CsvImporter.java index e07b82248..d68a03d80 100644 --- a/csv/src/main/java/io/github/dfa1/vortex/csv/CsvImporter.java +++ b/csv/src/main/java/io/github/dfa1/vortex/csv/CsvImporter.java @@ -1,7 +1,6 @@ package io.github.dfa1.vortex.csv; -import de.siegmar.fastcsv.reader.CsvReader; -import de.siegmar.fastcsv.reader.CsvRecord; +import io.github.dfa1.vortex.core.compute.Utf8Validation; import io.github.dfa1.vortex.core.model.ColumnName; import io.github.dfa1.vortex.core.model.DType; import io.github.dfa1.vortex.core.model.PType; @@ -9,11 +8,15 @@ import java.io.IOException; import java.io.InputStream; +import java.lang.invoke.MethodHandles; +import java.lang.invoke.VarHandle; +import java.nio.ByteOrder; import java.net.URI; import java.net.http.HttpClient; import java.net.http.HttpRequest; import java.net.http.HttpResponse; import java.nio.channels.FileChannel; +import java.nio.file.Files; import java.nio.file.Path; import java.nio.file.StandardOpenOption; import java.util.ArrayList; @@ -22,7 +25,14 @@ import java.util.List; import java.util.Map; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; + +import static java.nio.charset.StandardCharsets.UTF_8; /// Parses a CSV file and writes a Vortex file. /// @@ -71,7 +81,7 @@ public static void importCsv(Path csvPath, Path vortexPath) throws IOException { /// @throws IOException if reading or writing fails /// @throws IllegalArgumentException if the CSV file has no data rows public static void importCsv(Path csvPath, Path vortexPath, ImportOptions options) throws IOException { - try (CsvReader reader = csvReader(csvPath, options)) { + try (CsvByteReader reader = new CsvByteReader(Files.newInputStream(csvPath), options.delimiter())) { importCsv(reader, vortexPath, options); } } @@ -119,33 +129,96 @@ public static void importCsv(URI csvUri, Path vortexPath, ImportOptions options) } throw failure; } - try (InputStream body = response.body(); - CsvReader reader = CsvReader.builder() - .fieldSeparator(options.delimiter()) - .ofCsvRecord(body)) { + try (CsvByteReader reader = new CsvByteReader(response.body(), options.delimiter())) { importCsv(reader, vortexPath, options); } } /// The buffered first chunk of data rows (used for schema inference) plus the resolved /// header names and schema, produced by [#readFirstChunk]. - private record FirstChunk(List rows, DType.Struct schema) { + private record FirstChunk(ColumnChunk rows, DType.Struct schema) { } - private static void importCsv(CsvReader reader, Path vortexPath, ImportOptions options) + private static void importCsv(CsvByteReader reader, Path vortexPath, ImportOptions options) throws IOException { FirstChunk firstChunk = readFirstChunk(reader, options); DType.Struct schema = firstChunk.schema(); + // Closed in reverse: the chunk writer drains its last chunk before the Vortex writer closes. try (FileChannel channel = FileChannel.open( vortexPath, StandardOpenOption.CREATE, StandardOpenOption.WRITE, StandardOpenOption.TRUNCATE_EXISTING); - VortexWriter writer = VortexWriter.create(channel, schema, options.writeOptions())) { + VortexWriter writer = VortexWriter.create(channel, schema, options.writeOptions()); + ChunkWriter chunkWriter = new ChunkWriter(writer, schema, options.writeOptions().executor())) { // firstChunk isn't touched again after this call, so its buffered rows (up to - // chunkSize row arrays) become collectible before the streaming loop below runs, - // rather than staying reachable for the rest of a potentially large import. - long totalRows = writeFirstChunk(writer, firstChunk, options); - streamRemainingChunks(reader, writer, schema, options, totalRows); + // chunkSize row arrays) become collectible once written, rather than staying + // reachable for the rest of a potentially large import. + long totalRows = writeFirstChunk(chunkWriter, firstChunk, options); + streamRemainingChunks(reader, chunkWriter, options, totalRows); + chunkWriter.await(); + } + } + + /// Converts and writes each chunk on one background thread while the caller parses the next. + /// + /// Parsing is single-threaded and the import's critical path: converting a chunk's text to + /// columns and handing it to the writer used to stall it for about a third of a large import + /// while the encoding threads sat partly idle. One chunk is in flight at a time, so chunks are + /// written in file order and memory stays at two chunks of rows. + private static final class ChunkWriter implements AutoCloseable { + + private final VortexWriter writer; + private final DType.Struct schema; + private final Executor columnExecutor; + private final ExecutorService thread = Executors.newSingleThreadExecutor(); + private Future pending = CompletableFuture.completedFuture(null); + + ChunkWriter(VortexWriter writer, DType.Struct schema, Executor columnExecutor) { + this.writer = writer; + this.schema = schema; + this.columnExecutor = columnExecutor; + } + + /// Waits for the previous chunk, then starts writing `rows`. + void submit(ColumnChunk rows) throws IOException { + await(); + pending = thread.submit(() -> { + writer.writeChunk(buildChunk(schema, rows, columnExecutor)); + return null; + }); + } + + /// Waits for the chunk in flight, rethrowing its failure. + void await() throws IOException { + try { + pending.get(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("interrupted writing a chunk", e); + } catch (ExecutionException e) { + // The column conversion runs on futures, so its failure arrives wrapped once more. + Throwable cause = e.getCause(); + if (cause instanceof CompletionException && cause.getCause() != null) { + cause = cause.getCause(); + } + if (cause instanceof IOException io) { + throw io; + } + if (cause instanceof RuntimeException re) { + throw re; + } + if (cause instanceof Error error) { + throw error; + } + throw new IOException(cause); + } + } + + /// Waits for the chunk in flight (without rethrowing: the caller is already failing or has + /// awaited it) so the Vortex writer is never closed under a running write. + @Override + public void close() { + thread.close(); } } @@ -154,31 +227,29 @@ private static void importCsv(CsvReader reader, Path vortexPath, Impo /// those buffered rows unless [ImportOptions#schema()] already gives one. /// /// @throws IllegalArgumentException if the CSV file has no data rows - private static FirstChunk readFirstChunk(CsvReader reader, ImportOptions options) { - int chunkSize = options.chunkSize(); + private static FirstChunk readFirstChunk(CsvByteReader reader, ImportOptions options) throws IOException { String[] headers = null; - List rows = new ArrayList<>(chunkSize); - boolean expectHeader = options.header(); - - for (CsvRecord csvRecord : reader) { - if (expectHeader) { - headers = csvRecord.getFields().toArray(String[]::new); - expectHeader = false; - } else { - rows.add(csvRecord.getFields().toArray(String[]::new)); - if (rows.size() == chunkSize) { - break; + if (options.header()) { + ColumnChunk header = new ColumnChunk(1); + if (reader.readRecord(header)) { + headers = new String[header.columns()]; + for (int c = 0; c < headers.length; c++) { + headers[c] = header.text(c, 0); } } } + ColumnChunk rows = new ColumnChunk(options.chunkSize()); + while (!rows.isFull() && reader.readRecord(rows)) { + // fill the first chunk + } - if (rows.isEmpty()) { + if (rows.rows() == 0) { throw new IllegalArgumentException("CSV file has no data rows"); } // Generate synthetic column names when the file has no header row. if (headers == null) { - headers = generateHeaders(rows.getFirst().length); + headers = generateHeaders(rows.columns()); } DType.Struct schema = options.schema() != null @@ -190,43 +261,42 @@ private static FirstChunk readFirstChunk(CsvReader reader, ImportOpti /// Writes the buffered first chunk and reports progress once. /// /// @return the row count written so far (the first chunk's size) - private static long writeFirstChunk(VortexWriter writer, FirstChunk firstChunk, ImportOptions options) + private static long writeFirstChunk(ChunkWriter chunkWriter, FirstChunk firstChunk, ImportOptions options) throws IOException { - writer.writeChunk(buildChunk(firstChunk.schema(), firstChunk.rows(), options.writeOptions().executor())); - long totalRows = firstChunk.rows().size(); + chunkWriter.submit(firstChunk.rows()); + long totalRows = firstChunk.rows().rows(); reportProgress(options, totalRows); return totalRows; } /// Streams the rest of the file through the still-open `reader`, writing a chunk every /// [ImportOptions#chunkSize()] rows and reporting progress every [#PROGRESS_BATCH] rows. - private static void streamRemainingChunks(CsvReader reader, VortexWriter writer, - DType.Struct schema, ImportOptions options, long rowsSoFar) throws IOException { + private static void streamRemainingChunks(CsvByteReader reader, ChunkWriter chunkWriter, + ImportOptions options, long rowsSoFar) throws IOException { int chunkSize = options.chunkSize(); long totalRows = rowsSoFar; long lastReported = totalRows; - List chunk = new ArrayList<>(chunkSize); - for (CsvRecord csvRecord : reader) { - chunk.add(csvRecord.getFields().toArray(String[]::new)); + ColumnChunk chunk = new ColumnChunk(chunkSize); + while (reader.readRecord(chunk)) { totalRows++; - if (chunk.size() == chunkSize) { - writer.writeChunk(buildChunk(schema, chunk, options.writeOptions().executor())); - chunk.clear(); + if (chunk.isFull()) { + chunkWriter.submit(chunk); + chunk = new ColumnChunk(chunkSize); } if (totalRows - lastReported >= PROGRESS_BATCH) { reportProgress(options, totalRows); lastReported = totalRows; } } - if (!chunk.isEmpty()) { - writer.writeChunk(buildChunk(schema, chunk, options.writeOptions().executor())); + if (chunk.rows() > 0) { + chunkWriter.submit(chunk); } if (totalRows > lastReported) { reportProgress(options, totalRows); } } - private static DType.Struct inferSchemaFromRows(String[] headers, List rows) { + private static DType.Struct inferSchemaFromRows(String[] headers, ColumnChunk rows) { int colCount = headers.length; boolean[] canBeLong = new boolean[colCount]; boolean[] canBeDouble = new boolean[colCount]; @@ -235,9 +305,9 @@ private static DType.Struct inferSchemaFromRows(String[] headers, List Arrays.fill(canBeDouble, true); Arrays.fill(canBeBool, true); - for (String[] row : rows) { + for (int row = 0; row < rows.rows(); row++) { for (int c = 0; c < colCount; c++) { - probeCell(canBeLong, canBeDouble, canBeBool, c, safeGet(row, c)); + probeCell(canBeLong, canBeDouble, canBeBool, c, rows.text(c, row)); } } @@ -279,12 +349,6 @@ private static void reportProgress(ImportOptions options, long totalRows) { } } - private static CsvReader csvReader(Path path, ImportOptions options) throws IOException { - return CsvReader.builder() - .fieldSeparator(options.delimiter()) - .ofCsvRecord(path); - } - private static String[] generateHeaders(int colCount) { String[] headers = new String[colCount]; for (int i = 0; i < colCount; i++) { @@ -306,19 +370,14 @@ private static DType resolveType(boolean canBeLong, boolean canBeDouble, boolean return DType.UTF8; } - static Map buildChunk(DType.Struct schema, List rows) { - return buildChunk(schema, rows, Runnable::run); - } - /// Converts the text of each column on `executor`: the columns are independent, and with the /// parse this was the caller thread's work ahead of the (parallel) compression. - static Map buildChunk(DType.Struct schema, List rows, Executor executor) { - int n = rows.size(); + static Map buildChunk(DType.Struct schema, ColumnChunk rows, Executor executor) { List> columns = new ArrayList<>(); for (int c = 0; c < schema.fieldNames().size(); c++) { DType dtype = schema.fieldTypes().get(c); - int colIdx = c; - columns.add(CompletableFuture.supplyAsync(() -> buildColumn(dtype, rows, colIdx, n), executor)); + int column = c; + columns.add(CompletableFuture.supplyAsync(() -> buildColumn(dtype, rows, column), executor)); } Map chunk = new LinkedHashMap<>(); for (int c = 0; c < columns.size(); c++) { @@ -327,35 +386,48 @@ static Map buildChunk(DType.Struct schema, List ro return chunk; } - private static Object buildColumn(DType dtype, List rows, int colIdx, int n) { + /// Converts one column from its bytes, as the `String` parsers would convert its text: plain + /// ASCII numbers are parsed in place, anything else falls back to `Long.parseLong` and friends + /// on the decoded text, so their results, and their errors, are unchanged. A column the file + /// lacks (a schema with more fields than the CSV) is all empty fields. + private static Object buildColumn(DType dtype, ColumnChunk rows, int column) { + int n = rows.rows(); + byte[] bytes = column < rows.columns() ? rows.bytes(column) : null; return switch (dtype) { case DType.Primitive p when p.ptype() == PType.I64 -> { long[] arr = new long[n]; - for (int i = 0; i < n; i++) { - String v = safeGet(rows.get(i), colIdx); - arr[i] = v.isEmpty() ? 0L : Long.parseLong(v); + for (int i = 0; bytes != null && i < n; i++) { + arr[i] = parseLong(bytes, rows.start(column, i), rows.end(column, i)); } yield arr; } case DType.Primitive p when p.ptype() == PType.F64 -> { double[] arr = new double[n]; - for (int i = 0; i < n; i++) { - String v = safeGet(rows.get(i), colIdx); - arr[i] = v.isEmpty() ? 0.0 : Double.parseDouble(v); + for (int i = 0; bytes != null && i < n; i++) { + arr[i] = parseDouble(bytes, rows.start(column, i), rows.end(column, i)); } yield arr; } case DType.Bool _ -> { boolean[] arr = new boolean[n]; - for (int i = 0; i < n; i++) { - arr[i] = Boolean.parseBoolean(safeGet(rows.get(i), colIdx)); + for (int i = 0; bytes != null && i < n; i++) { + arr[i] = parseBoolean(bytes, rows.start(column, i), rows.end(column, i)); } yield arr; } case DType.Utf8 _ -> { - String[] arr = new String[n]; + byte[][] arr = new byte[n][]; + // one word-at-a-time pass over the column instead of a check per field: CSV text + // is usually all ASCII, which is always well-formed UTF-8 + boolean ascii = bytes == null || n == 0 || isAscii(bytes, 0, rows.end(column, n - 1)); for (int i = 0; i < n; i++) { - arr[i] = safeGet(rows.get(i), colIdx); + if (bytes == null) { + arr[i] = EMPTY; + } else if (ascii) { + arr[i] = Arrays.copyOfRange(bytes, rows.start(column, i), rows.end(column, i)); + } else { + arr[i] = utf8(bytes, rows.start(column, i), rows.end(column, i)); + } } yield arr; } @@ -363,10 +435,117 @@ private static Object buildColumn(DType dtype, List rows, int colIdx, }; } - private static String safeGet(String[] row, int idx) { - if (idx >= row.length || row[idx] == null) { - return ""; + private static final byte[] EMPTY = new byte[0]; + + /// `Long.parseLong` on the field's text, `0` for an empty field. Up to 18 ASCII digits cannot + /// overflow a long, so they are summed in place; anything else takes the `String` path. + static long parseLong(byte[] b, int from, int to) { + if (from == to) { + return 0L; + } + int i = from; + boolean negative = b[i] == '-'; + if (negative || b[i] == '+') { + i++; + } + if (i == to || to - i > 18) { + return Long.parseLong(new String(b, from, to - from, UTF_8)); + } + long value = 0; + for (; i < to; i++) { + int digit = b[i] - '0'; + if (digit < 0 || digit > 9) { + return Long.parseLong(new String(b, from, to - from, UTF_8)); + } + value = value * 10 + digit; + } + return negative ? -value : value; + } + + /// `Boolean.parseBoolean` on the field's text: `true` ignoring ASCII case. Text is compared + /// only when it could match, i.e. when it is 4 bytes, or longer and not ASCII (a non-ASCII + /// character could case-fold to one of `true`'s letters). + static boolean parseBoolean(byte[] b, int from, int to) { + int length = to - from; + if (length == 4) { + return (b[from] | 0x20) == 't' && (b[from + 1] | 0x20) == 'r' + && (b[from + 2] | 0x20) == 'u' && (b[from + 3] | 0x20) == 'e'; + } + if (length > 4 && !isAscii(b, from, to)) { + return Boolean.parseBoolean(new String(b, from, length, UTF_8)); + } + return false; + } + + /// The field's bytes as a Utf8 value: copied when well-formed UTF-8, else decoded and encoded + /// back, which turns each malformed sequence into U+FFFD as FastCSV's decoding did. + static byte[] utf8(byte[] b, int from, int to) { + if (from == to) { + return EMPTY; + } + if (isAscii(b, from, to) || Utf8Validation.isWellFormed(b, from, to)) { + return Arrays.copyOfRange(b, from, to); + } + return new String(b, from, to - from, UTF_8).getBytes(UTF_8); + } + + /// Whether no byte of `b[from, to)` has its high bit set, eight bytes at a time. + static boolean isAscii(byte[] b, int from, int to) { + int i = from; + long high = 0; + for (; i + Long.BYTES <= to; i += Long.BYTES) { + high |= (long) LONGS.get(b, i); + } + for (; i < to; i++) { + high |= b[i]; + } + return (high & 0x8080_8080_8080_8080L) == 0; + } + + private static final VarHandle LONGS = MethodHandles.byteArrayViewVarHandle(long[].class, ByteOrder.LITTLE_ENDIAN); + + /// Powers of ten a double holds exactly: 10^22 is the last. + private static final double[] EXACT_POWERS_OF_TEN = { + 1e0, 1e1, 1e2, 1e3, 1e4, 1e5, 1e6, 1e7, 1e8, 1e9, 1e10, 1e11, + 1e12, 1e13, 1e14, 1e15, 1e16, 1e17, 1e18, 1e19, 1e20, 1e21, 1e22}; + + /// `Double.parseDouble` on the field's text, `0.0` for an empty field. + /// + /// A plain decimal (`-12.345`) whose digits fit 2^53 is computed as `digits / 10^fractionDigits`: + /// both operands are exact doubles and IEEE division rounds correctly, so the result is the + /// correctly rounded value `parseDouble` returns (Clinger's fast path). Anything else (an + /// exponent, more digits, `NaN`, spaces, a bare `.5`) takes `parseDouble` itself. + static double parseDouble(byte[] b, int from, int to) { + if (from == to) { + return 0.0; + } + int i = from; + boolean negative = b[i] == '-'; + if (negative || b[i] == '+') { + i++; + } + long digits = 0; + int count = 0; + int fraction = -1; + for (; i < to; i++) { + int c = b[i]; + if (c == '.' && fraction < 0 && count > 0) { + fraction = 0; + continue; + } + int digit = c - '0'; + if (digit < 0 || digit > 9 || ++count > 18) { + return Double.parseDouble(new String(b, from, to - from, UTF_8)); + } + digits = digits * 10 + digit; + if (fraction >= 0) { + fraction++; + } + } + if (count == 0 || fraction == 0 || fraction > 22 || digits > (1L << 53)) { + return Double.parseDouble(new String(b, from, to - from, UTF_8)); } - return row[idx]; + double value = fraction > 0 ? digits / EXACT_POWERS_OF_TEN[fraction] : digits; + return negative ? -value : value; } } diff --git a/csv/src/main/java/io/github/dfa1/vortex/csv/ImportOptions.java b/csv/src/main/java/io/github/dfa1/vortex/csv/ImportOptions.java index 77a05b0d3..1316e8a4a 100644 --- a/csv/src/main/java/io/github/dfa1/vortex/csv/ImportOptions.java +++ b/csv/src/main/java/io/github/dfa1/vortex/csv/ImportOptions.java @@ -85,7 +85,8 @@ public ImportOptions withSchema(DType.Struct overrideSchema) { /// Override the field delimiter. /// - /// @param separator the delimiter character + /// @param separator the delimiter character: ASCII, and not a quote or line break (the import + /// reads bytes, where a non-ASCII character spans several) /// @return a copy of this options with the delimiter applied public ImportOptions withDelimiter(char separator) { return new ImportOptions(separator, chunkSize, header, schema, progressListener, writeOptions); diff --git a/csv/src/test/java/io/github/dfa1/vortex/csv/CsvByteReaderTest.java b/csv/src/test/java/io/github/dfa1/vortex/csv/CsvByteReaderTest.java new file mode 100644 index 000000000..9146f0688 --- /dev/null +++ b/csv/src/test/java/io/github/dfa1/vortex/csv/CsvByteReaderTest.java @@ -0,0 +1,175 @@ +package io.github.dfa1.vortex.csv; + +import de.siegmar.fastcsv.reader.CsvReader; +import de.siegmar.fastcsv.reader.StringArrayHandler; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.HexFormat; +import java.util.List; +import java.util.Random; +import java.util.stream.LongStream; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.catchThrowable; + +/// CsvByteReader replaced FastCSV in the import, so it must read every file the way FastCSV's +/// default reader did. FastCSV is the oracle: random files mixing quoted fields, escaped quotes, +/// every line ending, empty lines, multi-byte and malformed UTF-8, read by both. +class CsvByteReaderTest { + + static LongStream seeds() { + return LongStream.rangeClosed(1, 200); + } + + @ParameterizedTest + @MethodSource("seeds") + void readsRecordsAsFastCsv(long seed) { + // Given: a file FastCSV rejects must fail in both + Random random = new Random(seed); + char delimiter = random.nextBoolean() ? ',' : ';'; + byte[] csv = randomCsv(random, delimiter); + Object expected = outcome(() -> fastCsv(csv, delimiter)); + + // When: a tiny read size, so records and quotes straddle refills + Object result = outcome(() -> byteReader(new TrickleInputStream(csv, 1 + random.nextInt(7)), delimiter)); + + // Then + assertThat(result).as(HexFormat.of().formatHex(csv)).isEqualTo(expected); + } + + @ParameterizedTest + @ValueSource(strings = { + "a,b\n1,2,3\n", // more fields than the first record + "a,b\n1\n", // fewer + "a,b\n\"x\"y,2\n", // text after a closing quote + }) + void rejectsWhatFastCsvRejects(String csv) { + // Given + byte[] bytes = csv.getBytes(StandardCharsets.UTF_8); + Throwable fastCsvFailure = catchThrowable(() -> fastCsv(bytes, ',')); + + // When + Throwable result = catchThrowable(() -> byteReader(new ByteArrayInputStream(bytes), ',')); + + // Then + assertThat(fastCsvFailure).as("FastCSV rejects it too").isNotNull(); + assertThat(result).isInstanceOf(IOException.class); + } + + @Test + void aNonAsciiDelimiterIsRefused() { + // When + Throwable result = catchThrowable(() -> new CsvByteReader(InputStream.nullInputStream(), '§')); + + // Then: a delimiter byte inside a multi-byte character would split it + assertThat(result).isInstanceOf(IllegalArgumentException.class); + } + + private static byte[] randomCsv(Random random, char delimiter) { + String[] pieces = {"a", "bc", "1", "-2.5", " ", "\"", "\n", "\r", String.valueOf(delimiter), "é", "€", "😀", ""}; + String[] endings = {"\n", "\r\n", "\r"}; + ByteArrayOutputStream out = new ByteArrayOutputStream(); + if (random.nextInt(10) == 0) { + out.writeBytes(new byte[]{(byte) 0xef, (byte) 0xbb, (byte) 0xbf}); + } + int columns = 1 + random.nextInt(4); + int records = random.nextInt(30); + for (int r = 0; r < records; r++) { + if (random.nextInt(8) == 0) { + out.writeBytes(endings[random.nextInt(endings.length)].getBytes(StandardCharsets.UTF_8)); + } + for (int c = 0; c < columns; c++) { + if (c > 0) { + out.write(delimiter); + } + StringBuilder field = new StringBuilder(); + for (int k = random.nextInt(4); k > 0; k--) { + field.append(pieces[random.nextInt(pieces.length)]); + } + String text = field.toString(); + boolean mustQuote = text.indexOf(delimiter) >= 0 || text.contains("\"") + || text.contains("\n") || text.contains("\r"); + if (mustQuote || random.nextInt(5) == 0) { + out.writeBytes(("\"" + text.replace("\"", "\"\"") + "\"").getBytes(StandardCharsets.UTF_8)); + } else { + out.writeBytes(text.getBytes(StandardCharsets.UTF_8)); + if (random.nextInt(10) == 0) { + // malformed UTF-8 ending the field: a lone lead byte, or a stray continuation + out.write(random.nextBoolean() ? 0xc3 : 0x80); + } + } + } + if (r < records - 1 || random.nextBoolean()) { + out.writeBytes(endings[random.nextInt(endings.length)].getBytes(StandardCharsets.UTF_8)); + } + } + return out.toByteArray(); + } + + /// The records read, or `"rejected"` when reading fails: which exception differs by reader. + private static Object outcome(Read read) { + try { + return read.records(); + } catch (IOException | RuntimeException _) { + return "rejected"; + } + } + + @FunctionalInterface + private interface Read { + List> records() throws IOException; + } + + private static List> fastCsv(byte[] csv, char delimiter) throws IOException { + List> records = new ArrayList<>(); + try (CsvReader reader = CsvReader.builder().fieldSeparator(delimiter) + .build(StringArrayHandler.of(), new ByteArrayInputStream(csv))) { + for (String[] fields : reader) { + records.add(List.of(fields)); + } + } + return records; + } + + private static List> byteReader(InputStream csv, char delimiter) throws IOException { + ColumnChunk chunk = new ColumnChunk(10_000); + try (CsvByteReader reader = new CsvByteReader(csv, delimiter)) { + while (reader.readRecord(chunk)) { + // read every record + } + } + List> records = new ArrayList<>(); + for (int row = 0; row < chunk.rows(); row++) { + List fields = new ArrayList<>(); + for (int c = 0; c < chunk.columns(); c++) { + fields.add(chunk.text(c, row)); + } + records.add(fields); + } + return records; + } + + /// Returns at most `step` bytes per read, as a slow network body would. + private static final class TrickleInputStream extends ByteArrayInputStream { + private final int step; + + TrickleInputStream(byte[] bytes, int step) { + super(bytes); + this.step = step; + } + + @Override + public synchronized int read(byte[] b, int off, int len) { + return super.read(b, off, Math.min(len, step)); + } + } +} diff --git a/csv/src/test/java/io/github/dfa1/vortex/csv/CsvImporterTest.java b/csv/src/test/java/io/github/dfa1/vortex/csv/CsvImporterTest.java index d317e70d4..03680b2df 100644 --- a/csv/src/test/java/io/github/dfa1/vortex/csv/CsvImporterTest.java +++ b/csv/src/test/java/io/github/dfa1/vortex/csv/CsvImporterTest.java @@ -11,13 +11,20 @@ import io.github.dfa1.vortex.writer.WriteOptions; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.Path; -import java.util.List; import java.util.concurrent.ForkJoinPool; +import java.util.function.DoubleSupplier; +import java.util.function.LongSupplier; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.catchThrowable; class CsvImporterTest { @@ -119,22 +126,192 @@ void respectsSchemaOverride(@TempDir Path tmp) throws Exception { } @Test - void buildChunkOnAnExecutorKeepsTheColumnsAndTheirOrder() { + void chunksWrittenInTheBackgroundKeepFileOrder(@TempDir Path tmp) throws Exception { + // Given: two full chunks and a partial one, so each is converted while the next is parsed + int rows = 2 * ImportOptions.defaults().chunkSize() + 123; + StringBuilder text = new StringBuilder("id\n"); + for (int i = 0; i < rows; i++) { + text.append(i).append('\n'); + } + Path csv = tmp.resolve("ordered.csv"); + Files.writeString(csv, text); + Path vortex = tmp.resolve("ordered.vortex"); + + // When + CsvImporter.importCsv(csv, vortex); + + // Then: every row back once, in file order + long expected = 0; + try (VortexReader reader = VortexReader.open(vortex); + ScanIterator iter = reader.scan(ScanOptions.all())) { + while (iter.hasNext()) { + try (Chunk chunk = iter.next()) { + LongArray ids = chunk.column(ColumnName.of("id")); + for (int i = 0; i < chunk.rowCount(); i++) { + assertThat(ids.getLong(i)).isEqualTo(expected++); + } + } + } + } + assertThat(expected).isEqualTo(rows); + } + + @Test + void aBadValueAfterTheFirstChunkFailsWithItsOwnException(@TempDir Path tmp) throws Exception { + // Given: the first chunk infers a long column; a later chunk, converted on the background + // thread, holds a value that is not one + int rows = ImportOptions.defaults().chunkSize() + 10; + StringBuilder text = new StringBuilder("qty\n"); + for (int i = 0; i < rows; i++) { + text.append(i == rows - 1 ? "N/A" : String.valueOf(i)).append('\n'); + } + Path csv = tmp.resolve("late.csv"); + Files.writeString(csv, text); + Path vortex = tmp.resolve("late.vortex"); + + // When + Throwable result = catchThrowable(() -> CsvImporter.importCsv(csv, vortex)); + + // Then: the conversion's own exception, not the executor's wrapper around it + assertThat(result).isInstanceOf(NumberFormatException.class).hasMessageContaining("N/A"); + } + + @ParameterizedTest + @ValueSource(strings = {"0", "-0", "+5", "-42", "123456789012345678", "1234567890123456789", + "9223372036854775807", "-9223372036854775808", "9223372036854775808", "1_0", " 1", "1e3", + "-", "+", "--1", "١٢٣", "12a"}) + void parseLong_agreesWithLongParseLong(String text) { + // Given: 18 digits stay on the byte path, longer, signed-only and non-ASCII digits fall back + byte[] bytes = text.getBytes(StandardCharsets.UTF_8); + Object expected = parsed(() -> Long.parseLong(text)); + + // When + Object result = parsed(() -> CsvImporter.parseLong(bytes, 0, bytes.length)); + + // Then + assertThat(result).isEqualTo(expected); + } + + @ParameterizedTest + @ValueSource(strings = {"0", "-0", "-0.0", "+1.5", "364.99", "007.50", "1.", ".5", ".", "-", "1.2.3", "1e3", + " 1.5", "NaN", "Infinity", "0x1p3", "1.5d", "9007199254740993", "9007199254740992.5", + "123456789012345678", "1234567890123456789", "0.1234567890123456789012", "3.14159265358979323846"}) + void parseDouble_agreesWithDoubleParseDouble(String text) { + // Given: the fast path's edges (2^53, 22 fraction digits, 18 digits) and everything it hands back + byte[] bytes = text.getBytes(StandardCharsets.UTF_8); + Object expected = parsedDouble(() -> Double.parseDouble(text)); + + // When + Object result = parsedDouble(() -> CsvImporter.parseDouble(bytes, 0, bytes.length)); + + // Then + assertThat(result).isEqualTo(expected); + } + + @ParameterizedTest + @ValueSource(longs = {1, 2, 3, 4, 5}) + void parseDouble_agreesWithDoubleParseDoubleOnRandomDecimals(long seed) { + // Given: decimals of up to 18 digits with up to 22 fraction digits, bit-for-bit + java.util.Random random = new java.util.Random(seed); + for (int k = 0; k < 4_000; k++) { + long digits = random.nextLong() >>> (1 + random.nextInt(63)); + int fraction = random.nextInt(23); + String text = new java.math.BigDecimal(java.math.BigInteger.valueOf(digits), fraction).toPlainString(); + if (random.nextBoolean()) { + text = "-" + text; + } + byte[] bytes = text.getBytes(StandardCharsets.UTF_8); + + // When + double result = CsvImporter.parseDouble(bytes, 0, bytes.length); + + // Then + assertThat(Double.doubleToRawLongBits(result)).as(text) + .isEqualTo(Double.doubleToRawLongBits(Double.parseDouble(text))); + } + } + + @ParameterizedTest + @ValueSource(strings = {"", "a", "abcdefg", "abcdefgh", "abcdefghi", "abcdefghé", "éabcdefghijklmnop"}) + void isAscii_findsAHighByteAnywhere(String text) { + // Given: lengths around the 8-byte word, the high byte in the word part and in the tail + byte[] bytes = text.getBytes(StandardCharsets.UTF_8); + + // When + boolean result = CsvImporter.isAscii(bytes, 0, bytes.length); + + // Then + assertThat(result).isEqualTo(text.chars().allMatch(c -> c < 0x80)); + } + + @ParameterizedTest + @ValueSource(strings = {"true", "TRUE", "tRuE", "true ", "false", "t", "truee", "trué", "true"}) + void parseBoolean_agreesWithBooleanParseBoolean(String text) { + // Given + byte[] bytes = text.getBytes(StandardCharsets.UTF_8); + boolean expected = Boolean.parseBoolean(text); + + // When + boolean result = CsvImporter.parseBoolean(bytes, 0, bytes.length); + + // Then + assertThat(result).isEqualTo(expected); + } + + @Test + void utf8_replacesMalformedBytesAsTheDecoderDoes() { + // Given: a lone lead byte between ASCII, which FastCSV's decoding turned into U+FFFD + byte[] bytes = {'a', (byte) 0xc3, 'b'}; + + // When + byte[] result = CsvImporter.utf8(bytes, 0, bytes.length); + + // Then + assertThat(new String(result, StandardCharsets.UTF_8)).isEqualTo("a�b"); + } + + @Test + void buildChunkOnAnExecutorKeepsTheColumnsAndTheirOrder() throws IOException { // Given: a column of each kind, so every conversion path runs off the caller thread DType.Struct schema = DType.structBuilder() .field(ColumnName.of("id"), DType.I64) .field(ColumnName.of("price"), DType.F64) .field(ColumnName.of("name"), DType.UTF8) .build(); - List rows = List.of(new String[]{"1", "1.5", "a"}, new String[]{"2", "2.5", "b"}); + ColumnChunk rows = new ColumnChunk(2); + try (CsvByteReader reader = new CsvByteReader( + new ByteArrayInputStream("1,1.5,a\n2,2.5,b\n".getBytes(StandardCharsets.UTF_8)), ',')) { + while (reader.readRecord(rows)) { + // read both rows + } + } // When var result = CsvImporter.buildChunk(schema, rows, ForkJoinPool.commonPool()); - // Then + // Then: strings come out as their UTF-8 bytes, which the writer takes for a Utf8 column assertThat(result.keySet()).extracting(ColumnName::value).containsExactly("id", "price", "name"); assertThat((long[]) result.get(ColumnName.of("id"))).containsExactly(1L, 2L); assertThat((double[]) result.get(ColumnName.of("price"))).containsExactly(1.5, 2.5); - assertThat((String[]) result.get(ColumnName.of("name"))).containsExactly("a", "b"); + assertThat((byte[][]) result.get(ColumnName.of("name"))).isDeepEqualTo(new byte[][]{ + "a".getBytes(StandardCharsets.UTF_8), "b".getBytes(StandardCharsets.UTF_8)}); + } + + /// The parsed value's bits, or the exception's class and message when parsing fails. + private static Object parsedDouble(DoubleSupplier parse) { + try { + return Double.doubleToRawLongBits(parse.getAsDouble()); + } catch (NumberFormatException e) { + return e.getClass().getName() + ": " + e.getMessage(); + } + } + + /// The parsed value, or the exception's class and message when parsing fails. + private static Object parsed(LongSupplier parse) { + try { + return parse.getAsLong(); + } catch (NumberFormatException e) { + return e.getClass().getName() + ": " + e.getMessage(); + } } } diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/Repartitioner.java b/writer/src/main/java/io/github/dfa1/vortex/writer/Repartitioner.java index d232d7818..43df3c1f1 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/Repartitioner.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/Repartitioner.java @@ -6,6 +6,7 @@ import io.github.dfa1.vortex.writer.encode.NullableData; import java.lang.reflect.Array; +import java.nio.charset.StandardCharsets; import java.util.ArrayDeque; import java.util.ArrayList; import java.util.Arrays; @@ -174,7 +175,9 @@ private long nbytes(Object data, int offset, int rows) { case String[] a -> { long n = 16L * rows; for (int i = offset; i < offset + rows; i++) { - n += a[i] == null ? 0 : utf8Length(a[i]); + // Encoding to count beats a char-by-char count by 2x: the JDK encodes ASCII + // with a vectorized copy, and the throwaway array dies young. + n += a[i] == null ? 0 : a[i].getBytes(StandardCharsets.UTF_8).length; } yield n; } @@ -190,30 +193,6 @@ private long nbytes(Object data, int offset, int rows) { }; } - /// The length of `s` in UTF-8, counted without encoding it: sizing a batch must not cost a copy - /// of every string, as it runs on the caller's thread, ahead of the parallel compression. - /// An unpaired surrogate is one byte, as `getBytes` encodes it (`?`). - static int utf8Length(String s) { - int n = s.length(); - int bytes = n; - for (int i = 0; i < n; i++) { - char c = s.charAt(i); - if (c >= 0x800) { - if (Character.isHighSurrogate(c) && i + 1 < n && Character.isLowSurrogate(s.charAt(i + 1))) { - bytes += 2; - i++; - } else if (Character.isSurrogate(c)) { - continue; - } else { - bytes += 2; - } - } else if (c >= 0x80) { - bytes++; - } - } - return bytes; - } - private long elementBytes() { DType storage = dtype instanceof DType.Extension ext ? ext.storageDType() : dtype; return switch (storage) { diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/ZoneAccumulator.java b/writer/src/main/java/io/github/dfa1/vortex/writer/ZoneAccumulator.java index 587ec254b..7efa6b3d3 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/ZoneAccumulator.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/ZoneAccumulator.java @@ -276,25 +276,38 @@ private void zone(Object data, int off, int len) { Object min = null; Object max = null; boolean unknown = false; + // INT, FLOAT and BOOL track their bounds as primitives and box once per zone: boxing on every + // new bound allocated a Long per row of an ascending column such as a row id. switch (kind) { case INT -> { long[] longs = PrimitiveArrays.toLongs(slice(values, null, off, len), ptype, EncodingId.VORTEX_PRIMITIVE); boolean unsigned = ptype.isUnsigned(); + boolean seen = false; + long lo = 0; + long hi = 0; for (int i = 0; i < len; i++) { if (validity != null && !validity[off + i]) { nulls++; continue; } long v = longs[i]; - if (min == null || compare(v, (Long) min, unsigned) < 0) { - min = v; + if (!seen || compare(v, lo, unsigned) < 0) { + lo = v; } - if (max == null || compare(v, (Long) max, unsigned) > 0) { - max = v; + if (!seen || compare(v, hi, unsigned) > 0) { + hi = v; } + seen = true; + } + if (seen) { + min = lo; + max = hi; } } case F16, FLOAT -> { + boolean seen = false; + double lo = 0; + double hi = 0; for (int i = 0; i < len; i++) { if (validity != null && !validity[off + i]) { nulls++; @@ -309,29 +322,40 @@ private void zone(Object data, int off, int len) { nans++; continue; } - if (min == null || v < (Double) min) { - min = v; + if (!seen || v < lo) { + lo = v; } - if (max == null || v > (Double) max) { - max = v; + if (!seen || v > hi) { + hi = v; } + seen = true; } - if (kind == Kind.F16) { + if (seen && kind == Kind.F16) { // The F16 table column is stored as raw half-float bits like the data. - min = min == null ? null : (long) Float.floatToFloat16(((Double) min).floatValue()); - max = max == null ? null : (long) Float.floatToFloat16(((Double) max).floatValue()); + min = (long) Float.floatToFloat16((float) lo); + max = (long) Float.floatToFloat16((float) hi); + } else if (seen) { + min = lo; + max = hi; } } case BOOL -> { boolean[] a = (boolean[]) values; + boolean seen = false; + boolean lo = true; + boolean hi = false; for (int i = 0; i < len; i++) { if (validity != null && !validity[off + i]) { nulls++; continue; } - boolean v = a[off + i]; - min = min == null ? v : (Boolean) min && v; - max = max == null ? v : (Boolean) max || v; + lo &= a[off + i]; + hi |= a[off + i]; + seen = true; + } + if (seen) { + min = lo; + max = hi; } } case DECIMAL -> { diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/OnPairEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/OnPairEncodingEncoder.java index c7885b4a8..e6bdac7af 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/OnPairEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/OnPairEncodingEncoder.java @@ -308,10 +308,16 @@ void entryCreated() { } } - /// Byte trie over the dictionary tokens for greedy longest-prefix match. Edges live in one - /// [LongIntMap] keyed by `(node << 8) | byte`; node 0 is the root. + /// Byte trie over the dictionary tokens for greedy longest-prefix match; node 0 is the root. + /// + /// Every match steps through the root's child and nearly every one through a grandchild, so + /// those edges are direct-indexed by the first byte and by the first two (`node + 1`, so `0` + /// is absent); deeper edges live in an open-addressing table keyed by `(node << 8) | byte`. + /// Hashing every step was 30% of a default-options CSV import, most of it cache misses. private static final class Trie { - private final LongIntMap edges = new LongIntMap(8192); + private final int[] depth1 = new int[1 << 8]; + private final int[] depth2 = new int[1 << 16]; + private final EdgeTable edges = new EdgeTable(); private int[] nodeToken = {-1}; private int nodes = 1; private final List tokens = new ArrayList<>(); @@ -322,15 +328,14 @@ private static final class Trie { int insert(byte[] src, int from, int to) { int node = 0; for (int i = from; i < to; i++) { - long key = ((long) node << 8) | (src[i] & 0xFF); - int child = edges.get(key); + int child = child(node, src, from, i); if (child < 0) { child = nodes++; if (child == nodeToken.length) { nodeToken = Arrays.copyOf(nodeToken, 2 * child); Arrays.fill(nodeToken, child, nodeToken.length, -1); } - edges.put(key, child); + setChild(node, src, from, i, child); } node = child; } @@ -349,10 +354,28 @@ int size() { /// byte is a token, so the match is at least one byte long. long longestMatch(byte[] src, int from, int to) { int limit = Math.min(to, from + Encoder.MAX_TOKEN_SIZE); - int node = 0; long best = 0; - for (int i = from; i < limit; i++) { - node = edges.get(((long) node << 8) | (src[i] & 0xFF)); + // the first two steps peeled off the loop: direct-indexed, and no per-step branch on depth + int first = src[from] & 0xFF; + int node = depth1[first] - 1; + if (node < 0) { + return best; + } + if (nodeToken[node] >= 0) { + best = match(nodeToken[node], 1); + } + if (from + 1 >= limit) { + return best; + } + node = depth2[(first << 8) | (src[from + 1] & 0xFF)] - 1; + if (node < 0) { + return best; + } + if (nodeToken[node] >= 0) { + best = match(nodeToken[node], 2); + } + for (int i = from + 2; i < limit; i++) { + node = edges.get((node << 8) | (src[i] & 0xFF)); if (node < 0) { break; } @@ -363,6 +386,23 @@ long longestMatch(byte[] src, int from, int to) { return best; } + /// The child of `node`, reached by `src[i]` from a path that started at `src[from]`, or -1. + private int child(int node, byte[] src, int from, int i) { + return switch (i - from) { + case 0 -> depth1[src[i] & 0xFF] - 1; + case 1 -> depth2[((src[from] & 0xFF) << 8) | (src[i] & 0xFF)] - 1; + default -> edges.get((node << 8) | (src[i] & 0xFF)); + }; + } + + private void setChild(int node, byte[] src, int from, int i, int child) { + switch (i - from) { + case 0 -> depth1[src[i] & 0xFF] = child + 1; + case 1 -> depth2[((src[from] & 0xFF) << 8) | (src[i] & 0xFF)] = child + 1; + default -> edges.put((node << 8) | (src[i] & 0xFF), child); + } + } + /// The tokens in ascending unsigned-lexicographic order. byte[][] sortedTokens() { byte[][] sorted = tokens.toArray(new byte[0][]); @@ -382,4 +422,65 @@ static int length(long match) { return (int) (match & 0xFF); } } + + /// The trie's deep edges: open addressing over one `long[]` of `(key << 32) | (child + 1)`, so a + /// lookup reads one slot where [io.github.dfa1.vortex.core.compute.LongIntMap]'s parallel key + /// and value arrays read two, and these lookups are cache misses. A key is `(node << 8) | byte` + /// with node below 2^23 (at most 4096 tokens of 16 bytes), so it fits the upper half. + private static final class EdgeTable { + private static final long HASH_MULTIPLIER = 0x9E3779B97F4A7C15L; + + private long[] slots = new long[1 << 14]; + private int shift = Long.numberOfLeadingZeros(slots.length - 1L); + private int size; + + /// The child for `key`, or -1. + int get(int key) { + int mask = slots.length - 1; + int slot = slotOf(key); + while (true) { + long entry = slots[slot]; + if (entry == 0) { + return -1; + } + if ((int) (entry >>> 32) == key) { + return (int) entry - 1; + } + slot = (slot + 1) & mask; + } + } + + /// Adds the edge `key` -> `child`, which is not present yet. + void put(int key, int child) { + if (2 * (size + 1) > slots.length) { + grow(); + } + insert((long) key << 32 | (child + 1L)); + size++; + } + + private void insert(long entry) { + int mask = slots.length - 1; + int slot = slotOf((int) (entry >>> 32)); + while (slots[slot] != 0) { + slot = (slot + 1) & mask; + } + slots[slot] = entry; + } + + private void grow() { + long[] old = slots; + slots = new long[old.length * 4]; + shift = Long.numberOfLeadingZeros(slots.length - 1L); + for (long entry : old) { + if (entry != 0) { + insert(entry); + } + } + } + + private int slotOf(int key) { + return (int) ((key * HASH_MULTIPLIER) >>> shift); + } + } } diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/VarBinEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/VarBinEncodingEncoder.java index daccbf8d6..1d3617274 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/VarBinEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/VarBinEncodingEncoder.java @@ -184,6 +184,29 @@ static MinMax minMaxStats(byte[][] utf8) { new String(utf8[max], StandardCharsets.UTF_8)); } + /// [#minMaxStats(String[])] over each string's UTF-8 bytes, for an encoder that already holds + /// them one array per row; see [#minMaxStats(String[], VarBinBytes.Rows)] for why bytes. + /// + /// @param strings the string values (may contain `null`) + /// @param utf8 `strings` encoded row for row, `null` where `strings` is + /// @return the [MinMax] of encoded scalars, or `null` if all entries are `null` + static MinMax minMaxStats(String[] strings, byte[][] utf8) { + int min = -1; + int max = -1; + for (int i = 0; i < utf8.length; i++) { + if (utf8[i] == null) { + continue; + } + if (min < 0 || Arrays.compareUnsigned(utf8[i], utf8[min]) < 0) { + min = i; + } + if (max < 0 || Arrays.compareUnsigned(utf8[i], utf8[max]) > 0) { + max = i; + } + } + return min < 0 ? null : encodeMinMax(strings[min], strings[max]); + } + /// Encodes an already-computed string min/max as [#minMaxStats(String[])] would. /// /// @param minStr the minimum in UTF-8 order, or `null` when there is no non-null value diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ZstdEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ZstdEncodingEncoder.java index 8eadac46c..97d1d7b50 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ZstdEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ZstdEncodingEncoder.java @@ -112,20 +112,19 @@ public boolean acceptsNullable(DType dtype) { public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { requireBinding(ZSTD_BINDING_PRESENT); // Compression does not reorder anything, so the column's bounds are the input's bounds. - // Computed once here rather than in each of the four shape-specific paths below. - MinMax stats = ZoneMapStats.of(dtype, data); if (data instanceof NullableData nd) { if (dtype instanceof DType.Primitive dt) { - return encodeNullablePrimitive(dt, nd, ctx).withStats(stats); + return encodeNullablePrimitive(dt, nd, ctx).withStats(ZoneMapStats.of(dtype, data)); } if (dtype instanceof DType.Utf8 || dtype instanceof DType.Binary) { - return encodeNullableVarBin(nd, ctx).withStats(stats); + byte[][] encoded = VarBinBytes.toRawByteArrays(nd.values()); + return encodeNullableVarBin(encoded, nd.validity(), ctx).withStats(textStats(dtype, nd.values(), encoded)); } throw new VortexException(EncodingId.VORTEX_ZSTD, "NullableData is unsupported for dtype: " + dtype); } if (dtype instanceof DType.Primitive dt) { - return encodePrimitive(dt, data, ctx.arena()).withStats(stats); + return encodePrimitive(dt, data, ctx.arena()).withStats(ZoneMapStats.of(dtype, data)); } if (dtype instanceof DType.Utf8 || dtype instanceof DType.Binary) { byte[][] encoded = VarBinBytes.toRawByteArrays(data); @@ -133,7 +132,7 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { throw new VortexException(EncodingId.VORTEX_ZSTD, "non-nullable " + dtype + " contains null"); } - return encodeVarBin(encoded, ctx.arena()).withStats(stats); + return encodeVarBin(encoded, ctx.arena()).withStats(textStats(dtype, data, encoded)); } throw new VortexException(EncodingId.VORTEX_ZSTD, "unsupported dtype: " + dtype); } @@ -170,12 +169,23 @@ private EncodeResult encodeNullablePrimitive(DType.Primitive dt, NullableData nd return buildNullableResult(packed, uniformLayout(countValid(validity), byteWidth), validity, ctx); } - private EncodeResult encodeNullableVarBin(NullableData nd, EncodeContext ctx) { + private EncodeResult encodeNullableVarBin(byte[][] encoded, boolean[] validity, EncodeContext ctx) { // Strip null positions: only valid rows reach the compressed payload (mirrors the Rust // reference). The decoder scatters them back over the validity mask carried by child[0]. - byte[][] valid = stripNulls(VarBinBytes.toRawByteArrays(nd.values())); + byte[][] valid = stripNulls(encoded); MemorySegment packed = buildLengthPrefixed(valid, ctx.arena()); - return buildNullableResult(packed, varBinLayout(packed, valid.length), nd.validity(), ctx); + return buildNullableResult(packed, varBinLayout(packed, valid.length), validity, ctx); + } + + /// The zone-map bounds of a string column, ordered by the UTF-8 bytes already encoded for the + /// payload rather than by comparing the `String`s again: the vectorized byte compare is several + /// times faster on values sharing a prefix (timestamps, ids), where the `String` compare was the + /// largest single cost of a compact CSV import. Binary has no useful ordering, as in [ZoneMapStats]. + private static MinMax textStats(DType dtype, Object values, byte[][] encoded) { + if (values instanceof String[] strings) { + return VarBinEncodingEncoder.minMaxStats(strings, encoded); + } + return VarBinEncodingEncoder.textMinMax(dtype, values); } private EncodeResult buildNullableResult( diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/RepartitionerTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/RepartitionerTest.java index c446cac2f..49cee3f7b 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/RepartitionerTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/RepartitionerTest.java @@ -4,10 +4,7 @@ import io.github.dfa1.vortex.core.model.PType; import io.github.dfa1.vortex.writer.encode.NullableData; import org.junit.jupiter.api.Test; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.ValueSource; -import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.List; import java.util.stream.LongStream; @@ -108,19 +105,6 @@ private static long[] concat(List chunks) { return chunks.stream().flatMapToLong(c -> LongStream.of((long[]) c)).toArray(); } - @ParameterizedTest - @ValueSource(strings = {"", "ascii", "caf\u00e9", "\u20ac uro", "\uD83D\uDE00 emoji", "lone \uD83D high", "lone \uDE00 low", "\uD83D", "\uDE00\uD83D"}) - void utf8LengthMatchesTheEncoder(String text) { - // Given: ASCII, 2/3/4-byte characters and unpaired surrogates, which getBytes encodes as '?' - int expected = text.getBytes(StandardCharsets.UTF_8).length; - - // When - int result = Repartitioner.utf8Length(text); - - // Then - assertThat(result).isEqualTo(expected); - } - @Test void addWithPresizedPiecesCutsTheSameChunks() { // Given: strings long enough that a batch crosses the 1 MiB block target, so cuts depend on the sizes diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/VarBinEncodingEncoderTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/VarBinEncodingEncoderTest.java index ebe60b412..922f0d98a 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/VarBinEncodingEncoderTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/VarBinEncodingEncoderTest.java @@ -234,9 +234,34 @@ class Stats { @ParameterizedTest @ValueSource(longs = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10}) void minMaxStats_overRowBytes_matchesUtf8OrderOverStrings(long seed) { - // Given — the byte compare (#565) must pick the same min/max as Utf8Order on the - // strings. Mixed ASCII, U+E000.. and surrogate pairs, where UTF-16 order disagrees with - // UTF-8, plus shared prefixes and nulls (an empty row in the bytes, but not a value). + // Given — the byte compare (#565) must pick the same min/max as Utf8Order on the strings + String[] strings = mixedOrderStrings(seed); + MinMax expected = VarBinEncodingEncoder.minMaxStats(strings); + + // When + MinMax result = VarBinEncodingEncoder.minMaxStats(strings, VarBinBytes.toContiguous(strings)); + + // Then + assertThat(result).isEqualTo(expected); + } + + @ParameterizedTest + @ValueSource(longs = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10}) + void minMaxStats_overPerRowBytes_matchesUtf8OrderOverStrings(long seed) { + // Given — Zstd's per-row arrays, where a null row stays a null array + String[] strings = mixedOrderStrings(seed); + MinMax expected = VarBinEncodingEncoder.minMaxStats(strings); + + // When + MinMax result = VarBinEncodingEncoder.minMaxStats(strings, VarBinBytes.toRawByteArrays(strings)); + + // Then + assertThat(result).isEqualTo(expected); + } + + /// Mixed ASCII, U+E000.. and surrogate pairs, where UTF-16 order disagrees with UTF-8, plus + /// shared prefixes and nulls (an empty or null row in the bytes, but not a value). + private static String[] mixedOrderStrings(long seed) { Random random = new Random(seed); String[] alphabet = {"a", "b", "\u00e9", "\uE000", "\uFFFD", "\uD83D\uDE00", "\uD800\uDC00"}; String[] strings = new String[50]; @@ -250,13 +275,7 @@ void minMaxStats_overRowBytes_matchesUtf8OrderOverStrings(long seed) { } strings[i] = sb.toString(); } - MinMax expected = VarBinEncodingEncoder.minMaxStats(strings); - - // When - MinMax result = VarBinEncodingEncoder.minMaxStats(strings, VarBinBytes.toContiguous(strings)); - - // Then - assertThat(result).isEqualTo(expected); + return strings; } @Test