|
| 1 | +package io.github.dfa1.vortex.integration; |
| 2 | + |
| 3 | +import dev.vortex.api.Session; |
| 4 | +import dev.vortex.api.VortexWriter; |
| 5 | +import dev.vortex.arrow.ArrowAllocation; |
| 6 | +import dev.vortex.jni.NativeLoader; |
| 7 | +import io.github.dfa1.vortex.inspect.VortexInspector; |
| 8 | +import io.github.dfa1.vortex.reader.ReadRegistry; |
| 9 | +import io.github.dfa1.vortex.reader.ScanOptions; |
| 10 | +import io.github.dfa1.vortex.reader.VortexReader; |
| 11 | +import io.github.dfa1.vortex.reader.array.Array; |
| 12 | +import io.github.dfa1.vortex.reader.array.DoubleArray; |
| 13 | +import io.github.dfa1.vortex.reader.array.FloatArray; |
| 14 | +import io.github.dfa1.vortex.reader.array.MaskedArray; |
| 15 | +import org.apache.arrow.c.ArrowArray; |
| 16 | +import org.apache.arrow.c.ArrowSchema; |
| 17 | +import org.apache.arrow.c.Data; |
| 18 | +import org.apache.arrow.memory.BufferAllocator; |
| 19 | +import org.apache.arrow.vector.Float4Vector; |
| 20 | +import org.apache.arrow.vector.Float8Vector; |
| 21 | +import org.apache.arrow.vector.VectorSchemaRoot; |
| 22 | +import org.apache.arrow.vector.types.FloatingPointPrecision; |
| 23 | +import org.apache.arrow.vector.types.pojo.ArrowType; |
| 24 | +import org.apache.arrow.vector.types.pojo.Field; |
| 25 | +import org.apache.arrow.vector.types.pojo.Schema; |
| 26 | +import org.junit.jupiter.api.Test; |
| 27 | +import org.junit.jupiter.api.io.TempDir; |
| 28 | + |
| 29 | +import java.io.IOException; |
| 30 | +import java.nio.file.Path; |
| 31 | +import java.util.ArrayList; |
| 32 | +import java.util.HashMap; |
| 33 | +import java.util.List; |
| 34 | + |
| 35 | +import static org.assertj.core.api.Assertions.assertThat; |
| 36 | +import static org.junit.jupiter.api.Assumptions.assumeTrue; |
| 37 | + |
| 38 | +/// Ground-truth interop for `fastlanes.rle` over floating-point columns (issue #209). |
| 39 | +/// |
| 40 | +/// The Rust (JNI) writer compresses long constant runs of doubles to |
| 41 | +/// `fastlanes.rle`; before the fix the Java reader threw |
| 42 | +/// `fastlanes.rle: unsupported ptype F64`. These tests write such a column with |
| 43 | +/// the JNI writer, then confirm via [VortexInspector] that RLE was actually |
| 44 | +/// chosen ( `assumeTrue` — a test asserting on whatever encoding the compressor |
| 45 | +/// happened to pick would be meaningless), and finally assert the Java reader |
| 46 | +/// decodes every value exactly. The nullable case additionally exercises the |
| 47 | +/// validity re-wrap the F64 arm shares with the integer arms. |
| 48 | +class RleF64InteropIntegrationTest { |
| 49 | + |
| 50 | + private static final Session SESSION = Session.create(); |
| 51 | + private static final BufferAllocator ALLOCATOR = ArrowAllocation.rootAllocator(); |
| 52 | + |
| 53 | + private static final Schema F64_SCHEMA = new Schema(List.of( |
| 54 | + Field.notNullable("v", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)) |
| 55 | + )); |
| 56 | + private static final Schema F64_NULLABLE_SCHEMA = new Schema(List.of( |
| 57 | + Field.nullable("v", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)) |
| 58 | + )); |
| 59 | + private static final Schema F32_SCHEMA = new Schema(List.of( |
| 60 | + Field.notNullable("v", new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)) |
| 61 | + )); |
| 62 | + |
| 63 | + /// Number of rows: long constant runs (512-row blocks of one value) over a |
| 64 | + /// handful of distinct values is the canonical RLE-friendly shape and matches |
| 65 | + /// the real-world seoul weather columns from #209. |
| 66 | + private static final int ROWS = 8_192; |
| 67 | + private static final int RUN = 512; |
| 68 | + |
| 69 | + static { |
| 70 | + NativeLoader.loadJni(); |
| 71 | + } |
| 72 | + |
| 73 | + @Test |
| 74 | + void jniWriter_javaReader_f64RunLengthEncoded(@TempDir Path tmp) throws IOException { |
| 75 | + // Given — 8_192 doubles in 512-row constant runs cycling 5 distinct fractional |
| 76 | + // values; fractional values catch any accidental integer-narrowing decode bug. |
| 77 | + double[] unique = {-1.5, 2.25, 3.75, -4.125, 5.5}; |
| 78 | + double[] vals = new double[ROWS]; |
| 79 | + for (int i = 0; i < ROWS; i++) { |
| 80 | + vals[i] = unique[(i / RUN) % unique.length]; |
| 81 | + } |
| 82 | + Path file = tmp.resolve("rle_f64.vtx"); |
| 83 | + writeF64(file, vals); |
| 84 | + assumeRle(file); |
| 85 | + |
| 86 | + // When — Java reader decodes the whole column |
| 87 | + double[] result = readDoubleColumn(file); |
| 88 | + |
| 89 | + // Then — every value matches the input exactly |
| 90 | + assertThat(result).containsExactly(vals); |
| 91 | + } |
| 92 | + |
| 93 | + @Test |
| 94 | + void jniWriter_javaReader_f64NullableDecodesValues(@TempDir Path tmp) throws IOException { |
| 95 | + // Given — the same run shape but nullable, with whole 512-row runs set null. |
| 96 | + // Exercises the F64 RLE decode path on a nullable schema end-to-end. |
| 97 | + double[] unique = {10.5, 20.25, 30.75}; |
| 98 | + double[] vals = new double[ROWS]; |
| 99 | + boolean[] nulls = new boolean[ROWS]; |
| 100 | + for (int i = 0; i < ROWS; i++) { |
| 101 | + int block = i / RUN; |
| 102 | + nulls[i] = block % 3 == 1; // every third run is entirely null |
| 103 | + vals[i] = unique[block % unique.length]; |
| 104 | + } |
| 105 | + Path file = tmp.resolve("rle_f64_nullable.vtx"); |
| 106 | + writeF64Nullable(file, vals, nulls); |
| 107 | + assumeRle(file); |
| 108 | + |
| 109 | + // When — decode the whole column (unwrapping any validity mask) |
| 110 | + double[] result = readDoubleColumnUnwrapped(file); |
| 111 | + |
| 112 | + // Then — decode does not throw, row count is right, and every non-null input |
| 113 | + // value is exact. This test deliberately does not assert the null mask: the JNI |
| 114 | + // writer folds RLE validity into the indices child, and surfacing it as a |
| 115 | + // MaskedArray from the RLE decode path is the subject of the separate in-flight |
| 116 | + // fix #210 (the existing jniWriter_nullableColumn_decodesWithoutError follows the |
| 117 | + // same convention for the bitpacked case). |
| 118 | + assertThat(result).hasSize(ROWS); |
| 119 | + for (int i = 0; i < ROWS; i++) { |
| 120 | + if (!nulls[i]) { |
| 121 | + assertThat(result[i]).as("value@%d", i).isEqualTo(vals[i]); |
| 122 | + } |
| 123 | + } |
| 124 | + } |
| 125 | + |
| 126 | + @Test |
| 127 | + void jniWriter_javaReader_f32RunLengthEncoded(@TempDir Path tmp) throws IOException { |
| 128 | + // Given — the same run shape as F64 but 32-bit; guards the 4-byte value read. |
| 129 | + // Gated on RLE selection: the compressor may prefer another encoding for F32, |
| 130 | + // in which case the test is skipped rather than asserting on the wrong path. |
| 131 | + float[] unique = {-1.5f, 2.25f, 3.75f, -4.125f, 5.5f}; |
| 132 | + float[] vals = new float[ROWS]; |
| 133 | + for (int i = 0; i < ROWS; i++) { |
| 134 | + vals[i] = unique[(i / RUN) % unique.length]; |
| 135 | + } |
| 136 | + Path file = tmp.resolve("rle_f32.vtx"); |
| 137 | + writeF32(file, vals); |
| 138 | + assumeRle(file); |
| 139 | + |
| 140 | + // When |
| 141 | + float[] result = readFloatColumn(file); |
| 142 | + |
| 143 | + // Then |
| 144 | + assertThat(result).containsExactly(vals); |
| 145 | + } |
| 146 | + |
| 147 | + // ── helpers ─────────────────────────────────────────────────────────────── |
| 148 | + |
| 149 | + /// Skips the test unless the JNI compressor actually chose `fastlanes.rle`. |
| 150 | + private static void assumeRle(Path file) throws IOException { |
| 151 | + String report; |
| 152 | + try (VortexReader vf = VortexReader.open(file, ReadRegistry.loadAll())) { |
| 153 | + report = VortexInspector.inspect(vf); |
| 154 | + } |
| 155 | + assumeTrue(report.contains("fastlanes.rle"), |
| 156 | + () -> "compressor did not choose fastlanes.rle:\n" + report); |
| 157 | + } |
| 158 | + |
| 159 | + private static void writeF64(Path file, double[] vals) throws IOException { |
| 160 | + String uri = file.toAbsolutePath().toUri().toString(); |
| 161 | + try (VortexWriter writer = VortexWriter.create(SESSION, uri, F64_SCHEMA, new HashMap<>(), ALLOCATOR); |
| 162 | + VectorSchemaRoot root = VectorSchemaRoot.create(F64_SCHEMA, ALLOCATOR)) { |
| 163 | + Float8Vector vec = (Float8Vector) root.getVector("v"); |
| 164 | + vec.allocateNew(vals.length); |
| 165 | + for (int i = 0; i < vals.length; i++) { |
| 166 | + vec.setSafe(i, vals[i]); |
| 167 | + } |
| 168 | + root.setRowCount(vals.length); |
| 169 | + exportBatch(writer, root); |
| 170 | + } |
| 171 | + } |
| 172 | + |
| 173 | + private static void writeF64Nullable(Path file, double[] vals, boolean[] nulls) throws IOException { |
| 174 | + String uri = file.toAbsolutePath().toUri().toString(); |
| 175 | + try (VortexWriter writer = VortexWriter.create(SESSION, uri, F64_NULLABLE_SCHEMA, new HashMap<>(), ALLOCATOR); |
| 176 | + VectorSchemaRoot root = VectorSchemaRoot.create(F64_NULLABLE_SCHEMA, ALLOCATOR)) { |
| 177 | + Float8Vector vec = (Float8Vector) root.getVector("v"); |
| 178 | + vec.allocateNew(vals.length); |
| 179 | + for (int i = 0; i < vals.length; i++) { |
| 180 | + if (nulls[i]) { |
| 181 | + vec.setNull(i); |
| 182 | + } else { |
| 183 | + vec.setSafe(i, vals[i]); |
| 184 | + } |
| 185 | + } |
| 186 | + root.setRowCount(vals.length); |
| 187 | + exportBatch(writer, root); |
| 188 | + } |
| 189 | + } |
| 190 | + |
| 191 | + private static void writeF32(Path file, float[] vals) throws IOException { |
| 192 | + String uri = file.toAbsolutePath().toUri().toString(); |
| 193 | + try (VortexWriter writer = VortexWriter.create(SESSION, uri, F32_SCHEMA, new HashMap<>(), ALLOCATOR); |
| 194 | + VectorSchemaRoot root = VectorSchemaRoot.create(F32_SCHEMA, ALLOCATOR)) { |
| 195 | + Float4Vector vec = (Float4Vector) root.getVector("v"); |
| 196 | + vec.allocateNew(vals.length); |
| 197 | + for (int i = 0; i < vals.length; i++) { |
| 198 | + vec.setSafe(i, vals[i]); |
| 199 | + } |
| 200 | + root.setRowCount(vals.length); |
| 201 | + exportBatch(writer, root); |
| 202 | + } |
| 203 | + } |
| 204 | + |
| 205 | + private static void exportBatch(VortexWriter writer, VectorSchemaRoot root) throws IOException { |
| 206 | + try (ArrowArray arr = ArrowArray.allocateNew(ALLOCATOR); |
| 207 | + ArrowSchema schema = ArrowSchema.allocateNew(ALLOCATOR)) { |
| 208 | + Data.exportVectorSchemaRoot(ALLOCATOR, root, null, arr, schema); |
| 209 | + writer.writeBatch(arr.memoryAddress(), schema.memoryAddress()); |
| 210 | + } |
| 211 | + } |
| 212 | + |
| 213 | + private static double[] readDoubleColumn(Path file) throws IOException { |
| 214 | + try (var vf = VortexReader.open(file, ReadRegistry.loadAll()); |
| 215 | + var iter = vf.scan(ScanOptions.columns("v"))) { |
| 216 | + var out = new ArrayList<Double>(); |
| 217 | + iter.forEachRemaining(c -> { |
| 218 | + DoubleArray arr = c.column("v"); |
| 219 | + for (long i = 0; i < arr.length(); i++) { |
| 220 | + out.add(arr.getDouble(i)); |
| 221 | + } |
| 222 | + }); |
| 223 | + return out.stream().mapToDouble(Double::doubleValue).toArray(); |
| 224 | + } |
| 225 | + } |
| 226 | + |
| 227 | + private static double[] readDoubleColumnUnwrapped(Path file) throws IOException { |
| 228 | + try (var vf = VortexReader.open(file, ReadRegistry.loadAll()); |
| 229 | + var iter = vf.scan(ScanOptions.columns("v"))) { |
| 230 | + var out = new ArrayList<Double>(); |
| 231 | + iter.forEachRemaining(c -> { |
| 232 | + Array a = c.column("v"); |
| 233 | + DoubleArray arr = (DoubleArray) (a instanceof MaskedArray m ? m.inner() : a); |
| 234 | + for (long i = 0; i < arr.length(); i++) { |
| 235 | + out.add(arr.getDouble(i)); |
| 236 | + } |
| 237 | + }); |
| 238 | + return out.stream().mapToDouble(Double::doubleValue).toArray(); |
| 239 | + } |
| 240 | + } |
| 241 | + |
| 242 | + private static float[] readFloatColumn(Path file) throws IOException { |
| 243 | + try (var vf = VortexReader.open(file, ReadRegistry.loadAll()); |
| 244 | + var iter = vf.scan(ScanOptions.columns("v"))) { |
| 245 | + var out = new ArrayList<Float>(); |
| 246 | + iter.forEachRemaining(c -> { |
| 247 | + FloatArray arr = c.column("v"); |
| 248 | + for (long i = 0; i < arr.length(); i++) { |
| 249 | + out.add(arr.getFloat(i)); |
| 250 | + } |
| 251 | + }); |
| 252 | + float[] result = new float[out.size()]; |
| 253 | + for (int i = 0; i < result.length; i++) { |
| 254 | + result[i] = out.get(i); |
| 255 | + } |
| 256 | + return result; |
| 257 | + } |
| 258 | + } |
| 259 | +} |
0 commit comments