|
1 | 1 | package io.github.dfa1.vortex.writer; |
2 | 2 |
|
| 3 | +import static io.github.dfa1.vortex.core.io.PTypeIO.LE_FLOAT; |
3 | 4 | import static io.github.dfa1.vortex.core.io.PTypeIO.LE_INT; |
| 5 | +import static io.github.dfa1.vortex.core.io.PTypeIO.LE_SHORT; |
4 | 6 | import java.lang.foreign.MemorySegment; |
5 | 7 | import java.lang.foreign.ValueLayout; |
6 | 8 |
|
|
9 | 11 | import io.github.dfa1.vortex.reader.Layout; |
10 | 12 | import io.github.dfa1.vortex.reader.SegmentSpec; |
11 | 13 | import io.github.dfa1.vortex.reader.VortexReader; |
| 14 | +import io.github.dfa1.vortex.reader.array.DoubleArray; |
12 | 15 | import io.github.dfa1.vortex.reader.array.LongArray; |
13 | 16 | import io.github.dfa1.vortex.reader.array.MaskedArray; |
14 | 17 | import io.github.dfa1.vortex.reader.array.StructArray; |
@@ -382,6 +385,105 @@ void primitiveDictColumn_emitsNumericMinMaxZoneMapWrappingDict(@TempDir Path tmp |
382 | 385 | } |
383 | 386 | } |
384 | 387 |
|
| 388 | + @Test |
| 389 | + void zoneMaps_f64StatsPayloadDecodesPerZoneMinMaxSum(@TempDir Path tmp) throws IOException { |
| 390 | + // Given a zone-mapped F64 file (3 zones of 4 rows). Only I64 stats are otherwise covered, so |
| 391 | + // the float stat path — scalarDouble plus the F64 arms of statColumn/sumColumn — is untested: |
| 392 | + // a decode that returned 0.0 or read the wrong scalar field would slip through. Fractional |
| 393 | + // .5 values also make a truncating (int) decode observable. |
| 394 | + DType.Struct schema = new DType.Struct(List.of("v"), List.of(DType.F64), false); |
| 395 | + WriteOptions opts = new WriteOptions(4, true, 0.90, 0, true, false); |
| 396 | + Path file = tmp.resolve("zoned-f64.vtx"); |
| 397 | + try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE); |
| 398 | + var sut = VortexWriter.create(ch, schema, opts)) { |
| 399 | + for (int z = 0; z < 3; z++) { |
| 400 | + double[] v = new double[4]; |
| 401 | + for (int i = 0; i < 4; i++) { |
| 402 | + v[i] = z * 4.0 + i + 0.5; |
| 403 | + } |
| 404 | + sut.writeChunk(Map.of("v", v)); |
| 405 | + } |
| 406 | + } |
| 407 | + |
| 408 | + // When / Then min/max/sum per zone decode as exact doubles |
| 409 | + try (VortexReader reader = VortexReader.open(file)) { |
| 410 | + Layout zonesFlat = reader.layout().children().get(0).children().get(1); |
| 411 | + SegmentSpec spec = reader.footer().segmentSpecs().get(zonesFlat.segments().getFirst()); |
| 412 | + try (Arena arena = Arena.ofConfined()) { |
| 413 | + StructArray stats = (StructArray) reader.decodeFlatSegment(spec, f64StatsTableDtype(), 3, arena); |
| 414 | + DoubleArray max = (DoubleArray) ((MaskedArray) stats.field("max")).inner(); |
| 415 | + DoubleArray min = (DoubleArray) ((MaskedArray) stats.field("min")).inner(); |
| 416 | + DoubleArray sum = (DoubleArray) ((MaskedArray) stats.field("sum")).inner(); |
| 417 | + for (int z = 0; z < 3; z++) { |
| 418 | + double base = z * 4.0; |
| 419 | + assertThat(min.getDouble(z)).as("min zone %d", z).isEqualTo(base + 0.5); |
| 420 | + assertThat(max.getDouble(z)).as("max zone %d", z).isEqualTo(base + 3.5); |
| 421 | + assertThat(sum.getDouble(z)).as("sum zone %d", z) |
| 422 | + .isEqualTo((base + 0.5) + (base + 1.5) + (base + 2.5) + (base + 3.5)); |
| 423 | + } |
| 424 | + } |
| 425 | + } |
| 426 | + } |
| 427 | + |
| 428 | + @ParameterizedTest |
| 429 | + @EnumSource(value = PType.class, names = {"I8", "I16", "I32", "F32"}) |
| 430 | + void zoneMaps_perTypeStatsDecodePerZoneMinMax(PType ptype, @TempDir Path tmp) throws IOException { |
| 431 | + // Given a 2-zone column of `ptype`: zone 0 = [0, 1], zone 1 = [2, 3]. everyPrimitiveType only |
| 432 | + // checks the layout is zoned; this pins the decoded per-zone MIN/MAX for the I8/I16/I32 and |
| 433 | + // F32 statColumn arms (and the f32 scalar field read in scalarDouble), which the I64/F64 |
| 434 | + // value tests never reach. |
| 435 | + DType.Struct schema = new DType.Struct(List.of("v"), List.of(new DType.Primitive(ptype, false)), false); |
| 436 | + WriteOptions opts = new WriteOptions(2, true, 0.90, 0, false, false); |
| 437 | + Path file = tmp.resolve("ptype-stats-" + ptype + ".vtx"); |
| 438 | + try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE); |
| 439 | + var sut = VortexWriter.create(ch, schema, opts)) { |
| 440 | + sut.writeChunk(Map.of("v", sample(ptype, 0))); |
| 441 | + sut.writeChunk(Map.of("v", sample(ptype, 2))); |
| 442 | + } |
| 443 | + |
| 444 | + // When / Then min/max per zone match the source |
| 445 | + try (VortexReader reader = VortexReader.open(file)) { |
| 446 | + Layout zonesFlat = reader.layout().children().get(0).children().get(1); |
| 447 | + SegmentSpec spec = reader.footer().segmentSpecs().get(zonesFlat.segments().getFirst()); |
| 448 | + try (Arena arena = Arena.ofConfined()) { |
| 449 | + StructArray stats = (StructArray) reader.decodeFlatSegment(spec, perTypeStatsTableDtype(ptype), 2, arena); |
| 450 | + MemorySegment min = ((MaskedArray) stats.field("min")).inner().materialize(arena); |
| 451 | + MemorySegment max = ((MaskedArray) stats.field("max")).inner().materialize(arena); |
| 452 | + assertThat(readStat(min, ptype, 0)).as("min zone 0").isEqualTo(0.0); |
| 453 | + assertThat(readStat(max, ptype, 0)).as("max zone 0").isEqualTo(1.0); |
| 454 | + assertThat(readStat(min, ptype, 1)).as("min zone 1").isEqualTo(2.0); |
| 455 | + assertThat(readStat(max, ptype, 1)).as("max zone 1").isEqualTo(3.0); |
| 456 | + } |
| 457 | + } |
| 458 | + } |
| 459 | + |
| 460 | + /// Reads the per-zone stat at `idx` from a materialised min/max segment, widened to double for |
| 461 | + /// uniform assertion across the fixed-width primitive types. |
| 462 | + private static double readStat(MemorySegment seg, PType ptype, int idx) { |
| 463 | + return switch (ptype) { |
| 464 | + case I8 -> seg.get(ValueLayout.JAVA_BYTE, idx); |
| 465 | + case I16 -> seg.get(LE_SHORT, (long) idx * 2); |
| 466 | + case I32 -> seg.get(LE_INT, (long) idx * 4); |
| 467 | + case F32 -> seg.get(LE_FLOAT, (long) idx * 4); |
| 468 | + default -> throw new AssertionError(ptype); |
| 469 | + }; |
| 470 | + } |
| 471 | + |
| 472 | + /// Stats-table dtype for a numeric column of `ptype`: MAX/MIN in the column ptype, SUM widened |
| 473 | + /// to i64/u64/f64 per [VortexWriter] `zoneSumDtype`, plus NULL_COUNT. |
| 474 | + private static DType.Struct perTypeStatsTableDtype(PType ptype) { |
| 475 | + DType minMax = new DType.Primitive(ptype, true); |
| 476 | + DType sum = switch (ptype) { |
| 477 | + case U8, U16, U32, U64 -> new DType.Primitive(PType.U64, true); |
| 478 | + case I8, I16, I32, I64 -> new DType.Primitive(PType.I64, true); |
| 479 | + default -> new DType.Primitive(PType.F64, true); |
| 480 | + }; |
| 481 | + return new DType.Struct( |
| 482 | + List.of("max", "max_is_truncated", "min", "min_is_truncated", "sum", "null_count"), |
| 483 | + List.of(minMax, DType.BOOL, minMax, DType.BOOL, sum, new DType.Primitive(PType.U64, true)), |
| 484 | + false); |
| 485 | + } |
| 486 | + |
385 | 487 | /// Two values starting at `base` in the storage shape for `ptype`. |
386 | 488 | private static Object sample(PType ptype, int base) { |
387 | 489 | return switch (ptype) { |
@@ -416,6 +518,16 @@ private static DType.Struct numericStatsTableDtype() { |
416 | 518 | false); |
417 | 519 | } |
418 | 520 |
|
| 521 | + /// Stats-table dtype for a numeric F64 column: MAX, MIN, SUM (nullable F64), NULL_COUNT. |
| 522 | + private static DType.Struct f64StatsTableDtype() { |
| 523 | + DType nullableF64 = new DType.Primitive(PType.F64, true); |
| 524 | + return new DType.Struct( |
| 525 | + List.of("max", "max_is_truncated", "min", "min_is_truncated", "sum", "null_count"), |
| 526 | + List.of(nullableF64, DType.BOOL, nullableF64, DType.BOOL, |
| 527 | + nullableF64, new DType.Primitive(PType.U64, true)), |
| 528 | + false); |
| 529 | + } |
| 530 | + |
419 | 531 | /// Same as [#statsTableDtype()] but with Utf8 (string) MAX/MIN columns. |
420 | 532 | private static DType.Struct utf8StatsTableDtype() { |
421 | 533 | DType nullableUtf8 = new DType.Utf8(true); |
|
0 commit comments