diff --git a/CHANGELOG.md b/CHANGELOG.md index 19ee35e24..959934503 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `Chunk.prefetch` and `VortexHandle.prefetch` start the range requests of several columns together over HTTP, merged into few requests, and an HTTP handle keeps the segments it has read. ([#744](https://github.com/dfa1/vortex-java/issues/744)) ### Fixed +- An unexpected HTTP status now reports the server's own explanation and `Retry-After`, so a Hugging Face rate limit (`429`) says what to do instead of just the code. ([#743](https://github.com/dfa1/vortex-java/issues/743)) +- `vortex inspect` over HTTP fetches only the end of each segment, several at a time, instead of the whole file, and shows a progress bar. ([#743](https://github.com/dfa1/vortex-java/issues/743)) - `vortex view` no longer freezes while it fetches data over HTTP: it shows a loading spinner, `q` quits at once, and it fetches only the columns on screen. ([#744](https://github.com/dfa1/vortex-java/issues/744)) - A decimal min or max in a flat node's stats is decoded as a `BigDecimal` instead of text, so a filtered scan can prune on it and `columnStats()` reports it. ([#710](https://github.com/dfa1/vortex-java/issues/710)) - A `vortex.uuid` column's file statistics carry its min and max, as Rust writes them from 0.92, and `WriteSummary` reports them as `UUID`s. ([#730](https://github.com/dfa1/vortex-java/pull/730)) diff --git a/cli/src/main/java/io/github/dfa1/vortex/cli/InspectCommand.java b/cli/src/main/java/io/github/dfa1/vortex/cli/InspectCommand.java index 3d11a70e6..c144fb4f6 100644 --- a/cli/src/main/java/io/github/dfa1/vortex/cli/InspectCommand.java +++ b/cli/src/main/java/io/github/dfa1/vortex/cli/InspectCommand.java @@ -1,5 +1,6 @@ package io.github.dfa1.vortex.cli; +import io.github.dfa1.vortex.core.error.VortexException; import io.github.dfa1.vortex.inspect.HtmlReport; import io.github.dfa1.vortex.inspect.InspectorTree; import io.github.dfa1.vortex.inspect.VortexInspector; @@ -24,13 +25,19 @@ static int run(String[] args) { if (handle == null) { return ExitStatus.FILE_NOT_FOUND; } + InspectorTree tree = InspectorTree.build(handle, (done, total) -> ProgressBar.render(done, total, "segments")); + ProgressBar.clear(); System.out.print(html - ? HtmlReport.render(InspectorTree.build(handle), lastSegment(target)) - : VortexInspector.inspect(handle)); + ? HtmlReport.render(tree, lastSegment(target)) + : VortexInspector.render(tree)); return ExitStatus.OK; } catch (IOException | RuntimeException e) { System.err.println("error: " + CliHandles.describe(e)); - e.printStackTrace(System.err); + // a VortexException is an expected failure (bad file, HTTP error) already fully described + // above; only an unexpected error needs its stack + if (!(e instanceof VortexException)) { + e.printStackTrace(System.err); + } return ExitStatus.ERROR; } } diff --git a/cli/src/main/java/io/github/dfa1/vortex/cli/ProgressBar.java b/cli/src/main/java/io/github/dfa1/vortex/cli/ProgressBar.java index c7f685f55..5a3fcc50c 100644 --- a/cli/src/main/java/io/github/dfa1/vortex/cli/ProgressBar.java +++ b/cli/src/main/java/io/github/dfa1/vortex/cli/ProgressBar.java @@ -17,10 +17,19 @@ private ProgressBar() { /// @param done rows processed so far /// @param total total rows expected static void render(long done, long total) { + render(done, total, "rows"); + } + + /// Renders the same bar counting `unit` instead of rows. + /// + /// @param done units processed so far + /// @param total units expected + /// @param unit what is counted, e.g. `segments` + static void render(long done, long total, String unit) { int pct = total > 0 ? (int) (done * 100L / total) : 100; int filled = pct * WIDTH / 100; String bar = "=".repeat(filled) + (filled < WIDTH ? ">" : "") + " ".repeat(Math.max(0, WIDTH - 1 - filled)); - System.err.printf("\r [%s] %3d%% %,d / %,d rows", bar, pct, done, total); + System.err.printf("\r [%s] %3d%% %,d / %,d %s", bar, pct, done, total, unit); System.err.flush(); } diff --git a/docs/reference.md b/docs/reference.md index c03d39c43..6c3195a24 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -85,7 +85,9 @@ Interface common to both handle implementations below — everything scan-shaped that only reads (e.g. `ParquetExporter.exportParquet(VortexHandle, Path)`) doesn't care whether the source is a local file or a remote URL. Implements `Closeable`. `prefetch(Collection)` is a read-ahead hint for the segments about to be read: a no-op on `VortexReader`, concurrent merged range -requests on `VortexHttpReader`. Call it from the thread that opened the handle. +requests on `VortexHttpReader`. Call it from the thread that opened the handle. `segmentTail(SegmentSpec)` +reads only the end of a segment (its array node, which is all the inspector needs); over HTTP that is one +small range request, with `prefetchTails(Collection)` as its read-ahead hint. | Implementation | Backing storage | |--------------------|-----------------------------------------------| diff --git a/inspector/src/main/java/io/github/dfa1/vortex/inspect/InspectorTree.java b/inspector/src/main/java/io/github/dfa1/vortex/inspect/InspectorTree.java index f22604ba5..3b4ef2c61 100644 --- a/inspector/src/main/java/io/github/dfa1/vortex/inspect/InspectorTree.java +++ b/inspector/src/main/java/io/github/dfa1/vortex/inspect/InspectorTree.java @@ -149,7 +149,7 @@ private static Node shallowNode(Layout layout, Optional fieldName) { /// compression, or missing data. /// /// Callers should cache the result — every call triggers a fresh - /// [io.github.dfa1.vortex.reader.VortexHandle#rawSegment], which is a network round-trip on remote handles. + /// [io.github.dfa1.vortex.reader.VortexHandle#segmentTail], which is a network round-trip on remote handles. /// /// @param node node to resolve /// @param handle open file handle @@ -164,7 +164,7 @@ public static Peek peek(Node node, VortexHandle handle) { if (spec.compression().code != 0) { return Peek.EMPTY; } - MemorySegment seg = handle.rawSegment(spec); + MemorySegment seg = handle.segmentTail(spec); return peekFlatRoot(seg, handle.footer().arraySpecs()); } @@ -180,7 +180,13 @@ public static InspectorTree build(VortexHandle handle, Progress progress) { Layout layout = handle.layout(); DType dtype = handle.dtype(); - int total = countPeekableSegments(layout, footer); + List peekable = new ArrayList<>(); + collectPeekableSegments(layout, footer, peekable); + int total = peekable.size(); + handle.prefetchTails(peekable); + List zoneTables = new ArrayList<>(); + collectZoneTableSegments(layout, footer, zoneTables); + handle.prefetch(zoneTables); int[] counter = {0}; List colNames = rootChildLabels(dtype); @@ -332,7 +338,7 @@ private static Node buildNode(Layout layout, Optional fieldName, VortexH int segIdx = layout.segments().getFirst(); SegmentSpec spec = handle.footer().segmentSpecs().get(segIdx); if (spec.compression().code == 0) { - MemorySegment seg = handle.rawSegment(spec); + MemorySegment seg = handle.segmentTail(spec); Peek peek = peekFlatRoot(seg, arraySpecs); if (peek.encoding() != null) { localUsed.addAll(peek.nestedEncodings()); @@ -354,18 +360,28 @@ private static Node buildNode(Layout layout, Optional fieldName, VortexH return new Node(layout, fieldName, Set.copyOf(localUsed), stats, List.copyOf(children)); } - private static int countPeekableSegments(Layout layout, Footer footer) { - int n = 0; + /// The segments of every zone-map table, which `columnZones` reads one by one after the peeks. + private static void collectZoneTableSegments(Layout layout, Footer footer, List out) { + if (layout.isZoned() && layout.children().size() >= 2) { + for (int segment : layout.children().get(1).segments()) { + out.add(footer.segmentSpecs().get(segment)); + } + } + for (Layout child : layout.children()) { + collectZoneTableSegments(child, footer, out); + } + } + + private static void collectPeekableSegments(Layout layout, Footer footer, List out) { if (layout.isFlat() && !layout.segments().isEmpty()) { SegmentSpec spec = footer.segmentSpecs().get(layout.segments().getFirst()); if (spec.compression().code == 0) { - n++; + out.add(spec); } } for (Layout child : layout.children()) { - n += countPeekableSegments(child, footer); + collectPeekableSegments(child, footer, out); } - return n; } /// Callback used by [#build(VortexHandle, Progress)] to report how many diff --git a/inspector/src/test/java/io/github/dfa1/vortex/inspect/InspectorTreeTest.java b/inspector/src/test/java/io/github/dfa1/vortex/inspect/InspectorTreeTest.java index dab360499..c6f4f9327 100644 --- a/inspector/src/test/java/io/github/dfa1/vortex/inspect/InspectorTreeTest.java +++ b/inspector/src/test/java/io/github/dfa1/vortex/inspect/InspectorTreeTest.java @@ -211,8 +211,8 @@ void buildShallow_skipsAllSlicesAndStillNamesColumns() { assertThat(sut.usedEncodings()).isEmpty(); assertThat(sut.root().children().get(0).usedEncodings()).isEmpty(); assertThat(sut.root().children().getFirst().stats()).isEqualTo(io.github.dfa1.vortex.core.model.ArrayStats.empty()); - // rawSegment is reserved for lazy peek; shallow build must never call it - org.mockito.Mockito.verify(handle, org.mockito.Mockito.never()).rawSegment( + // segmentTail is reserved for lazy peek; shallow build must never call it + org.mockito.Mockito.verify(handle, org.mockito.Mockito.never()).segmentTail( org.mockito.Mockito.any()); } @@ -230,7 +230,7 @@ void peek_nonFlatNode_returnsEmptyWithoutSlicing() { // Then assertThat(result).isSameAs(InspectorTree.Peek.EMPTY); - org.mockito.Mockito.verify(handle, org.mockito.Mockito.never()).rawSegment( + org.mockito.Mockito.verify(handle, org.mockito.Mockito.never()).segmentTail( org.mockito.Mockito.any()); } @@ -251,7 +251,7 @@ void peek_compressedFlatSegment_returnsEmptyWithoutSlicing() { // Then assertThat(result).isSameAs(InspectorTree.Peek.EMPTY); - org.mockito.Mockito.verify(handle, org.mockito.Mockito.never()).rawSegment( + org.mockito.Mockito.verify(handle, org.mockito.Mockito.never()).segmentTail( org.mockito.Mockito.any()); } @@ -287,7 +287,7 @@ void build_flatSegmentWithNestedEncoding_reportsRootAndNestedEncodings() { SegmentSpec spec = new SegmentSpec(0, 100, (byte) 0, CompressionScheme.NONE); givenHandle(dtype, root, arraySpecs, List.of(spec)); byte[] fb = nestedArrayFlatBuffer(0, 1); - given(handle.rawSegment(spec)).willReturn(wrapAsSegment(fb, arena)); + given(handle.segmentTail(spec)).willReturn(wrapAsSegment(fb, arena)); // When InspectorTree sut = InspectorTree.build(handle); @@ -308,7 +308,7 @@ void build_flatSegmentEncodingTreeAtDepthLimit_clearsGuard() { SegmentSpec spec = new SegmentSpec(0, 100, (byte) 0, CompressionScheme.NONE); givenHandle(dtype, root, List.of("vortex.flat"), List.of(spec)); byte[] fb = deeplyNestedArrayFlatBuffer(62); - given(handle.rawSegment(spec)).willReturn(wrapAsSegment(fb, arena)); + given(handle.segmentTail(spec)).willReturn(wrapAsSegment(fb, arena)); // When InspectorTree sut = InspectorTree.build(handle); @@ -329,7 +329,7 @@ void build_flatSegmentEncodingTreeOneOverDepthLimit_throwsVortexException() { SegmentSpec spec = new SegmentSpec(0, 100, (byte) 0, CompressionScheme.NONE); givenHandle(dtype, root, List.of("vortex.flat"), List.of(spec)); byte[] fb = deeplyNestedArrayFlatBuffer(63); - given(handle.rawSegment(spec)).willReturn(wrapAsSegment(fb, arena)); + given(handle.segmentTail(spec)).willReturn(wrapAsSegment(fb, arena)); // When / Then assertThatThrownBy(() -> InspectorTree.build(handle)) @@ -350,7 +350,7 @@ void build_flatSegmentEncodingIndexOutOfBounds_throwsVortexException() { SegmentSpec spec = new SegmentSpec(0, 100, (byte) 0, CompressionScheme.NONE); givenHandle(dtype, root, List.of("vortex.flat"), List.of(spec)); byte[] fb = nestedArrayFlatBuffer(0, 999); - given(handle.rawSegment(spec)).willReturn(wrapAsSegment(fb, arena)); + given(handle.segmentTail(spec)).willReturn(wrapAsSegment(fb, arena)); // When / Then assertThatThrownBy(() -> InspectorTree.build(handle)) @@ -373,7 +373,7 @@ void build_flatSegmentEncodingTreeOverNodeCountLimit_throwsVortexException() { SegmentSpec spec = new SegmentSpec(0, 100, (byte) 0, CompressionScheme.NONE); givenHandle(dtype, root, List.of("vortex.flat"), List.of(spec)); byte[] fb = wideArrayFlatBuffer(InspectorTree.MAX_ENCODING_TREE_NODES + 1); - given(handle.rawSegment(spec)).willReturn(wrapAsSegment(fb, arena)); + given(handle.segmentTail(spec)).willReturn(wrapAsSegment(fb, arena)); // When / Then assertThatThrownBy(() -> InspectorTree.build(handle)) diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/VortexHandle.java b/reader/src/main/java/io/github/dfa1/vortex/reader/VortexHandle.java index c33062574..68594bf32 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/VortexHandle.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/VortexHandle.java @@ -61,6 +61,24 @@ public interface VortexHandle extends Closeable { /// `spec.offset() + spec.length()` MemorySegment rawSegment(SegmentSpec spec); + /// Returns the end of the given segment: at least its array-node FlatBuffer and the 4-byte + /// length that follows it, which is all the inspector reads to name an encoding and its stats. + /// May be the whole segment. On an HTTP handle it fetches a few KiB instead of the segment. + /// + /// @param spec the segment to read the end of + /// @return a read-only [MemorySegment] ending at the same byte as the segment + default MemorySegment segmentTail(SegmentSpec spec) { + return rawSegment(spec); + } + + /// Hints that `specs` are about to be read with [#segmentTail(SegmentSpec)], so a remote + /// handle can fetch them concurrently. Does nothing on a memory-mapped handle. Call it from + /// the thread that reads the segments, the one that opened the handle. + /// + /// @param specs the segments whose ends are about to be read + default void prefetchTails(Collection specs) { + } + /// Hints that `specs` are about to be read with [#rawSegment(SegmentSpec)], so a remote handle /// can fetch them concurrently instead of one round trip each. Does nothing on a /// memory-mapped handle, where a read is already free. Call it from the thread that reads the diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/VortexHttpReader.java b/reader/src/main/java/io/github/dfa1/vortex/reader/VortexHttpReader.java index b4ed8c69a..ed6c6ad6f 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/VortexHttpReader.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/VortexHttpReader.java @@ -18,6 +18,8 @@ import java.net.http.HttpClient; import java.net.http.HttpRequest; import java.net.http.HttpResponse; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; import java.util.ArrayList; import java.util.Collection; import java.util.Comparator; @@ -45,10 +47,18 @@ public final class VortexHttpReader implements VortexHandle { /// all metadata blobs for typical Vortex files. static final MemorySize TAIL_SIZE = MemorySize.ofKiB(65); + /// First guess for [#segmentTail(SegmentSpec)]: an array-node FlatBuffer is a few hundred + /// bytes, so one request nearly always gets it with its length. + private static final MemorySize TAIL_PEEK_SIZE = MemorySize.ofKiB(4); + /// Cap on concurrent prefetch requests: enough to hide a ~200 ms round trip across a /// screenful of columns without opening a connection per segment. private static final int MAX_IN_FLIGHT = 8; + /// Cap on concurrent [#prefetchTails(Collection)] requests: a tail is a few KiB, so the + /// round trip dominates and a deeper queue than [#MAX_IN_FLIGHT] pays off. + private static final int MAX_TAILS_IN_FLIGHT = 32; + /// Most heap one [#prefetch(Collection)] call queues: a fetched response waits as a `byte[]` /// until its segments are read, unlike the off-heap arena a read copies it into. private static final long MAX_PREFETCH_BYTES = MemorySize.ofMiB(64).bytes(); @@ -85,12 +95,16 @@ public final class VortexHttpReader implements VortexHandle { private final LayoutRegistry layoutRegistry; /// Range requests started by [#prefetch(Collection)] and not yet consumed by [#rawSegment]. private final Map prefetched; + + /// Tail requests started by [#prefetchTails(Collection)] and not yet consumed by [#segmentTail]. + private final Map> tails; /// Every segment already read, so reading it again (scrolling back) costs no request. Never /// evicted and never expired: the file is immutable for the life of this handle, and the bytes /// live in [#arena] until close anyway, so keeping the reference adds no memory. Bounded by the /// file size. Touched only from the arena's owning thread, like [#rawSegment]. private final Map fetched; private final Semaphore inFlight; + private final Semaphore tailInFlight; private final ExecutorService fetcher; private boolean closed; @@ -109,9 +123,11 @@ private VortexHttpReader( this.registry = registry; this.layoutRegistry = layoutRegistry; this.prefetched = new ConcurrentHashMap<>(); + this.tails = new ConcurrentHashMap<>(); this.fetched = new HashMap<>(); this.closed = false; this.inFlight = new Semaphore(MAX_IN_FLIGHT); + this.tailInFlight = new Semaphore(MAX_TAILS_IN_FLIGHT); this.fetcher = Executors.newVirtualThreadPerTaskExecutor(); } @@ -294,7 +310,22 @@ private static TailFetch fetchTail(URI uri, HttpClient client) throws IOExceptio return new TailFetch(body, 0L, body.length); } - throw new VortexException("HTTP " + status + " fetching tail of " + uri); + throw new VortexException(httpError(resp, "fetching tail of " + uri)); + } + + /// The message for an unexpected status: the code, and for the server's own explanation, as + /// Hugging Face gives on a `429` (who is limited, what to do), `x-error-message` and + /// `Retry-After` when present. Both come from an untrusted server, so the text is capped. + private static String httpError(HttpResponse resp, String what) { + StringBuilder message = new StringBuilder("HTTP ").append(resp.statusCode()).append(' ').append(what); + resp.headers().firstValue("x-error-message").ifPresent(m -> message.append(": ").append(cap(m))); + resp.headers().firstValue("Retry-After") + .ifPresent(seconds -> message.append(" (Retry-After: ").append(cap(seconds)).append(')')); + return message.toString(); + } + + private static String cap(String text) { + return text.length() <= 300 ? text : text.substring(0, 300) + "…"; } /// Sends one suffix-Range request for the last `suffixLength` bytes of `uri`. @@ -371,7 +402,7 @@ private static byte[] fetchRange(URI uri, long from, long to, HttpClient client) HttpResponse resp = sendRange(req, client, uri, from, to - from + 1); int status = resp.statusCode(); if (status != 206 && status != 200) { - throw new VortexException("HTTP " + status + " fetching range from " + uri); + throw new VortexException(httpError(resp, "fetching range from " + uri)); } byte[] body = resp.body(); long expected = to - from + 1; @@ -483,6 +514,76 @@ public MemorySegment rawSegment(SegmentSpec spec) { return readOnly; } + /// Fetches only the end of the segment: a first guess of [#TAIL_PEEK_SIZE] covers the + /// array-node FlatBuffer of almost every segment, and a second request completes it when the + /// trailing length says it is larger. Uses the request [#prefetchTails(Collection)] started + /// for this segment, if any. + /// + /// @param spec the segment to read the end of + /// @return a read-only [MemorySegment] ending at the same byte as the segment + @Override + public MemorySegment segmentTail(SegmentSpec spec) { + CompletableFuture started = tails.remove(spec); + byte[] tail; + if (started == null) { + tail = fetchTail(spec); + } else { + try { + tail = started.get(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new VortexException("interrupted fetching the end of a segment from " + uri); + } catch (ExecutionException e) { + if (e.getCause() instanceof VortexException vortex) { + throw vortex; + } + throw new VortexException("failed to fetch the end of a segment from " + uri, e.getCause()); + } + } + return MemorySegment.ofArray(tail).asReadOnly(); + } + + /// Starts the [#segmentTail(SegmentSpec)] requests for `specs` on virtual threads, at most + /// [#MAX_TAILS_IN_FLIGHT] at a time. A tail is a few KiB, so queueing one per segment of a large + /// file stays small on the heap. Call it from the thread that opened this handle; it does + /// nothing once the handle is closed. + /// + /// @param specs the segments whose ends are about to be read + @Override + public void prefetchTails(Collection specs) { + if (closed) { + return; + } + for (SegmentSpec spec : specs) { + tails.computeIfAbsent(spec, s -> CompletableFuture.supplyAsync(() -> fetchTail(s), fetcher)); + } + } + + private byte[] fetchTail(SegmentSpec spec) { + long end = spec.offset() + spec.length(); + long first = Math.min(spec.length(), TAIL_PEEK_SIZE.bytes()); + byte[] tail = fetchBounded(end - first, end); + if (tail.length >= Integer.BYTES) { + long needed = Integer.BYTES + Integer.toUnsignedLong( + ByteBuffer.wrap(tail).order(ByteOrder.LITTLE_ENDIAN).getInt(tail.length - Integer.BYTES)); + if (needed > tail.length && needed <= spec.length()) { + tail = fetchBounded(end - needed, end); + } + } + return tail; + } + + private byte[] fetchBounded(long from, long end) { + tailInFlight.acquireUninterruptibly(); + try { + return fetchRange(uri, from, end - 1, client); + } catch (IOException e) { + throw new VortexException("failed to fetch [%d, %d) from %s: %s".formatted(from, end, uri, e.getMessage())); + } finally { + tailInFlight.release(); + } + } + /// Starts the range requests for `specs` on virtual threads, at most [#MAX_IN_FLIGHT] at a /// time, so the [#rawSegment(SegmentSpec)] calls that follow find their bytes already on the /// way instead of paying one round trip each. Allocation still happens on the calling @@ -593,6 +694,7 @@ public void close() { closed = true; fetcher.shutdownNow(); prefetched.clear(); + tails.clear(); fetched.clear(); } diff --git a/reader/src/test/java/io/github/dfa1/vortex/reader/VortexHttpReaderTailFetchTest.java b/reader/src/test/java/io/github/dfa1/vortex/reader/VortexHttpReaderTailFetchTest.java index b73992331..2843f661e 100644 --- a/reader/src/test/java/io/github/dfa1/vortex/reader/VortexHttpReaderTailFetchTest.java +++ b/reader/src/test/java/io/github/dfa1/vortex/reader/VortexHttpReaderTailFetchTest.java @@ -162,6 +162,22 @@ void open_416RetryStill416_throwsWithStatus() throws Exception { .hasMessageContaining("HTTP 416"); } + @Test + void open_429_reportsTheServersExplanationAndRetryAfter() throws Exception { + // Given — Hugging Face's anonymous rate limit answers with a message saying what to do + // and when to retry; the bare status code hides both from the user. + doReturn(responseWithHeaders(429, Map.of( + "x-error-message", List.of("We had to rate limit your IP. Pass a HF_TOKEN."), + "retry-after", List.of("31")), new byte[0])).when(client).send(any(), any()); + + // When / Then + assertThatThrownBy(() -> VortexHttpReader.open(URI, ReadRegistry.empty(), client)) + .isInstanceOf(VortexException.class) + .hasMessageContaining("HTTP 429") + .hasMessageContaining("Pass a HF_TOKEN.") + .hasMessageContaining("Retry-After: 31"); + } + // ── helpers ─────────────────────────────────────────────────────────────── private static byte[] fixtureBytes() throws IOException { @@ -174,8 +190,16 @@ private static byte[] fixtureBytes() throws IOException { } } - @SuppressWarnings("unchecked") private static HttpResponse response(int status, String contentRange, byte[] body) { + Map> headers = Map.of(); + if (contentRange != null) { + headers = Map.of("content-range", List.of(contentRange)); + } + return responseWithHeaders(status, headers, body); + } + + @SuppressWarnings("unchecked") + private static HttpResponse responseWithHeaders(int status, Map> headers, byte[] body) { return new HttpResponse<>() { @Override public int statusCode() { @@ -189,10 +213,7 @@ public byte[] body() { @Override public HttpHeaders headers() { - Map> map = contentRange == null - ? Map.of() - : Map.of("content-range", List.of(contentRange)); - return HttpHeaders.of(map, (k, v) -> true); + return HttpHeaders.of(headers, (k, v) -> true); } @Override diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/WriterHttpTailTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/WriterHttpTailTest.java new file mode 100644 index 000000000..b9cf40f21 --- /dev/null +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/WriterHttpTailTest.java @@ -0,0 +1,150 @@ +package io.github.dfa1.vortex.writer; + +import io.github.dfa1.vortex.core.model.ColumnName; +import io.github.dfa1.vortex.core.model.DType; +import io.github.dfa1.vortex.core.model.DTypes; +import io.github.dfa1.vortex.reader.ReadRegistry; +import io.github.dfa1.vortex.reader.SegmentSpec; +import io.github.dfa1.vortex.reader.VortexHttpReader; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.api.io.TempDir; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +import java.lang.foreign.ValueLayout; +import java.util.ArrayList; +import java.util.Random; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpHeaders; +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.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicLong; +import javax.net.ssl.SSLSession; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.BDDMockito.given; + +/// `segmentTail` reads only the end of a segment over HTTP, where the inspector finds the array node, +/// so `inspect` on a remote file no longer downloads every segment to print metadata (#743). +@ExtendWith(MockitoExtension.class) +class WriterHttpTailTest { + + private static final URI FILE_URI = URI.create("http://example.com/tail.vortex"); + private static final ColumnName V = ColumnName.of("v"); + + @Mock + private HttpClient client; + + @Test + void segmentTail_matchesTheEndOfTheSegmentAndFetchesLittle(@TempDir Path tmp) throws Exception { + // Given — a multi-megabyte file with many segments + DType.Struct schema = new DType.Struct(List.of(V), List.of(DTypes.I64), false); + WriteOptions opts = WriteOptions.cascading(0).withZoneMaps(true).withGlobalDict(false).withoutEditions(); + Path file = tmp.resolve("tail.vtx"); + try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE); + var writer = VortexWriter.create(ch, schema, opts)) { + for (int chunk = 0; chunk < 4; chunk++) { + long[] values = new long[300_000]; + Random random = new Random(chunk); + for (int r = 0; r < values.length; r++) { + values[r] = random.nextLong(); + } + writer.writeChunk(Map.of(V, values)); + } + } + byte[] bytes = Files.readAllBytes(file); + AtomicLong fetched = new AtomicLong(); + given(client.send(any(), any())).willAnswer(inv -> { + HttpResponse response = rangeResponse(inv.getArgument(0), bytes); + fetched.addAndGet(response.body().length); + return response; + }); + + try (var reader = VortexHttpReader.open(FILE_URI, ReadRegistry.loadAll(), client)) { + List specs = reader.footer().segmentSpecs().stream().filter(s -> s.length() > 4).toList(); + long afterOpen = fetched.get(); + + // When + reader.prefetchTails(specs); + List result = new ArrayList<>(); + for (SegmentSpec spec : specs) { + result.add(reader.segmentTail(spec).toArray(ValueLayout.JAVA_BYTE)); + } + + // Then — each tail is the last bytes of its segment, and far less than the segments themselves was fetched + long segmentBytes = 0; + for (int i = 0; i < specs.size(); i++) { + SegmentSpec spec = specs.get(i); + int end = (int) (spec.offset() + spec.length()); + assertThat(result.get(i)).isEqualTo(Arrays.copyOfRange(bytes, end - result.get(i).length, end)); + segmentBytes += spec.length(); + } + assertThat(segmentBytes).isGreaterThan(1_000_000); + assertThat(fetched.get() - afterOpen).isLessThan(segmentBytes / 10); + } + } + + private static HttpResponse rangeResponse(HttpRequest req, byte[] file) { + String[] parts = req.headers().firstValue("Range").orElseThrow().substring("bytes=".length()).split("-", 2); + int from = Integer.parseInt(parts[0].isEmpty() ? "0" : parts[0]); + int to = Math.min(Integer.parseInt(parts[1].isEmpty() ? "0" : parts[1]), file.length - 1); + if (parts[0].isEmpty()) { + from = Math.max(0, file.length - Integer.parseInt(parts[1])); + to = file.length - 1; + } + byte[] body = Arrays.copyOfRange(file, from, to + 1); + String contentRange = "bytes " + from + "-" + to + "/" + file.length; + return new HttpResponse<>() { + @Override + public int statusCode() { + return 206; + } + + @Override + public byte[] body() { + return body; + } + + @Override + public HttpHeaders headers() { + return HttpHeaders.of(Map.of("content-range", List.of(contentRange)), (k, v) -> true); + } + + @Override + public HttpRequest request() { + return null; + } + + @Override + public Optional> previousResponse() { + return Optional.empty(); + } + + @Override + public Optional sslSession() { + return Optional.empty(); + } + + @Override + public URI uri() { + return FILE_URI; + } + + @Override + public HttpClient.Version version() { + return HttpClient.Version.HTTP_1_1; + } + }; + } +}