diff --git a/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java b/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java index e59a7972587..c740a4b70f5 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java +++ b/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java @@ -8,6 +8,7 @@ import dev.vortex.io.NativeWritable; import dev.vortex.jni.NativeWriter; import java.io.IOException; +import java.util.HashMap; import java.util.Map; import java.util.Objects; import java.util.concurrent.atomic.AtomicBoolean; @@ -19,6 +20,9 @@ /** * Writer for Vortex files. * + *

Open one with {@link #builder(Session, String, Schema, BufferAllocator)} to write to a URI, or + * {@link #builder(Session, NativeWritable, Schema, BufferAllocator)} to write into a caller-provided byte sink. + * *

Batches are accepted via the Arrow C Data Interface: callers export an Arrow record batch to an {@code ArrowArray} * / {@code ArrowSchema} pair and pass the addresses to {@link #writeBatch(long, long)}. The writer accepts up to four * in-flight batches on the session's runtime thread before back-pressuring the caller. @@ -44,53 +48,133 @@ private VortexWriter(long pointer) { } /** - * Create a writer that streams records into the file at {@code uri}. The path may be a full URI or a plain local - * filesystem path. The Arrow schema describes the exact layout of every batch written. + * Start configuring a writer that streams records into the file at {@code uri} through a native storage client. The + * path may be a full URI or a plain local filesystem path. + * + * @param arrowSchema describes the exact layout of every batch written */ - public static VortexWriter create( - Session session, String uri, Schema arrowSchema, Map options, BufferAllocator allocator) - throws IOException { - Objects.requireNonNull(session, "session"); - Objects.requireNonNull(uri, "uri"); - Objects.requireNonNull(arrowSchema, "arrowSchema"); - Objects.requireNonNull(allocator, "allocator"); - ArrowSchema ffi = ArrowSchema.allocateNew(allocator); - try { - Data.exportSchema(allocator, arrowSchema, null, ffi); - long ptr = NativeWriter.create(session.nativePointer(), uri, ffi.memoryAddress(), options); - if (ptr <= 0) { - throw new IOException("failed to create writer for uri " + uri + " (ptr=" + ptr + ")"); - } - return new VortexWriter(ptr); - } finally { - ffi.close(); - } + public static Builder builder(Session session, String uri, Schema arrowSchema, BufferAllocator allocator) { + return new Builder(session, Objects.requireNonNull(uri, "uri"), null, arrowSchema, allocator); } /** - * Create a writer that streams the file into a caller-provided byte sink instead of a native storage client. This - * is the integration point for external I/O abstractions (for example Iceberg's {@code PositionOutputStream}). + * Start configuring a writer that streams the file into a caller-provided byte sink instead of a native storage + * client. This is the integration point for external I/O abstractions (for example Iceberg's + * {@code PositionOutputStream}). * *

The native side writes and flushes the sink but never closes it: after {@link #close()} returns, all bytes * have been written and flushed, and the caller must close the sink to finalize the file. + * + * @param arrowSchema describes the exact layout of every batch written */ - public static VortexWriter create( - Session session, NativeWritable writable, Schema arrowSchema, BufferAllocator allocator) - throws IOException { - Objects.requireNonNull(session, "session"); - Objects.requireNonNull(writable, "writable"); - Objects.requireNonNull(arrowSchema, "arrowSchema"); - Objects.requireNonNull(allocator, "allocator"); - ArrowSchema ffi = ArrowSchema.allocateNew(allocator); - try { - Data.exportSchema(allocator, arrowSchema, null, ffi); - long ptr = NativeWriter.createStream(session.nativePointer(), writable, ffi.memoryAddress()); - if (ptr <= 0) { - throw new IOException("failed to create stream writer (ptr=" + ptr + ")"); + public static Builder builder( + Session session, NativeWritable writable, Schema arrowSchema, BufferAllocator allocator) { + return new Builder(session, null, Objects.requireNonNull(writable, "writable"), arrowSchema, allocator); + } + + /** + * Configures and opens a {@link VortexWriter}. Everything required to open one is fixed by + * {@link VortexWriter#builder}; the setters here are all optional. + * + *

Not thread-safe, and not reusable: each {@link #build()} opens one file. + */ + public static final class Builder { + private final Session session; + private final String uri; + private final NativeWritable writable; + private final Schema arrowSchema; + private final BufferAllocator allocator; + private final Map options = new HashMap<>(); + private final Map metadata = new HashMap<>(); + + private Builder( + Session session, String uri, NativeWritable writable, Schema arrowSchema, BufferAllocator allocator) { + this.session = Objects.requireNonNull(session, "session"); + this.uri = uri; + this.writable = writable; + this.arrowSchema = Objects.requireNonNull(arrowSchema, "arrowSchema"); + this.allocator = Objects.requireNonNull(allocator, "allocator"); + } + + /** + * Object-store credentials and options, replacing any set so far. Only the URI destination reaches a storage + * client, so this is rejected on a writer built over a {@link NativeWritable}. + */ + public Builder options(Map newOptions) { + checkOptionsApply(); + Objects.requireNonNull(newOptions, "options"); + options.clear(); + newOptions.forEach(this::putOption); + return this; + } + + /** Add a single object-store option. See {@link #options(Map)}. */ + public Builder putOption(String key, String value) { + checkOptionsApply(); + Objects.requireNonNull(key, "key"); + Objects.requireNonNull(value, "value"); + options.put(key, value); + return this; + } + + /** + * User-defined metadata segments to store in the file footer, replacing any set so far. + * + *

Values are opaque bytes, returned verbatim by {@code NativeFiles.readMetadata}. The native writer caps + * both the number of segments per file and the length of each key; the limits are enforced (and named in the + * error) by {@link #build()}, rather than when the file is finalized. + */ + public Builder metadata(Map newMetadata) { + Objects.requireNonNull(newMetadata, "metadata"); + metadata.clear(); + newMetadata.forEach(this::putMetadata); + return this; + } + + /** Add a single metadata segment. See {@link #metadata(Map)}. */ + public Builder putMetadata(String key, byte[] value) { + Objects.requireNonNull(key, "key"); + Objects.requireNonNull(value, "value"); + metadata.put(key, value); + return this; + } + + /** + * Open the writer. + * + * @throws IOException if the writer cannot be opened, including when the native side rejects the configured + * metadata + */ + public VortexWriter build() throws IOException { + ArrowSchema ffi = ArrowSchema.allocateNew(allocator); + try { + Data.exportSchema(allocator, arrowSchema, null, ffi); + final long pointer; + try { + pointer = uri != null + ? NativeWriter.create(session.nativePointer(), uri, ffi.memoryAddress(), options, metadata) + : NativeWriter.createStream( + session.nativePointer(), writable, ffi.memoryAddress(), metadata); + } catch (RuntimeException e) { + // Native failures arrive as RuntimeException. Wrap them so callers can handle every + // failure to open a writer as an IOException, as writeBatch and finish already do. + throw new IOException("failed to create writer for " + destination(), e); + } + if (pointer <= 0) { + throw new IOException("failed to create writer for " + destination() + " (ptr=" + pointer + ")"); + } + return new VortexWriter(pointer); + } finally { + ffi.close(); } - return new VortexWriter(ptr); - } finally { - ffi.close(); + } + + private String destination() { + return uri != null ? "uri " + uri : "stream"; + } + + private void checkOptionsApply() { + Preconditions.checkState(uri != null, "options apply to uri destinations only"); } } diff --git a/java/vortex-jni/src/main/java/dev/vortex/jni/NativeFiles.java b/java/vortex-jni/src/main/java/dev/vortex/jni/NativeFiles.java index fae03d50186..8a969276e34 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/jni/NativeFiles.java +++ b/java/vortex-jni/src/main/java/dev/vortex/jni/NativeFiles.java @@ -3,13 +3,16 @@ package dev.vortex.jni; +import com.google.common.base.Preconditions; import dev.vortex.api.Session; +import dev.vortex.io.NativeReadable; import java.util.List; import java.util.Map; +import java.util.Objects; /** - * Static utilities for discovering and deleting Vortex files on an object store. The caller supplies a {@link Session}; - * its runtime handle is forwarded to the underlying object store. + * Static utilities for discovering, deleting, and inspecting the metadata of Vortex files on an object store. The + * caller supplies a {@link Session}; its runtime handle is forwarded to the underlying object store. */ public final class NativeFiles { static { @@ -28,7 +31,42 @@ public static void delete(Session session, String[] uris, Map op delete(session.nativePointer(), uris, options); } + /** + * Read the user-defined metadata segments written into the Vortex file at {@code uri}, as opaque bytes keyed by the + * names the {@code VortexWriter} builder was given. Returns an empty map for a file that carries no metadata. + * + *

This opens the file on its own. Metadata lives in dedicated segments that opening a file does not read by + * default, so a segment that falls outside the file-tail read costs an extra round trip. + */ + public static Map readMetadata(Session session, String uri, Map options) { + return readMetadata(session.nativePointer(), uri, options); + } + + /** + * Read the user-defined metadata segments of a Vortex file through a caller-provided {@link NativeReadable}, + * instead of a native storage client. See {@link #readMetadata(Session, String, Map)}. + * + *

The readable must stay open for the duration of the call; native code never closes it. + * + *

{@link NativeReadable#name()} keys the file in the session's footer cache, exactly as it does when the + * readable is scanned, so reading metadata and then scanning the same file share one footer read. + */ + public static Map readMetadata(Session session, NativeReadable readable) { + Objects.requireNonNull(readable, "readable"); + String name = readable.name(); + Preconditions.checkArgument(name != null && !name.isEmpty(), "readable reported an empty name"); + long length = readable.length(); + Preconditions.checkArgument(length >= 0, "readable for %s reported negative length", name); + return readMetadataFromReadable(session.nativePointer(), readable, name, length); + } + private static native List listFiles(long sessionPointer, String uri, Map options); private static native void delete(long sessionPointer, String[] uris, Map options); + + private static native Map readMetadata( + long sessionPointer, String uri, Map options); + + private static native Map readMetadataFromReadable( + long sessionPointer, Object readable, String name, long length); } diff --git a/java/vortex-jni/src/main/java/dev/vortex/jni/NativeWriter.java b/java/vortex-jni/src/main/java/dev/vortex/jni/NativeWriter.java index f0c0a7471ad..44bbd3b5323 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/jni/NativeWriter.java +++ b/java/vortex-jni/src/main/java/dev/vortex/jni/NativeWriter.java @@ -14,15 +14,26 @@ public final class NativeWriter { private NativeWriter() {} - /** Open a writer at {@code uri} that accepts batches matching the Arrow schema at {@code arrowSchemaAddress}. */ + /** + * Open a writer at {@code uri} that accepts batches matching the Arrow schema at {@code arrowSchemaAddress}. + * + * @param metadata user-defined metadata segments to store in the file footer (may be null) + */ public static native long create( - long sessionPointer, String uri, long arrowSchemaAddress, Map options); + long sessionPointer, + String uri, + long arrowSchemaAddress, + Map options, + Map metadata); /** * Open a writer that streams the file into a caller-provided {@link dev.vortex.io.NativeWritable}. The native side * writes and flushes but never closes the writable; the caller must close it after {@link #close}. + * + * @param metadata user-defined metadata segments to store in the file footer (may be null) */ - public static native long createStream(long sessionPointer, Object writable, long arrowSchemaAddress); + public static native long createStream( + long sessionPointer, Object writable, long arrowSchemaAddress, Map metadata); /** * Write a batch directly from Arrow C Data Interface addresses. diff --git a/java/vortex-jni/src/test/java/dev/vortex/api/GeoTypesTest.java b/java/vortex-jni/src/test/java/dev/vortex/api/GeoTypesTest.java index bdc13f5b3b4..381a07e2945 100644 --- a/java/vortex-jni/src/test/java/dev/vortex/api/GeoTypesTest.java +++ b/java/vortex-jni/src/test/java/dev/vortex/api/GeoTypesTest.java @@ -15,7 +15,6 @@ import java.nio.ByteOrder; import java.nio.file.Path; import java.util.ArrayList; -import java.util.HashMap; import java.util.List; import java.util.Map; import org.apache.arrow.c.ArrowArray; @@ -68,7 +67,8 @@ static void setup() throws IOException { BufferAllocator allocator = ArrowAllocation.rootAllocator(); Schema schema = new Schema(List.of(wkbField("geom"))); Session session = Session.create(); - try (VortexWriter writer = VortexWriter.create(session, writePath, schema, new HashMap<>(), allocator); + try (VortexWriter writer = VortexWriter.builder(session, writePath, schema, allocator) + .build(); VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { VarBinaryVector geomVec = (VarBinaryVector) root.getVector("geom"); geomVec.allocateNew(WKB_POINTS.size()); diff --git a/java/vortex-jni/src/test/java/dev/vortex/api/TestMinimal.java b/java/vortex-jni/src/test/java/dev/vortex/api/TestMinimal.java index 3a6f9ca0883..eb440bca7e9 100644 --- a/java/vortex-jni/src/test/java/dev/vortex/api/TestMinimal.java +++ b/java/vortex-jni/src/test/java/dev/vortex/api/TestMinimal.java @@ -14,7 +14,6 @@ import java.nio.file.Files; import java.nio.file.Path; import java.util.ArrayList; -import java.util.HashMap; import java.util.List; import java.util.Objects; import org.apache.arrow.c.ArrowArray; @@ -86,7 +85,8 @@ static void setup() throws IOException { Field.notNullable("Salary", ArrowType.Decimal.createDecimal(9, 2, 128)), Field.nullable("State", ArrowType.Utf8View.INSTANCE))); Session session = Session.create(); - try (VortexWriter writer = VortexWriter.create(session, writePath, schema, new HashMap<>(), allocator); + try (VortexWriter writer = VortexWriter.builder(session, writePath, schema, allocator) + .build(); VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { ViewVarCharVector nameVec = (ViewVarCharVector) root.getVector("Name"); DecimalVector salaryVec = (DecimalVector) root.getVector("Salary"); diff --git a/java/vortex-jni/src/test/java/dev/vortex/io/NativeIOBridgeTest.java b/java/vortex-jni/src/test/java/dev/vortex/io/NativeIOBridgeTest.java index 5a6d40e05db..e19bac01e3d 100644 --- a/java/vortex-jni/src/test/java/dev/vortex/io/NativeIOBridgeTest.java +++ b/java/vortex-jni/src/test/java/dev/vortex/io/NativeIOBridgeTest.java @@ -4,6 +4,7 @@ package dev.vortex.io; import static java.nio.charset.StandardCharsets.UTF_8; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -16,6 +17,7 @@ import dev.vortex.api.VortexWriteSummary; import dev.vortex.api.VortexWriter; import dev.vortex.arrow.ArrowAllocation; +import dev.vortex.jni.NativeFiles; import dev.vortex.jni.NativeLoader; import java.io.EOFException; import java.io.IOException; @@ -26,6 +28,8 @@ import java.nio.file.Path; import java.nio.file.StandardOpenOption; import java.util.List; +import java.util.Map; +import java.util.Set; import org.apache.arrow.c.ArrowArray; import org.apache.arrow.c.ArrowSchema; import org.apache.arrow.c.Data; @@ -132,7 +136,8 @@ private void writePeopleFile(Session session, BufferAllocator allocator, Path pa Schema schema = personSchema(); VortexWriteSummary summary; try (StreamWritable writable = new StreamWritable(path)) { - try (VortexWriter writer = VortexWriter.create(session, writable, schema, allocator); + try (VortexWriter writer = VortexWriter.builder(session, writable, schema, allocator) + .build(); VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { VarCharVector nameVec = (VarCharVector) root.getVector("name"); IntVector ageVec = (IntVector) root.getVector("age"); @@ -198,6 +203,49 @@ public void testWritableThenReadableRoundTrip() throws IOException { } } + @Test + public void testMetadataRoundTripThroughBridge() throws IOException { + Path path = tempDir.resolve("bridge_metadata.vortex"); + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + Schema schema = personSchema(); + byte[] value = "{\"type\":\"struct\",\"fields\":[]}".getBytes(UTF_8); + + try (StreamWritable writable = new StreamWritable(path)) { + try (VortexWriter writer = VortexWriter.builder(session, writable, schema, allocator) + .putMetadata("iceberg.schema", value) + .build(); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + VarCharVector nameVec = (VarCharVector) root.getVector("name"); + IntVector ageVec = (IntVector) root.getVector("age"); + nameVec.allocateNew(1); + ageVec.allocateNew(1); + nameVec.setSafe(0, "Alice".getBytes(UTF_8)); + ageVec.setSafe(0, 30); + root.setRowCount(1); + try (ArrowArray arrowArray = ArrowArray.allocateNew(allocator); + ArrowSchema arrowSchemaFfi = ArrowSchema.allocateNew(allocator)) { + Data.exportVectorSchemaRoot(allocator, root, null, arrowArray, arrowSchemaFfi); + writer.writeBatch(arrowArray.memoryAddress(), arrowSchemaFfi.memoryAddress()); + } + } + } + + // Both directions go through caller-provided I/O: no native storage client is involved. + try (FileChannelReadable readable = new FileChannelReadable(path)) { + Map metadata = NativeFiles.readMetadata(session, readable); + assertEquals(Set.of("iceberg.schema"), metadata.keySet()); + assertArrayEquals(value, metadata.get("iceberg.schema")); + } + + // The readable's name keys the session footer cache, so this read resolves the metadata + // against the footer cached by the read above rather than reading it again. + try (FileChannelReadable readable = new FileChannelReadable(path)) { + Map metadata = NativeFiles.readMetadata(session, readable); + assertArrayEquals(value, metadata.get("iceberg.schema")); + } + } + @Test public void testMultipleReadables() throws IOException { Path first = tempDir.resolve("bridge_a.vortex"); @@ -256,7 +304,8 @@ public void testLargeRoundTrip() throws IOException { int batches = 20; int rowsPerBatch = 5_000; try (StreamWritable writable = new StreamWritable(path)) { - try (VortexWriter writer = VortexWriter.create(session, writable, schema, allocator); + try (VortexWriter writer = VortexWriter.builder(session, writable, schema, allocator) + .build(); VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { VarCharVector nameVec = (VarCharVector) root.getVector("name"); IntVector ageVec = (IntVector) root.getVector("age"); diff --git a/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java b/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java index 59132b78168..e324348f8a5 100644 --- a/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java +++ b/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java @@ -6,7 +6,9 @@ import static java.nio.charset.StandardCharsets.UTF_8; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import dev.vortex.api.DataSource; @@ -112,7 +114,9 @@ public void testCreateWriter() throws IOException { Map options = new HashMap<>(); Session session = Session.create(); - try (VortexWriter writer = VortexWriter.create(session, writePath, personSchema(), options, allocator)) { + try (VortexWriter writer = VortexWriter.builder(session, writePath, personSchema(), allocator) + .options(options) + .build()) { assertNotNull(writer); } @@ -128,7 +132,9 @@ public void testCreateWriterPlainLocalPath() throws IOException { Map options = new HashMap<>(); Session session = Session.create(); - try (VortexWriter writer = VortexWriter.create(session, writePath, personSchema(), options, allocator)) { + try (VortexWriter writer = VortexWriter.builder(session, writePath, personSchema(), allocator) + .options(options) + .build()) { assertNotNull(writer); } @@ -144,7 +150,9 @@ public void testCreateWriterCreatesParentDirectories() throws IOException { Map options = new HashMap<>(); Session session = Session.create(); - try (VortexWriter writer = VortexWriter.create(session, writePath, personSchema(), options, allocator)) { + try (VortexWriter writer = VortexWriter.builder(session, writePath, personSchema(), allocator) + .options(options) + .build()) { assertNotNull(writer); } @@ -162,7 +170,8 @@ public void testWriteBatch() throws IOException { Session session = Session.create(); VortexWriteSummary summary; long bytesWhileOpen; - try (VortexWriter writer = VortexWriter.create(session, writePath, schema, new HashMap<>(), allocator); + try (VortexWriter writer = VortexWriter.builder(session, writePath, schema, allocator) + .build(); VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { VarCharVector nameVec = (VarCharVector) root.getVector("name"); IntVector ageVec = (IntVector) root.getVector("age"); @@ -225,6 +234,95 @@ public void testWriteBatch() throws IOException { } } + /** Write a single three-row batch of {@link #personSchema()} data. */ + private static void writePeopleBatch(VortexWriter writer, BufferAllocator allocator) throws IOException { + try (VectorSchemaRoot root = VectorSchemaRoot.create(personSchema(), allocator)) { + VarCharVector nameVec = (VarCharVector) root.getVector("name"); + IntVector ageVec = (IntVector) root.getVector("age"); + nameVec.allocateNew(3); + ageVec.allocateNew(3); + nameVec.setSafe(0, "Alice".getBytes(UTF_8)); + nameVec.setSafe(1, "Bob".getBytes(UTF_8)); + nameVec.setSafe(2, "Carol".getBytes(UTF_8)); + ageVec.setSafe(0, 30); + ageVec.setSafe(1, 25); + ageVec.setSafe(2, 40); + root.setRowCount(3); + + try (ArrowArray arrowArray = ArrowArray.allocateNew(allocator); + ArrowSchema arrowSchemaFfi = ArrowSchema.allocateNew(allocator)) { + Data.exportVectorSchemaRoot(allocator, root, null, arrowArray, arrowSchemaFfi); + writer.writeBatch(arrowArray.memoryAddress(), arrowSchemaFfi.memoryAddress()); + } + } + } + + @Test + public void testFileMetadataRoundTrip() throws IOException { + Path outputPath = tempDir.resolve("test_metadata.vortex"); + String writePath = outputPath.toAbsolutePath().toUri().toString(); + + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + byte[] schemaJson = "{\"type\":\"struct\",\"fields\":[]}".getBytes(UTF_8); + byte[] deleteType = "position".getBytes(UTF_8); + Map metadata = Map.of("iceberg.schema", schemaJson, "delete-type", deleteType); + + try (VortexWriter writer = VortexWriter.builder(session, writePath, personSchema(), allocator) + .metadata(metadata) + .build()) { + writePeopleBatch(writer, allocator); + } + + Map read = NativeFiles.readMetadata(session, writePath, new HashMap<>()); + assertEquals(metadata.keySet(), read.keySet()); + assertArrayEquals(schemaJson, read.get("iceberg.schema")); + assertArrayEquals(deleteType, read.get("delete-type")); + + // Metadata is stored out of band: the file still scans as usual. + DataSource ds = DataSource.open(session, writePath); + assertEquals(new DataSource.RowCount.Exact(3L), ds.rowCount()); + } + + @Test + public void testFileWithoutMetadataReadsEmpty() throws IOException { + Path outputPath = tempDir.resolve("test_no_metadata.vortex"); + String writePath = outputPath.toAbsolutePath().toUri().toString(); + + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + try (VortexWriter writer = VortexWriter.builder(session, writePath, personSchema(), allocator) + .build()) { + writePeopleBatch(writer, allocator); + } + + assertTrue(NativeFiles.readMetadata(session, writePath, new HashMap<>()).isEmpty()); + } + + @Test + public void testInvalidMetadataKeyRejectedAtCreate() { + Path outputPath = tempDir.resolve("test_bad_metadata.vortex"); + String writePath = outputPath.toAbsolutePath().toUri().toString(); + + BufferAllocator allocator = ArrowAllocation.rootAllocator(); + Session session = Session.create(); + // Keys are capped well below this; the writer must reject the set before any bytes + // are produced, rather than failing the first writeBatch with a send error. + Map metadata = Map.of("k".repeat(1024), new byte[] {1}); + + IOException thrown = assertThrows( + IOException.class, () -> VortexWriter.builder(session, writePath, personSchema(), allocator) + .metadata(metadata) + .build()); + Throwable cause = thrown.getCause(); + assertNotNull(cause, "native failure should be retained as the cause"); + assertTrue( + cause.getMessage().contains("metadata key"), + "error should identify the offending key, got: " + cause.getMessage()); + // Rejected before the sink is touched, so no partial file is left behind. + assertFalse(Files.exists(outputPath)); + } + @Test public void testParquetVariantRoundTrip() throws IOException { Path outputPath = tempDir.resolve("test_parquet_variant.vortex"); @@ -234,7 +332,8 @@ public void testParquetVariantRoundTrip() throws IOException { Schema schema = parquetVariantSchema(); Session session = Session.create(); - try (VortexWriter writer = VortexWriter.create(session, writePath, schema, new HashMap<>(), allocator); + try (VortexWriter writer = VortexWriter.builder(session, writePath, schema, allocator) + .build(); VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { populateParquetVariantRoot(root); diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexDataWriter.java b/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexDataWriter.java index 3a4390cb8b6..5a48547a7e3 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexDataWriter.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/write/VortexDataWriter.java @@ -123,8 +123,9 @@ public final class VortexDataWriter implements DataWriter, AutoClos var arrowSchema = SparkToArrowSchema.convert(schema); this.session = VortexSparkSession.get(options.asCaseSensitiveMap()); - this.vortexWriter = - VortexWriter.create(session, filePath, arrowSchema, options.asCaseSensitiveMap(), allocator); + this.vortexWriter = VortexWriter.builder(session, filePath, arrowSchema, allocator) + .options(options.asCaseSensitiveMap()) + .build(); this.vectorSchemaRoot = VectorSchemaRoot.create(arrowSchema, allocator); logger.debug("Initialized VortexDataWriter for {}", filePath); diff --git a/vortex-file/src/multi/mod.rs b/vortex-file/src/multi/mod.rs index 019c59a4113..e94f9069ad1 100644 --- a/vortex-file/src/multi/mod.rs +++ b/vortex-file/src/multi/mod.rs @@ -19,6 +19,7 @@ pub use uri::parse_uri_or_path; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_io::VortexReadAt; use vortex_io::filesystem::FileListing; use vortex_io::filesystem::FileSystemRef; use vortex_layout::LayoutReaderRef; @@ -232,20 +233,37 @@ async fn open_file( ) -> VortexResult { tracing::trace!(path = %file.path, "opening vortex file"); - // Open the reader first so we can use its URI as the cache key. - // The URI includes the full path (with any filesystem prefix), making it unique - // even when different PrefixFileSystem instances strip paths to the same relative name. let source = fs.open_read(&file.path).await?; + open_cached(session, source, &file.path, file.size, open_options_fn).await +} + +/// Open a single Vortex file through the session's footer cache, so that a later open of the +/// same file skips the footer read. +/// +/// The cache is keyed by the source's [`uri`](vortex_io::VortexReadAt::uri) where it reports one, +/// since that includes the full path (with any filesystem prefix) and so stays unique even when +/// different filesystems strip paths to the same relative name. `fallback_key` identifies the file +/// for sources that report no URI, and must be stable and unique within the session — two +/// different files sharing a key would read each other's footer. +/// +/// Caching the footer is independent of [`VortexOpenOptions::include_metadata`]: the footer holds +/// only metadata *locators*, and each open resolves the segments it was asked for. +pub async fn open_cached( + session: &VortexSession, + source: Arc, + fallback_key: &str, + file_size: Option, + open_options_fn: &(dyn Fn(VortexOpenOptions) -> VortexOpenOptions + Send + Sync), +) -> VortexResult { let cache_key = source .uri() - .map(|u| u.to_string()) - .unwrap_or_else(|| file.path.clone()); + .map_or_else(|| fallback_key.to_owned(), |uri| uri.to_string()); - // Build open options. The DashMap Ref from multi_file() must not live across an await, + // Build open options. The cache guard from multi_file() must not live across an await, // so we scope the cache lookup in a block. let options = { let mut options = open_options_fn(session.open_options()); - if let Some(size) = file.size { + if let Some(size) = file_size { options = options.with_file_size(size); } if let Some(footer) = session.multi_file().get_footer(&cache_key) { @@ -256,7 +274,7 @@ async fn open_file( let vortex_file = options.open(source).await?; - // Store footer in cache (scoped to avoid holding the Ref across subsequent code). + // Store footer in cache (scoped to avoid holding the guard across subsequent code). session .multi_file() .put_footer(&cache_key, vortex_file.footer().clone()); diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 58d1cb7cd4f..2d013b521c8 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -160,6 +160,15 @@ impl VortexWriteOptions { } self } + + /// Check the configured metadata segments against the `MAX_METADATA_*` limits. + /// + /// [`Self::write`] performs the same check, but only once the sink is already being written + /// to. Callers that accept metadata from elsewhere (FFI bindings, for example) can use this to + /// reject an invalid set before any bytes are produced. + pub fn validate_metadata(&self) -> VortexResult<()> { + validate_metadata_segments(&self.metadata) + } } impl VortexWriteOptions { @@ -601,3 +610,56 @@ impl WriteSummary { .collect()) } } + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_buffer::ByteBuffer; + + use super::*; + + fn write_options_with_keys(keys: &[String]) -> VortexWriteOptions { + vortex_array::array_session() + .write_options() + .with_metadata_segments( + keys.iter() + .map(|key| (key.clone(), ByteBuffer::copy_from(b"value"))), + ) + } + + #[rstest] + #[case::empty_key(vec![String::new()], "non-empty")] + #[case::oversized_key(vec!["k".repeat(MAX_METADATA_KEY_BYTES + 1)], "keys must be at most")] + // The cap is on bytes, not characters. + #[case::oversized_multibyte_key( + vec!["é".repeat(MAX_METADATA_KEY_BYTES / "é".len() + 1)], + "keys must be at most" + )] + #[case::too_many_segments( + (0..=MAX_METADATA_SEGMENTS).map(|idx| format!("key-{idx}")).collect(), + "at most 16 metadata segments" + )] + fn validate_metadata_rejects(#[case] keys: Vec, #[case] expected: &str) { + let Err(error) = write_options_with_keys(&keys).validate_metadata() else { + panic!("metadata must be rejected for {keys:?}"); + }; + assert!( + error.to_string().contains(expected), + "error should mention {expected:?}, got: {error}" + ); + } + + #[test] + fn validate_metadata_accepts_the_limits() -> VortexResult<()> { + // Distinct keys, each exactly at the key-length cap. + let keys = (0..MAX_METADATA_SEGMENTS) + .map(|idx| format!("{idx:0>width$}", width = MAX_METADATA_KEY_BYTES)) + .collect::>(); + write_options_with_keys(&keys).validate_metadata() + } + + #[test] + fn validate_metadata_accepts_no_metadata() -> VortexResult<()> { + write_options_with_keys(&[]).validate_metadata() + } +} diff --git a/vortex-jni/src/file.rs b/vortex-jni/src/file.rs index 343215e533a..1044ed7dc01 100644 --- a/vortex-jni/src/file.rs +++ b/vortex-jni/src/file.rs @@ -1,10 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Small utility JNI entry points for listing and deleting Vortex files via an object store. +//! Small utility JNI entry points for listing, deleting, and reading the user-defined +//! metadata of Vortex files. + +use std::sync::Arc; use futures::StreamExt; use jni::EnvUnowned; +use jni::objects::JByteArray; use jni::objects::JClass; use jni::objects::JMap; use jni::objects::JObject; @@ -13,15 +17,20 @@ use jni::objects::JString; use jni::sys::jlong; use jni::sys::jobject; use object_store::path::Path; +use vortex::buffer::ByteBuffer; use vortex::error::VortexResult; use vortex::error::vortex_err; +use vortex::file::multi::open_cached; use vortex::file::multi::parse_uri_or_path; +use vortex::io::VortexReadAt; use vortex::io::runtime::BlockingRuntime; use vortex::io::session::RuntimeSessionExt; +use vortex::session::VortexSession; use vortex::utils::aliases::hash_map::HashMap; use crate::RUNTIME; use crate::errors::try_or_throw; +use crate::io::java_readable; use crate::object_store::object_store_fs; use crate::session::session_ref; @@ -46,6 +55,161 @@ pub(crate) fn extract_properties( Ok(properties) } +/// Extract a Java `Map` of user-defined metadata segments. +/// +/// Values are opaque bytes, so they are copied out of the Java arrays rather than +/// interpreted. A null map yields an empty set of segments. +pub(crate) fn extract_metadata( + env: &mut jni::Env, + metadata: &JObject, +) -> Result, crate::errors::JNIError> { + let mut segments = HashMap::new(); + if metadata.is_null() { + return Ok(segments); + } + + let metadata_ref = env.new_local_ref(metadata)?; + let map = env.cast_local::(metadata_ref)?; + let mut iterator = map.iter(env)?; + while let Some(entry) = iterator.next(env)? { + let key_obj = entry.key(env)?; + let val_obj = entry.value(env)?; + let key_str = env.cast_local::(key_obj)?; + let key = key_str.try_to_string(env)?; + if val_obj.is_null() { + return Err(vortex_err!("null metadata value for key '{key}'").into()); + } + let bytes = env.cast_local::(val_obj)?; + segments.insert(key, ByteBuffer::from(env.convert_byte_array(&bytes)?)); + // Each entry costs three local refs, so release them as we go rather than relying on + // the frame's guaranteed capacity. + env.delete_local_ref(bytes); + env.delete_local_ref(key_str); + env.delete_local_ref(entry); + } + + Ok(segments) +} + +/// Build a `java.util.HashMap` from user-defined metadata segments. +fn metadata_to_java( + env: &mut jni::Env, + segments: Vec<(String, ByteBuffer)>, +) -> Result { + let map = env.new_object( + jni::jni_str!("java/util/HashMap"), + jni::jni_sig!("()V"), + &[], + )?; + let raw = map.as_raw(); + let map = env.cast_local::(map)?; + for (key, value) in segments { + let key = env.new_string(key)?; + let value = env.byte_array_from_slice(value.as_slice())?; + // A file may hold several segments and each entry costs two local refs, so release + // them as we go rather than relying on the frame's guaranteed capacity. + if let Some(previous) = map.put(env, key.as_ref(), value.as_ref())? { + env.delete_local_ref(previous); + } + env.delete_local_ref(value); + env.delete_local_ref(key); + } + + Ok(raw) +} + +/// Open a Vortex file with metadata resolution enabled and collect its metadata segments. +/// +/// `cache_key` identifies the file in the session footer cache. Both entry points name the file +/// the same way the scan path does — the object-store path, or the `NativeReadable`'s name — so a +/// metadata read and a later scan of the same file share one footer read. +fn read_metadata_segments( + session: &VortexSession, + source: Arc, + cache_key: &str, + file_size: Option, +) -> VortexResult> { + RUNTIME.block_on(async move { + let file = open_cached(session, source, cache_key, file_size, &|options| { + options.include_metadata() + }) + .await?; + Ok(file + .metadata_segments() + .map(|(key, value)| (key.to_string(), value.clone())) + .collect()) + }) +} + +/// Read the user-defined metadata segments of the Vortex file at `uri`. +/// +/// Returns a `java.util.HashMap`, empty when the file carries no metadata. +/// Metadata lives in its own segments, so this may issue a read beyond the file tail. +#[unsafe(no_mangle)] +pub extern "system" fn Java_dev_vortex_jni_NativeFiles_readMetadata( + mut env: EnvUnowned, + _class: JClass, + session_ptr: jlong, + uri: JString, + options: JObject, +) -> jobject { + try_or_throw(&mut env, |env| { + let session = unsafe { session_ref(session_ptr) }; + let uri: String = uri.try_to_string(env)?; + let url = parse_uri_or_path(&uri)?; + let properties = extract_properties(env, &options)?; + + let fs = object_store_fs(&url, &properties, session.handle())?; + // `FileSystem` keys are literal, already-decoded paths, so decode as `listFiles` does. + let path = Path::from_url_path(url.path()) + .map_err(|_| vortex_err!("cannot parse uri as object_store Path"))? + .to_string(); + let source = RUNTIME.block_on(async { fs.open_read(&path).await })?; + + let segments = read_metadata_segments(session, source, &path, None)?; + metadata_to_java(env, segments) + }) +} + +/// Read the user-defined metadata segments of a Vortex file through a caller-provided +/// `dev.vortex.io.NativeReadable`, so no native storage client is created. +/// +/// `name` is the readable's `name()`, which the interface requires to be stable and unique; it +/// keys the file in the session footer cache, as it does for the scan path. +#[unsafe(no_mangle)] +pub extern "system" fn Java_dev_vortex_jni_NativeFiles_readMetadataFromReadable( + mut env: EnvUnowned, + _class: JClass, + session_ptr: jlong, + readable: JObject, + name: JString, + length: jlong, +) -> jobject { + try_or_throw(&mut env, |env| { + let session = unsafe { session_ref(session_ptr) }; + if readable.is_null() { + throw_runtime!("null readable"); + } + let name: String = name.try_to_string(env)?; + if name.is_empty() { + throw_runtime!("readable name must not be empty"); + } + let length = + u64::try_from(length).map_err(|_| vortex_err!("negative readable length: {length}"))?; + + let vm = env.get_java_vm()?; + let readable = Arc::new(env.new_global_ref(&readable)?); + let source = java_readable(vm, readable, length, session.handle()); + + // `MultiFileDataSource` strips leading slashes and the scan path keys Java readables by + // that same normalized form, so normalize here too — otherwise a metadata read and a + // scan of the same readable would land on two different cache entries. + let cache_key = name.trim_start_matches('/'); + let segments = read_metadata_segments(session, source, cache_key, Some(length))?; + metadata_to_java(env, segments) + }) +} + /// List Vortex files under the given URI prefix. Returns a `java.util.ArrayList`. #[unsafe(no_mangle)] pub extern "system" fn Java_dev_vortex_jni_NativeFiles_listFiles( diff --git a/vortex-jni/src/io/mod.rs b/vortex-jni/src/io/mod.rs index fa20a6a4685..d1f2c14df44 100644 --- a/vortex-jni/src/io/mod.rs +++ b/vortex-jni/src/io/mod.rs @@ -32,6 +32,7 @@ mod write; use jni::Env; use jni::JavaVM; pub(crate) use read_at::JavaFileSystem; +pub(crate) use read_at::java_readable; use vortex::error::VortexError; use vortex::error::VortexResult; pub(crate) use write::JavaWrite; diff --git a/vortex-jni/src/io/read_at.rs b/vortex-jni/src/io/read_at.rs index a08652c7cc1..2c45fc347cd 100644 --- a/vortex-jni/src/io/read_at.rs +++ b/vortex-jni/src/io/read_at.rs @@ -195,6 +195,26 @@ impl VortexReadAt for JavaReadable { } } +/// Wrap a single Java `NativeReadable` as a [`VortexReadAt`], without registering it in a +/// [`JavaFileSystem`]. +/// +/// For one-shot reads (opening a file to inspect its metadata, say) where there is no glob to +/// resolve and no other file to share an upcall budget with. +pub(crate) fn java_readable( + vm: JavaVM, + readable: Arc>>, + len: u64, + handle: Handle, +) -> Arc { + Arc::new(JavaReadable::new( + vm, + readable, + len, + handle, + UpcallLimiter::new(DEFAULT_CONCURRENCY), + )) +} + struct JavaFileEntry { readable: Arc>>, size: u64, diff --git a/vortex-jni/src/writer.rs b/vortex-jni/src/writer.rs index db250ec17be..52f888e4db5 100644 --- a/vortex-jni/src/writer.rs +++ b/vortex-jni/src/writer.rs @@ -69,6 +69,7 @@ use vortex_parquet_variant::ParquetVariant; use crate::RUNTIME; use crate::errors::JNIError; use crate::errors::try_or_throw; +use crate::file::extract_metadata; use crate::file::extract_properties; use crate::io::JavaWrite; use crate::object_store::make_object_store; @@ -401,6 +402,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_create( uri: JString, arrow_schema_addr: jlong, options: JObject, + metadata: JObject, ) -> jlong { try_or_throw(&mut env, |env| { if session_ptr == 0 { @@ -417,11 +419,18 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_create( let file_path: String = uri.try_to_string(env)?; let properties: HashMap = extract_properties(env, &options)?; + let metadata = extract_metadata(env, &metadata)?; let resolved = resolve_store(&file_path, &properties)?; let (tx, rx) = mpsc::channel(WRITE_CHANNEL_CAPACITY); let stream = ArrayStreamAdapter::new(write_schema.clone(), rx); let strategy = write_strategy_for_schema(&write_schema); - let write_options = session.write_options().with_strategy(Arc::clone(&strategy)); + let write_options = session + .write_options() + .with_strategy(Arc::clone(&strategy)) + .with_metadata_segments(metadata); + // The same check runs inside `write`, but only once the write task is under way, where + // it would surface as an opaque send failure on the first batch. + write_options.validate_metadata()?; let (bytes_written, handle) = match resolved { ResolvedStore::Path(path) => { @@ -480,6 +489,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_createStream( session_ptr: jlong, writable: JObject, arrow_schema_addr: jlong, + metadata: JObject, ) -> jlong { try_or_throw(&mut env, |env| { if session_ptr == 0 { @@ -497,12 +507,18 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_createStream( let arrow_schema = Arc::new(Schema::try_from(ffi_schema)?); let write_schema = session.arrow().from_arrow_schema(arrow_schema.as_ref())?; + let metadata = extract_metadata(env, &metadata)?; let vm = env.get_java_vm()?; let writable = Arc::new(env.new_global_ref(&writable)?); let (tx, rx) = mpsc::channel(WRITE_CHANNEL_CAPACITY); let stream = ArrayStreamAdapter::new(write_schema.clone(), rx); let strategy = write_strategy_for_schema(&write_schema); - let write_options = session.write_options().with_strategy(Arc::clone(&strategy)); + let write_options = session + .write_options() + .with_strategy(Arc::clone(&strategy)) + .with_metadata_segments(metadata); + // See the note in `create`: validate before the write task can start. + write_options.validate_metadata()?; let mut write = CountingVortexWrite::new(JavaWrite::new(vm, writable)); let bytes_written = write.counter();