From b1117156ad5dd01524f3cf11a55a82a3ebd1d156 Mon Sep 17 00:00:00 2001 From: rayshrey Date: Wed, 15 Jul 2026 17:52:38 +0530 Subject: [PATCH 1/2] Replace Writer's DashMap on Rust side with pointer handle owned by Java side Signed-off-by: rayshrey --- .../parquet/bridge/NativeParquetWriter.java | 101 +++- .../parquet/bridge/ParquetFileMetadata.java | 2 +- .../opensearch/parquet/bridge/RustBridge.java | 47 +- .../parquet/engine/ParquetIndexingEngine.java | 22 +- .../opensearch/parquet/vsr/VSRManager.java | 13 +- .../parquet/writer/ParquetWriter.java | 20 +- .../src/main/rust/src/ffm.rs | 37 +- .../src/main/rust/src/test_utils.rs | 86 ++- .../src/main/rust/src/tests/mod.rs | 534 ++++++++---------- .../src/main/rust/src/writer.rs | 402 +++++++------ .../rust/tests/writer_integration_tests.rs | 78 ++- .../NativeParquetWriterRecoveryTests.java | 99 ++++ .../bridge/NativeParquetWriterTests.java | 27 +- 13 files changed, 850 insertions(+), 618 deletions(-) create mode 100644 sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/bridge/NativeParquetWriterRecoveryTests.java diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeParquetWriter.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeParquetWriter.java index ee18effdb070a..a438865c26e56 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeParquetWriter.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeParquetWriter.java @@ -14,6 +14,7 @@ import org.opensearch.plugin.stats.StatsRecorder; import java.io.IOException; +import java.lang.ref.Cleaner; import java.util.concurrent.atomic.AtomicBoolean; /** @@ -39,6 +40,20 @@ public class NativeParquetWriter { private final ParquetShardStatsTracker stats; private volatile boolean initialized = false; + /** Reclaims leaked native writers if this object is GC'd without an explicit close/flush. */ + private static final Cleaner CLEANER = Cleaner.create(); + + /** + * Opaque native handle (a {@code Box} pointer) minted by {@link #initialize}; + * 0 until initialized. Every native call that dereferences it runs inside a {@code synchronized} + * method, so they are serialized (write/finalize/free/memory never race on the same handle). + */ + private volatile long handle = 0L; + /** Set once the handle has been consumed (by finalize) or freed (by close/Cleaner). */ + private final AtomicBoolean released = new AtomicBoolean(false); + /** Cleaner registration; {@code clean()} on close deregisters the GC backstop. */ + private Cleaner.Cleanable cleanable; + /** * Creates a new NativeParquetWriter handle. Does not create the native writer — * call {@link #initialize(String, long, ParquetSortConfig, long)} before the first write. @@ -71,12 +86,16 @@ public NativeParquetWriter(String filePath) { * @throws IOException if the native writer creation fails * @throws IllegalStateException if already initialized */ - public void initialize(String indexName, long schemaAddress, ParquetSortConfig sortConfig, long writerGeneration) throws IOException { + public synchronized void initialize(String indexName, long schemaAddress, ParquetSortConfig sortConfig, long writerGeneration) + throws IOException { if (initialized) { throw new IllegalStateException("Writer already initialized: " + filePath); } - RustBridge.createWriter(filePath, indexName, schemaAddress, sortConfig, writerGeneration); + handle = RustBridge.createWriter(filePath, indexName, schemaAddress, sortConfig, writerGeneration); initialized = true; + // GC backstop: if this writer is dropped without close()/flush(), free the native handle. + // HandleCleanup holds only the primitive handle + shared released flag (never `this`). + cleanable = CLEANER.register(this, new HandleCleanup(handle, released)); } /** @@ -96,7 +115,7 @@ public boolean isInitialized() { * @throws IOException if the write fails or the writer is flushed * @throws IllegalStateException if the writer has not been initialized */ - public void write(long arrayAddress, long schemaAddress) throws IOException { + public synchronized void write(long arrayAddress, long schemaAddress) throws IOException { if (writerFlushed.get()) { throw new IOException("Cannot write to flushed Parquet writer: " + filePath); } @@ -104,7 +123,7 @@ public void write(long arrayAddress, long schemaAddress) throws IOException { throw new IllegalStateException("Writer not initialized: " + filePath); } StatsRecorder.recordOutcome( - () -> RustBridge.write(filePath, arrayAddress, schemaAddress), + () -> RustBridge.write(handle, arrayAddress, schemaAddress), stats::addNativeWriteTimeMillis, stats::incNativeWriteTotal, stats::incNativeWriteFailures @@ -119,18 +138,24 @@ public void write(long arrayAddress, long schemaAddress) throws IOException { * @return the file metadata, or null if the writer was never initialized * @throws IOException if the finalization fails */ - public ParquetFileMetadata flush() throws IOException { + public synchronized ParquetFileMetadata flush() throws IOException { if (writerFlushed.compareAndSet(false, true)) { if (initialized) { - StatsRecorder.recordOutcome(() -> { - RustBridge.WriterFinalizeResult result = RustBridge.finalizeWriter(filePath); - if (result != null) { - metadata.set(result.metadata()); - if (result.rowIdMapping() != null) { - rowIdMapping.set(result.rowIdMapping()); + try { + StatsRecorder.recordOutcome(() -> { + RustBridge.WriterFinalizeResult result = RustBridge.finalizeWriter(handle); + if (result != null) { + metadata.set(result.metadata()); + if (result.rowIdMapping() != null) { + rowIdMapping.set(result.rowIdMapping()); + } } - } - }, stats::addNativeFinalizeTimeMillis, stats::incNativeFinalizeTotal, stats::incNativeFinalizeFailures); + }, stats::addNativeFinalizeTimeMillis, stats::incNativeFinalizeTotal, stats::incNativeFinalizeFailures); + } finally { + // finalize_writer reclaims the native Box regardless of outcome, so the + // handle is consumed — mark released so close()/Cleaner never double-free. + released.set(true); + } } } return metadata.get(); @@ -153,4 +178,54 @@ public RowIdMapping getRowIdMapping() { return rowIdMapping.get(); } + /** + * Returns the native memory currently reserved by this writer, or 0 if it was never + * initialized or has already been finalized/freed. Blocks if a native write/finalize is in + * flight (stats tolerate waiting; the write path is unaffected functionally). + */ + public synchronized long getNativeBytesUsed() { + if (initialized == false || released.get()) { + return 0L; + } + return RustBridge.getWriterMemoryUsage(handle); + } + + /** + * Releases the native writer if it was not already consumed by {@link #flush()}. Idempotent — + * safe to call multiple times and after flush. This is the guaranteed cleanup path invoked on + * writer teardown; the {@link Cleaner} is only a backstop for the case where {@code close()} is + * never called (e.g. the writer is dropped after a failure). + */ + public synchronized void close() { + if (handle != 0L && released.compareAndSet(false, true)) { + RustBridge.freeWriter(handle); + } + if (cleanable != null) { + cleanable.clean(); + } + } + + /** + * Cleaner action reclaiming the native writer if the {@link NativeParquetWriter} becomes + * unreachable without an explicit {@link #close()}/{@link #flush()}. It captures only the + * primitive handle and the shared {@code released} flag — never the enclosing writer — so it + * cannot keep the writer reachable, and the CAS guarantees free happens at most once. + */ + private static final class HandleCleanup implements Runnable { + private final long handle; + private final AtomicBoolean released; + + HandleCleanup(long handle, AtomicBoolean released) { + this.handle = handle; + this.released = released; + } + + @Override + public void run() { + if (handle != 0L && released.compareAndSet(false, true)) { + RustBridge.freeWriter(handle); + } + } + } + } diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/ParquetFileMetadata.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/ParquetFileMetadata.java index 9daf1f636851a..153a3bccba957 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/ParquetFileMetadata.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/ParquetFileMetadata.java @@ -11,7 +11,7 @@ /** * Metadata extracted from a Parquet file after the native writer is closed. * - *

Returned by {@link RustBridge#finalizeWriter(String)} and {@link RustBridge#getFileMetadata(String)}. + *

Returned by {@link RustBridge#finalizeWriter(long)} and {@link RustBridge#getFileMetadata(String)}. * Contains the Parquet format version, total row count, the creator identifier string * embedded in the file footer, and the whole-file CRC32 checksum computed during write. * diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/RustBridge.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/RustBridge.java index a60d0e758c63e..43718cab40eb6 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/RustBridge.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/RustBridge.java @@ -43,7 +43,8 @@ public class RustBridge { private static final MethodHandle FINALIZE_WRITER; private static final MethodHandle GET_FILE_METADATA; private static final MethodHandle GET_COLUMN_METADATA; - private static final MethodHandle GET_FILTERED_BYTES; + private static final MethodHandle GET_WRITER_MEMORY_USAGE; + private static final MethodHandle FREE_WRITER; private static final MethodHandle ON_SETTINGS_UPDATE; private static final MethodHandle REMOVE_SETTINGS; private static final MethodHandle MERGE_FILES; @@ -83,18 +84,16 @@ public class RustBridge { lib.find("parquet_write").orElseThrow(), FunctionDescriptor.of( ValueLayout.JAVA_LONG, - ValueLayout.ADDRESS, - ValueLayout.JAVA_LONG, - ValueLayout.JAVA_LONG, - ValueLayout.JAVA_LONG + ValueLayout.JAVA_LONG, // handle + ValueLayout.JAVA_LONG, // array_address + ValueLayout.JAVA_LONG // schema_address ) ); FINALIZE_WRITER = linker.downcallHandle( lib.find("parquet_finalize_writer").orElseThrow(), FunctionDescriptor.of( ValueLayout.JAVA_LONG, - ValueLayout.ADDRESS, - ValueLayout.JAVA_LONG, + ValueLayout.JAVA_LONG, // handle ValueLayout.ADDRESS, ValueLayout.ADDRESS, ValueLayout.ADDRESS, @@ -120,9 +119,13 @@ public class RustBridge { ValueLayout.ADDRESS // num_row_groups_out ) ); - GET_FILTERED_BYTES = linker.downcallHandle( - lib.find("parquet_get_filtered_native_bytes_used").orElseThrow(), - FunctionDescriptor.of(ValueLayout.JAVA_LONG, ValueLayout.ADDRESS, ValueLayout.JAVA_LONG) + GET_WRITER_MEMORY_USAGE = linker.downcallHandle( + lib.find("parquet_get_writer_memory_usage").orElseThrow(), + FunctionDescriptor.of(ValueLayout.JAVA_LONG, ValueLayout.JAVA_LONG) + ); + FREE_WRITER = linker.downcallHandle( + lib.find("parquet_free_writer").orElseThrow(), + FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG) ); GET_COLUMN_METADATA = linker.downcallHandle( lib.find("parquet_get_column_metadata").orElseThrow(), @@ -292,7 +295,7 @@ public class RustBridge { public static void initLogger() {} - static void createWriter(String file, String indexName, long schemaAddress, ParquetSortConfig sortConfig, long writerGeneration) + static long createWriter(String file, String indexName, long schemaAddress, ParquetSortConfig sortConfig, long writerGeneration) throws IOException { try (var call = new NativeCall()) { var f = call.str(file); @@ -300,7 +303,7 @@ static void createWriter(String file, String indexName, long schemaAddress, Parq var sorts = call.strArray(sortConfig.sortColumns().toArray(new String[0])); var reverseArray = marshalBoolList(call, sortConfig.reverseSorts()); var nullsFirstArray = marshalBoolList(call, sortConfig.nullsFirst()); - call.invokeIO( + return call.invokeIO( CREATE_WRITER, f.segment(), f.len(), @@ -319,10 +322,9 @@ static void createWriter(String file, String indexName, long schemaAddress, Parq } } - static void write(String file, long arrayAddress, long schemaAddress) throws IOException { + static void write(long handle, long arrayAddress, long schemaAddress) throws IOException { try (var call = new NativeCall()) { - var f = call.str(file); - call.invokeIO(WRITE, f.segment(), f.len(), arrayAddress, schemaAddress); + call.invokeIO(WRITE, handle, arrayAddress, schemaAddress); } } @@ -332,9 +334,8 @@ static void write(String file, long arrayAddress, long schemaAddress) throws IOE record WriterFinalizeResult(ParquetFileMetadata metadata, RowIdMapping rowIdMapping) { } - static WriterFinalizeResult finalizeWriter(String file) throws IOException { + static WriterFinalizeResult finalizeWriter(long handle) throws IOException { try (var call = new NativeCall()) { - var f = call.str(file); var versionOut = call.intOut(); var numRowsOut = call.longOut(); var crc32Out = call.longOut(); @@ -344,8 +345,7 @@ static WriterFinalizeResult finalizeWriter(String file) throws IOException { var sortPermLenOut = call.longOut(); long rc = call.invokeIO( FINALIZE_WRITER, - f.segment(), - f.len(), + handle, versionOut, numRowsOut, out.data(), @@ -433,13 +433,16 @@ public static String getColumnMetadata(String file) throws IOException { } } - public static long getFilteredNativeBytesUsed(String pathPrefix) { + public static long getWriterMemoryUsage(long handle) { try (var call = new NativeCall()) { - var p = call.str(pathPrefix); - return call.invoke(GET_FILTERED_BYTES, p.segment(), p.len()); + return call.invoke(GET_WRITER_MEMORY_USAGE, handle); } } + public static void freeWriter(long handle) { + NativeCall.invokeVoid(FREE_WRITER, handle); + } + public static void onSettingsUpdate(NativeSettings nativeSettings) throws IOException { try (var call = new NativeCall()) { var idx = call.str(nativeSettings.getIndexName()); diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/engine/ParquetIndexingEngine.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/engine/ParquetIndexingEngine.java index 3216a6a5bfee5..95b6cb512ad59 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/engine/ParquetIndexingEngine.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/engine/ParquetIndexingEngine.java @@ -47,6 +47,8 @@ import java.util.Collection; import java.util.List; import java.util.Map; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.Supplier; import static org.opensearch.parquet.ParquetDataFormatPlugin.PARQUET_DATA_FORMAT; @@ -89,6 +91,13 @@ public class ParquetIndexingEngine implements IndexingExecutionEngine activeWriters = ConcurrentHashMap.newKeySet(); + /** * Creates a new ParquetIndexingEngine. * @@ -245,7 +254,7 @@ public Writer createWriter(WriterConfig config) { long mappingVersion = mappingVersionSupplier.get(); Schema schema = getOrBuildSchema(); Path filePath = buildParquetFilePath(shardPath, config.writerGeneration(), null); - return new ParquetWriter( + ParquetWriter writer = new ParquetWriter( filePath.toString(), config.writerGeneration(), 0L, @@ -256,8 +265,11 @@ public Writer createWriter(WriterConfig config) { indexSettings, threadPool, checksumStrategy, - statsTracker + statsTracker, + activeWriters::remove ); + activeWriters.add(writer); + return writer; } /** Parquet indexing uses only native (off-heap) memory via Arrow buffers and Rust writers, no JVM heap. */ @@ -268,7 +280,11 @@ public long getHeapBytesUsed() { @Override public long getNativeBytesUsed() { - return bufferPool.getTotalAllocatedBytes() + RustBridge.getFilteredNativeBytesUsed(shardPath.getDataPath().toString()); + long total = bufferPool.getTotalAllocatedBytes(); + for (ParquetWriter writer : activeWriters) { + total += writer.getNativeBytesUsed(); + } + return total; } @Override diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/vsr/VSRManager.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/vsr/VSRManager.java index 5d64224a2bfc2..9b9185d031210 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/vsr/VSRManager.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/vsr/VSRManager.java @@ -370,14 +370,25 @@ public void close() { throw new RuntimeException("Failed to close VSRManager: " + e.getMessage(), e); } finally { try { + if (writer != null) { + writer.close(); + } vsrPool.close(); } catch (Exception e) { - logger.error("Error releasing VSR pool during close for {}: {}", fileName, e.getMessage()); + logger.error("Error releasing Writer/VSR pool during close for {}: {}", fileName, e.getMessage()); } managedVSR.set(null); } } + /** + * Returns the native (Rust-side) memory reserved by this manager's writer, or 0 if the writer + * was never initialized or has been finalized/freed. + */ + public long getNativeBytesUsed() { + return writer == null ? 0L : writer.getNativeBytesUsed(); + } + /** * Initializes the native writer on first use, using the schema from the given VSR. */ diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/writer/ParquetWriter.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/writer/ParquetWriter.java index 24736d59524bb..c7ad13da8ffa4 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/writer/ParquetWriter.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/writer/ParquetWriter.java @@ -32,6 +32,7 @@ import java.io.IOException; import java.nio.file.Path; +import java.util.function.Consumer; import java.util.function.Supplier; /** @@ -58,6 +59,7 @@ public class ParquetWriter implements Writer { private final FormatChecksumStrategy checksumStrategy; private final Supplier schemaSupplier; private final ParquetShardStatsTracker stats; + private final Consumer onClose; private volatile long mappingVersion; private volatile WriterState state = WriterState.ACTIVE; private long acceptedRows = 0L; @@ -90,7 +92,8 @@ public ParquetWriter( IndexSettings indexSettings, ThreadPool threadPool, FormatChecksumStrategy checksumStrategy, - ParquetShardStatsTracker stats + ParquetShardStatsTracker stats, + Consumer onClose ) { this.file = file; this.writerGeneration = writerGeneration; @@ -99,6 +102,7 @@ public ParquetWriter( this.checksumStrategy = checksumStrategy; this.schemaSupplier = schemaSupplier; this.stats = stats; + this.onClose = onClose; this.vsrManager = new VSRManager( file, indexSettings, @@ -137,7 +141,8 @@ public ParquetWriter( indexSettings, threadPool, checksumStrategy, - new ParquetShardStatsTracker() + new ParquetShardStatsTracker(), + null ); } @@ -255,12 +260,23 @@ public void updateMappingVersion(long newVersion) { } } + /** + * Returns the native (Rust-side) memory currently reserved by this writer's native Parquet + * writer, or 0 if uninitialized/finalized. Summed by the engine for shard native-memory stats. + */ + public long getNativeBytesUsed() { + return vsrManager.getNativeBytesUsed(); + } + @Override public void close() throws IOException { try { vsrManager.close(); } finally { state = WriterState.CLOSED; + if (onClose != null) { + onClose.accept(this); + } } } diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/ffm.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/ffm.rs index 36f94820a37ff..db92a137864cc 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/ffm.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/ffm.rs @@ -14,7 +14,7 @@ use std::slice; use std::str; -use native_bridge_common::{ffm_safe, log_debug}; +use native_bridge_common::ffm_safe; use crate::field_config::FieldConfig; use crate::merge; @@ -104,22 +104,18 @@ pub unsafe extern "C" fn parquet_create_writer( nulls_first, writer_generation, ) - .map(|_| 0) + .map(|ptr| ptr as i64) .map_err(|e| e.to_string()) } #[ffm_safe] #[no_mangle] pub unsafe extern "C" fn parquet_write( - file_ptr: *const u8, - file_len: i64, + handle: i64, array_address: i64, schema_address: i64, ) -> i64 { - let filename = str_from_raw(file_ptr, file_len) - .map_err(|e| format!("parquet_write: {}", e))? - .to_string(); - NativeParquetWriter::write_data(filename, array_address, schema_address) + NativeParquetWriter::write_data(handle as *mut _, array_address, schema_address) .map(|_| 0) .map_err(|e| e.to_string()) } @@ -128,8 +124,7 @@ pub unsafe extern "C" fn parquet_write( #[ffm_safe] #[no_mangle] pub unsafe extern "C" fn parquet_finalize_writer( - file_ptr: *const u8, - file_len: i64, + handle: i64, version_out: *mut i32, num_rows_out: *mut i64, created_by_buf: *mut u8, @@ -140,10 +135,7 @@ pub unsafe extern "C" fn parquet_finalize_writer( sort_perm_ptr_out: *mut i64, sort_perm_len_out: *mut i64, ) -> i64 { - let filename = str_from_raw(file_ptr, file_len) - .map_err(|e| format!("parquet_finalize_writer: {}", e))? - .to_string(); - match NativeParquetWriter::finalize_writer(filename) { + match NativeParquetWriter::finalize_writer(handle as *mut _) { Ok(Some(result)) => { let fm = result.metadata.file_metadata(); if !version_out.is_null() { @@ -303,14 +295,15 @@ pub unsafe extern "C" fn parquet_get_column_metadata( } #[no_mangle] -pub unsafe extern "C" fn parquet_get_filtered_native_bytes_used( - prefix_ptr: *const u8, - prefix_len: i64, -) -> i64 { - let prefix = str_from_raw(prefix_ptr, prefix_len) - .unwrap_or("") - .to_string(); - NativeParquetWriter::get_filtered_writer_memory_usage(prefix).unwrap_or(0) as i64 +pub unsafe extern "C" fn parquet_get_writer_memory_usage(handle: i64) -> i64 { + NativeParquetWriter::get_writer_memory_usage(handle as *const _) as i64 +} + +/// Best-effort teardown of an un-finalized writer. Idempotency is enforced Java-side (the writer's +/// released-flag CAS + Cleaner), so this is only ever called once per live handle. Never fails. +#[no_mangle] +pub unsafe extern "C" fn parquet_free_writer(handle: i64) { + NativeParquetWriter::free_writer(handle as *mut _); } // --------------------------------------------------------------------------- diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/test_utils.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/test_utils.rs index 28efa40fbb249..e368f43c582f5 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/test_utils.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/test_utils.rs @@ -19,6 +19,82 @@ use tempfile::tempdir; use crate::writer::NativeParquetWriter; +use lazy_static::lazy_static; +use std::collections::HashMap; +use std::sync::Mutex; + +lazy_static! { + /// Test-only map from output filename to the opaque native writer handle. Lets the tests keep + /// a filename-based API while production owns writers via Java-side handles. + static ref TEST_HANDLES: Mutex> = Mutex::new(HashMap::new()); +} + +/// Test shim: create a writer and remember its handle keyed by filename. +pub fn test_create_writer( + filename: String, + index_name: String, + schema_address: i64, + sort_columns: Vec, + reverse_sorts: Vec, + nulls_first: Vec, + writer_generation: i64, +) -> Result<(), Box> { + let handle = NativeParquetWriter::create_writer( + filename.clone(), + index_name, + schema_address, + sort_columns, + reverse_sorts, + nulls_first, + writer_generation, + )?; + TEST_HANDLES.lock().unwrap().insert(filename, handle as i64); + Ok(()) +} + +/// Test shim: write to the writer previously created for `filename`. +pub fn test_write_data( + filename: String, + array_address: i64, + schema_address: i64, +) -> Result<(), Box> { + let handle = TEST_HANDLES + .lock() + .unwrap() + .get(&filename) + .copied() + .unwrap_or(0); + NativeParquetWriter::write_data( + handle as *mut crate::writer::WriterState, + array_address, + schema_address, + ) +} + +/// Test shim: finalize (and forget) the writer previously created for `filename`. +pub fn test_finalize_writer( + filename: String, +) -> Result, Box> { + let handle = TEST_HANDLES.lock().unwrap().remove(&filename).unwrap_or(0); + NativeParquetWriter::finalize_writer(handle as *mut crate::writer::WriterState) +} + +/// Test shim: native memory reserved by the writer previously created for `filename` (0 if none). +pub fn test_writer_memory_usage(filename: &str) -> usize { + let handle = TEST_HANDLES + .lock() + .unwrap() + .get(filename) + .copied() + .unwrap_or(0); + NativeParquetWriter::get_writer_memory_usage(handle as *const crate::writer::WriterState) +} + +/// Test shim: whether a (non-finalized) writer currently exists for `filename`. +pub fn test_has_writer(filename: &str) -> bool { + TEST_HANDLES.lock().unwrap().contains_key(filename) +} + pub fn create_test_ffi_schema() -> (Arc, i64) { let schema = Arc::new(Schema::new(vec![ Field::new("id", DataType::Int32, false), @@ -75,7 +151,7 @@ pub fn get_temp_file_path(name: &str) -> (tempfile::TempDir, String) { pub fn create_writer_and_assert_success(filename: &str) -> (Arc, i64) { let (schema, schema_ptr) = create_test_ffi_schema(); - let result = NativeParquetWriter::create_writer( + let result = test_create_writer( filename.to_string(), "test-index".to_string(), schema_ptr, @@ -94,7 +170,7 @@ pub fn create_sorted_writer_and_assert_success( reverse: bool, ) -> (Arc, i64) { let (schema, schema_ptr) = create_test_ffi_schema(); - let result = NativeParquetWriter::create_writer( + let result = test_create_writer( filename.to_string(), "test-index".to_string(), schema_ptr, @@ -108,13 +184,13 @@ pub fn create_sorted_writer_and_assert_success( } pub fn close_writer_and_cleanup_schema(filename: &str, schema_ptr: i64) { - let _ = NativeParquetWriter::finalize_writer(filename.to_string()); + let _ = test_finalize_writer(filename.to_string()); cleanup_ffi_schema(schema_ptr); } pub fn write_ffi_data_to_writer(filename: &str) -> (i64, i64) { let (array_ptr, data_schema_ptr) = create_test_ffi_data().unwrap(); - let result = NativeParquetWriter::write_data(filename.to_string(), array_ptr, data_schema_ptr); + let result = test_write_data(filename.to_string(), array_ptr, data_schema_ptr); assert!(result.is_ok()); (array_ptr, data_schema_ptr) } @@ -141,7 +217,7 @@ pub fn close_writer_and_get_metadata( filename: &str, schema_ptr: i64, ) -> crate::writer::FinalizeResult { - let result = NativeParquetWriter::finalize_writer(filename.to_string()); + let result = test_finalize_writer(filename.to_string()); cleanup_ffi_schema(schema_ptr); result.unwrap().unwrap() } diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/tests/mod.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/tests/mod.rs index 356116528ab31..8da143df09dbb 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/tests/mod.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/tests/mod.rs @@ -24,7 +24,7 @@ use std::io::Read; fn test_create_writer_success() { let (_temp_dir, filename) = get_temp_file_path("test.parquet"); let (_schema, schema_ptr) = create_writer_and_assert_success(&filename); - assert!(NativeParquetWriter::has_writer(&filename)); + assert!(test_has_writer(&filename)); close_writer_and_cleanup_schema(&filename, schema_ptr); } @@ -32,7 +32,7 @@ fn test_create_writer_success() { fn test_create_writer_invalid_path() { let invalid_path = "/invalid/path/that/does/not/exist/test.parquet"; let (_schema, schema_ptr) = create_test_ffi_schema(); - let result = NativeParquetWriter::create_writer( + let result = test_create_writer( invalid_path.to_string(), "test-index".to_string(), schema_ptr, @@ -48,7 +48,7 @@ fn test_create_writer_invalid_path() { #[test] fn test_create_writer_invalid_schema_pointer() { let (_temp_dir, filename) = get_temp_file_path("invalid_schema.parquet"); - let result = NativeParquetWriter::create_writer( + let result = test_create_writer( filename, "test-index".to_string(), 0, @@ -69,7 +69,9 @@ fn test_create_writer_multiple_times_same_file() { let (_temp_dir, filename) = get_temp_file_path("duplicate.parquet"); let (_schema, schema_ptr) = create_writer_and_assert_success(&filename); let (_, schema_ptr2) = create_test_ffi_schema(); - let result2 = NativeParquetWriter::create_writer( + // With Java-owned handles there is no filename-keyed native registry, so re-creating a writer + // for the same file is allowed — this is what lets a shard recover after an un-closed writer. + let result2 = test_create_writer( filename.clone(), "test-index".to_string(), schema_ptr2, @@ -78,11 +80,7 @@ fn test_create_writer_multiple_times_same_file() { vec![], 0, ); - assert!(result2.is_err()); - assert!(result2 - .unwrap_err() - .to_string() - .contains("Writer already exists")); + assert!(result2.is_ok()); cleanup_ffi_schema(schema_ptr2); close_writer_and_cleanup_schema(&filename, schema_ptr); } @@ -92,7 +90,7 @@ fn test_write_data_success() { let (_temp_dir, filename) = get_temp_file_path("write_success.parquet"); let (_schema, schema_ptr) = create_writer_and_assert_success(&filename); let (array_ptr, data_schema_ptr) = create_test_ffi_data().unwrap(); - let result = NativeParquetWriter::write_data(filename.clone(), array_ptr, data_schema_ptr); + let result = test_write_data(filename.clone(), array_ptr, data_schema_ptr); assert!(result.is_ok()); cleanup_ffi_data(array_ptr, data_schema_ptr); close_writer_and_cleanup_schema(&filename, schema_ptr); @@ -101,10 +99,10 @@ fn test_write_data_success() { #[test] fn test_write_data_no_writer() { let (array_ptr, schema_ptr) = create_test_ffi_data().unwrap(); - let result = - NativeParquetWriter::write_data("nonexistent.parquet".to_string(), array_ptr, schema_ptr); + // No writer was created for this filename, so the test shim resolves a null (0) handle and the + // native side rejects it. + let result = test_write_data("nonexistent.parquet".to_string(), array_ptr, schema_ptr); assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("Writer not found")); cleanup_ffi_data(array_ptr, schema_ptr); } @@ -123,13 +121,13 @@ fn test_write_data_multiple_batches() { fn test_write_data_invalid_pointers() { let (_temp_dir, filename) = get_temp_file_path("invalid_ffi.parquet"); let (_schema, schema_ptr) = create_writer_and_assert_success(&filename); - let result = NativeParquetWriter::write_data(filename.clone(), 0, 0); + let result = test_write_data(filename.clone(), 0, 0); assert!(result.is_err()); assert!(result .unwrap_err() .to_string() .contains("Invalid FFI addresses")); - let result = NativeParquetWriter::write_data(filename.clone(), 0, schema_ptr); + let result = test_write_data(filename.clone(), 0, schema_ptr); assert!(result.is_err()); assert!(result .unwrap_err() @@ -143,7 +141,7 @@ fn test_write_data_incompatible_schema() { let (_temp_dir, filename) = get_temp_file_path("write_mismatch.parquet"); let (_schema, schema_ptr) = create_writer_and_assert_success(&filename); let (array_ptr, data_schema_ptr) = create_mismatched_ffi_data().unwrap(); - let result = NativeParquetWriter::write_data(filename.clone(), array_ptr, data_schema_ptr); + let result = test_write_data(filename.clone(), array_ptr, data_schema_ptr); assert!(result.is_err()); cleanup_ffi_data(array_ptr, data_schema_ptr); close_writer_and_cleanup_schema(&filename, schema_ptr); @@ -155,7 +153,7 @@ fn test_finalize_writer_success() { let (_schema, schema_ptr) = create_writer_and_assert_success(&filename); let (array_ptr, data_schema_ptr) = write_ffi_data_to_writer(&filename); cleanup_ffi_data(array_ptr, data_schema_ptr); - let result = NativeParquetWriter::finalize_writer(filename.clone()); + let result = test_finalize_writer(filename.clone()); assert!(result.is_ok()); assert!(Path::new(&filename).exists()); cleanup_ffi_schema(schema_ptr); @@ -167,10 +165,10 @@ fn test_finalize_writer_with_data_returns_correct_metadata() { let (_schema, schema_ptr) = create_writer_and_assert_success(&filename); for _ in 0..2 { let (array_ptr, data_schema_ptr) = create_test_ffi_data().unwrap(); - NativeParquetWriter::write_data(filename.clone(), array_ptr, data_schema_ptr).unwrap(); + test_write_data(filename.clone(), array_ptr, data_schema_ptr).unwrap(); cleanup_ffi_data(array_ptr, data_schema_ptr); } - let result = NativeParquetWriter::finalize_writer(filename.clone()); + let result = test_finalize_writer(filename.clone()); assert!(result.is_ok()); let metadata = result.unwrap().unwrap(); assert_eq!(metadata.metadata.file_metadata().num_rows(), 6); @@ -184,9 +182,11 @@ fn test_finalize_writer_with_data_returns_correct_metadata() { #[test] fn test_close_nonexistent_writer() { - let result = NativeParquetWriter::finalize_writer("nonexistent.parquet".to_string()); - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("Writer not found")); + // No writer registered for this file -> the shim resolves a null handle and finalize is a + // no-op that returns None (rather than an error). + let result = test_finalize_writer("nonexistent.parquet".to_string()); + assert!(result.is_ok()); + assert!(result.unwrap().is_none()); } #[test] @@ -195,14 +195,12 @@ fn test_close_multiple_times_same_file() { let (_schema, schema_ptr) = create_writer_and_assert_success(&filename); let (array_ptr, data_schema_ptr) = write_ffi_data_to_writer(&filename); cleanup_ffi_data(array_ptr, data_schema_ptr); - let result1 = NativeParquetWriter::finalize_writer(filename.clone()); + let result1 = test_finalize_writer(filename.clone()); assert!(result1.is_ok()); - let result2 = NativeParquetWriter::finalize_writer(filename); - assert!(result2.is_err()); - assert!(result2 - .unwrap_err() - .to_string() - .contains("Writer not found")); + // Second finalize: the handle was already consumed/forgotten -> null handle -> no-op None. + let result2 = test_finalize_writer(filename); + assert!(result2.is_ok()); + assert!(result2.unwrap().is_none()); cleanup_ffi_schema(schema_ptr); } @@ -217,7 +215,7 @@ fn test_complete_writer_lifecycle() { cleanup_ffi_data(array_ptr, data_schema_ptr); } - let close_result = NativeParquetWriter::finalize_writer(filename.clone()); + let close_result = test_finalize_writer(filename.clone()); assert!(close_result.is_ok()); assert!(close_result.unwrap().is_some()); @@ -235,15 +233,15 @@ fn test_sorted_writer_ascending() { let (ap1, sp1) = create_test_ffi_data_with_ids(vec![30, 10, 50], vec![Some("C"), Some("A"), Some("E")]) .unwrap(); - NativeParquetWriter::write_data(filename.clone(), ap1, sp1).unwrap(); + test_write_data(filename.clone(), ap1, sp1).unwrap(); cleanup_ffi_data(ap1, sp1); let (ap2, sp2) = create_test_ffi_data_with_ids(vec![20, 40], vec![Some("B"), Some("D")]).unwrap(); - NativeParquetWriter::write_data(filename.clone(), ap2, sp2).unwrap(); + test_write_data(filename.clone(), ap2, sp2).unwrap(); cleanup_ffi_data(ap2, sp2); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); let ids = read_parquet_file_sorted_ids(&filename); assert_eq!( @@ -263,15 +261,15 @@ fn test_sorted_writer_descending() { let (ap1, sp1) = create_test_ffi_data_with_ids(vec![30, 10, 50], vec![Some("C"), Some("A"), Some("E")]) .unwrap(); - NativeParquetWriter::write_data(filename.clone(), ap1, sp1).unwrap(); + test_write_data(filename.clone(), ap1, sp1).unwrap(); cleanup_ffi_data(ap1, sp1); let (ap2, sp2) = create_test_ffi_data_with_ids(vec![20, 40], vec![Some("B"), Some("D")]).unwrap(); - NativeParquetWriter::write_data(filename.clone(), ap2, sp2).unwrap(); + test_write_data(filename.clone(), ap2, sp2).unwrap(); cleanup_ffi_data(ap2, sp2); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); let ids = read_parquet_file_sorted_ids(&filename); assert_eq!( @@ -291,15 +289,15 @@ fn test_unsorted_writer_preserves_insertion_order() { let (ap1, sp1) = create_test_ffi_data_with_ids(vec![30, 10, 50], vec![Some("C"), Some("A"), Some("E")]) .unwrap(); - NativeParquetWriter::write_data(filename.clone(), ap1, sp1).unwrap(); + test_write_data(filename.clone(), ap1, sp1).unwrap(); cleanup_ffi_data(ap1, sp1); let (ap2, sp2) = create_test_ffi_data_with_ids(vec![20, 40], vec![Some("B"), Some("D")]).unwrap(); - NativeParquetWriter::write_data(filename.clone(), ap2, sp2).unwrap(); + test_write_data(filename.clone(), ap2, sp2).unwrap(); cleanup_ffi_data(ap2, sp2); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); let ids = read_parquet_file_sorted_ids(&filename); assert_eq!( @@ -320,18 +318,15 @@ fn test_ipc_staging_sorted_writer_creates_and_cleans_up_staging_file() { // With eager sort-and-write, no staging file exists until a chunk is flushed. // The writer accumulates in memory. Verify the writer is open. - assert!( - NativeParquetWriter::has_writer(&filename), - "Writer should be open" - ); + assert!(test_has_writer(&filename), "Writer should be open"); let (ap, sp) = create_test_ffi_data_with_ids(vec![30, 10, 20], vec![Some("C"), Some("A"), Some("B")]) .unwrap(); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); cleanup_ffi_data(ap, sp); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); // The final Parquet file should exist assert!( @@ -352,7 +347,7 @@ fn test_ipc_staging_has_writer_returns_true() { let (_schema, schema_ptr) = create_sorted_writer_and_assert_success(&filename, "id", false); assert!( - NativeParquetWriter::has_writer(&filename), + test_has_writer(&filename), "has_writer should return true for IPC writer" ); @@ -365,7 +360,8 @@ fn test_ipc_staging_duplicate_writer_rejected() { let (_schema, schema_ptr) = create_sorted_writer_and_assert_success(&filename, "id", false); let (_, schema_ptr2) = create_test_ffi_schema(); - let result = NativeParquetWriter::create_writer( + // Handle-owned lifecycle: re-creating a sorted (IPC) writer for the same file is allowed. + let result = test_create_writer( filename.clone(), "test-index".to_string(), schema_ptr2, @@ -374,11 +370,7 @@ fn test_ipc_staging_duplicate_writer_rejected() { vec![false], 0, ); - assert!(result.is_err()); - assert!(result - .unwrap_err() - .to_string() - .contains("Writer already exists")); + assert!(result.is_ok()); cleanup_ffi_schema(schema_ptr2); close_writer_and_cleanup_schema(&filename, schema_ptr); @@ -390,7 +382,7 @@ fn test_ipc_staging_empty_data_produces_valid_parquet() { let (_schema, schema_ptr) = create_sorted_writer_and_assert_success(&filename, "id", false); // Finalize without writing any data - let result = NativeParquetWriter::finalize_writer(filename.clone()); + let result = test_finalize_writer(filename.clone()); assert!(result.is_ok()); assert!( Path::new(&filename).exists(), @@ -411,20 +403,20 @@ fn test_ipc_staging_multi_batch_sort() { // Write multiple batches with interleaved values let (ap1, sp1) = create_test_ffi_data_with_ids(vec![50, 10], vec![Some("E"), Some("A")]).unwrap(); - NativeParquetWriter::write_data(filename.clone(), ap1, sp1).unwrap(); + test_write_data(filename.clone(), ap1, sp1).unwrap(); cleanup_ffi_data(ap1, sp1); let (ap2, sp2) = create_test_ffi_data_with_ids(vec![30, 20], vec![Some("C"), Some("B")]).unwrap(); - NativeParquetWriter::write_data(filename.clone(), ap2, sp2).unwrap(); + test_write_data(filename.clone(), ap2, sp2).unwrap(); cleanup_ffi_data(ap2, sp2); let (ap3, sp3) = create_test_ffi_data_with_ids(vec![40, 60], vec![Some("D"), Some("F")]).unwrap(); - NativeParquetWriter::write_data(filename.clone(), ap3, sp3).unwrap(); + test_write_data(filename.clone(), ap3, sp3).unwrap(); cleanup_ffi_data(ap3, sp3); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); let ids = read_parquet_file_sorted_ids(&filename); assert_eq!( @@ -444,10 +436,10 @@ fn test_ipc_staging_descending_sort() { let (ap, sp) = create_test_ffi_data_with_ids(vec![10, 30, 20], vec![Some("A"), Some("C"), Some("B")]) .unwrap(); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); cleanup_ffi_data(ap, sp); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); let ids = read_parquet_file_sorted_ids(&filename); assert_eq!( @@ -472,18 +464,18 @@ fn test_ipc_and_parquet_writers_coexist() { let (ap1, dp1) = create_test_ffi_data_with_ids(vec![30, 10, 20], vec![Some("C"), Some("A"), Some("B")]) .unwrap(); - NativeParquetWriter::write_data(sorted_file.clone(), ap1, dp1).unwrap(); + test_write_data(sorted_file.clone(), ap1, dp1).unwrap(); cleanup_ffi_data(ap1, dp1); let (ap2, dp2) = create_test_ffi_data_with_ids(vec![30, 10, 20], vec![Some("C"), Some("A"), Some("B")]) .unwrap(); - NativeParquetWriter::write_data(unsorted_file.clone(), ap2, dp2).unwrap(); + test_write_data(unsorted_file.clone(), ap2, dp2).unwrap(); cleanup_ffi_data(ap2, dp2); // Finalize both - NativeParquetWriter::finalize_writer(sorted_file.clone()).unwrap(); - NativeParquetWriter::finalize_writer(unsorted_file.clone()).unwrap(); + test_finalize_writer(sorted_file.clone()).unwrap(); + test_finalize_writer(unsorted_file.clone()).unwrap(); // Sorted file should be sorted let sorted_ids = read_parquet_file_sorted_ids(&sorted_file); @@ -512,7 +504,7 @@ fn test_ipc_staging_concurrent_sorted_writers() { let filename = file_path.to_string_lossy().to_string(); let (_schema, schema_ptr) = create_test_ffi_schema(); - if NativeParquetWriter::create_writer( + if test_create_writer( filename.clone(), "test-index".to_string(), schema_ptr, @@ -528,13 +520,11 @@ fn test_ipc_staging_concurrent_sorted_writers() { vec![Some("C"), Some("A"), Some("B")], ) .unwrap(); - let write_ok = NativeParquetWriter::write_data(filename.clone(), ap, sp).is_ok(); + let write_ok = test_write_data(filename.clone(), ap, sp).is_ok(); cleanup_ffi_data(ap, sp); if write_ok { - if let Ok(Some(metadata)) = - NativeParquetWriter::finalize_writer(filename.clone()) - { + if let Ok(Some(metadata)) = test_finalize_writer(filename.clone()) { if metadata.metadata.file_metadata().num_rows() == 3 { let ids = read_parquet_file_sorted_ids(&filename); if ids == vec![10, 20, 30] { @@ -564,11 +554,11 @@ fn test_ipc_staging_complete_lifecycle_with_sync() { for batch_ids in [vec![50, 30], vec![10, 40], vec![20, 60]] { let names: Vec> = batch_ids.iter().map(|_| Some("x")).collect(); let (ap, sp) = create_test_ffi_data_with_ids(batch_ids, names).unwrap(); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); cleanup_ffi_data(ap, sp); } - let result = NativeParquetWriter::finalize_writer(filename.clone()); + let result = test_finalize_writer(filename.clone()); assert!(result.is_ok()); let metadata = result.unwrap().unwrap(); assert_eq!(metadata.metadata.file_metadata().num_rows(), 6); @@ -589,12 +579,11 @@ fn test_ipc_staging_complete_lifecycle_with_sync() { fn test_get_filtered_writer_memory_usage_with_writers() { let (_temp_dir, filename1) = get_temp_file_path("test1.parquet"); let (_temp_dir2, filename2) = get_temp_file_path("test2.parquet"); - let prefix = _temp_dir.path().to_string_lossy().to_string(); let (_schema1, schema_ptr1) = create_writer_and_assert_success(&filename1); let (_schema2, schema_ptr2) = create_writer_and_assert_success(&filename2); - let result = NativeParquetWriter::get_filtered_writer_memory_usage(prefix); - assert!(result.is_ok()); - assert!(result.unwrap() >= 0); + // Memory is now reported per writer handle (summed on the Java side); assert both are queryable. + let total = test_writer_memory_usage(&filename1) + test_writer_memory_usage(&filename2); + let _ = total; close_writer_and_cleanup_schema(&filename1, schema_ptr1); close_writer_and_cleanup_schema(&filename2, schema_ptr2); } @@ -622,11 +611,11 @@ fn test_crc32_matches_reread_with_data() { for _ in 0..3 { let (array_ptr, data_schema_ptr) = create_test_ffi_data().unwrap(); - NativeParquetWriter::write_data(filename.clone(), array_ptr, data_schema_ptr).unwrap(); + test_write_data(filename.clone(), array_ptr, data_schema_ptr).unwrap(); cleanup_ffi_data(array_ptr, data_schema_ptr); } - let result = NativeParquetWriter::finalize_writer(filename.clone()); + let result = test_finalize_writer(filename.clone()); assert!(result.is_ok()); let finalize_result = result.unwrap().unwrap(); let streaming_crc32 = finalize_result.crc32; @@ -644,16 +633,16 @@ fn test_crc32_matches_reread_with_data() { fn test_crc32_differs_for_different_content() { let (_temp_dir1, filename1) = get_temp_file_path("crc32_diff_a.parquet"); let (_schema1, schema_ptr1) = create_writer_and_assert_success(&filename1); - let result1 = NativeParquetWriter::finalize_writer(filename1.clone()); + let result1 = test_finalize_writer(filename1.clone()); let crc32_empty = result1.unwrap().unwrap().crc32; cleanup_ffi_schema(schema_ptr1); let (_temp_dir2, filename2) = get_temp_file_path("crc32_diff_b.parquet"); let (_schema2, schema_ptr2) = create_writer_and_assert_success(&filename2); let (array_ptr, data_schema_ptr) = create_test_ffi_data().unwrap(); - NativeParquetWriter::write_data(filename2.clone(), array_ptr, data_schema_ptr).unwrap(); + test_write_data(filename2.clone(), array_ptr, data_schema_ptr).unwrap(); cleanup_ffi_data(array_ptr, data_schema_ptr); - let result2 = NativeParquetWriter::finalize_writer(filename2.clone()); + let result2 = test_finalize_writer(filename2.clone()); let crc32_with_data = result2.unwrap().unwrap().crc32; cleanup_ffi_schema(schema_ptr2); @@ -677,7 +666,7 @@ fn test_concurrent_writer_creation() { let filename = file_path.to_string_lossy().to_string(); let (_schema, schema_ptr) = create_test_ffi_schema(); - if NativeParquetWriter::create_writer( + if test_create_writer( filename.clone(), "test-index".to_string(), schema_ptr, @@ -690,9 +679,9 @@ fn test_concurrent_writer_creation() { { success_count.fetch_add(1, Ordering::SeqCst); let (ap, sp) = create_test_ffi_data().unwrap(); - let _ = NativeParquetWriter::write_data(filename.clone(), ap, sp); + let _ = test_write_data(filename.clone(), ap, sp); cleanup_ffi_data(ap, sp); - let _ = NativeParquetWriter::finalize_writer(filename); + let _ = test_finalize_writer(filename); } cleanup_ffi_schema(schema_ptr); }); @@ -722,7 +711,9 @@ fn test_concurrent_close_operations_same_file() { let success_count = Arc::clone(&success_count); let handle = thread::spawn(move || { - if NativeParquetWriter::finalize_writer(filename).is_ok() { + // Exactly one thread wins the handle (the test map's remove is atomic) and performs the + // real finalize (Ok(Some)); the others observe a null handle and no-op (Ok(None)). + if matches!(test_finalize_writer(filename), Ok(Some(_))) { success_count.fetch_add(1, Ordering::SeqCst); } }); @@ -751,7 +742,7 @@ fn test_concurrent_writes_same_file() { let handle = thread::spawn(move || { let (array_ptr, data_schema_ptr) = create_test_ffi_data().unwrap(); - if NativeParquetWriter::write_data(filename, array_ptr, data_schema_ptr).is_ok() { + if test_write_data(filename, array_ptr, data_schema_ptr).is_ok() { success_count.fetch_add(1, Ordering::SeqCst); } cleanup_ffi_data(array_ptr, data_schema_ptr); @@ -793,9 +784,7 @@ fn test_concurrent_writes_different_files() { let handle = thread::spawn(move || { for _ in 0..2 { let (array_ptr, data_schema_ptr) = create_test_ffi_data().unwrap(); - if NativeParquetWriter::write_data(filename.clone(), array_ptr, data_schema_ptr) - .is_ok() - { + if test_write_data(filename.clone(), array_ptr, data_schema_ptr).is_ok() { success_count.fetch_add(1, Ordering::SeqCst); } cleanup_ffi_data(array_ptr, data_schema_ptr); @@ -832,7 +821,7 @@ fn test_bloom_filter_false_propagates_through_settings_store() { let (_temp_dir, filename) = get_temp_file_path("bloom_test.parquet"); let (_schema, schema_ptr) = create_test_ffi_schema(); - let result = NativeParquetWriter::create_writer( + let result = test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -863,7 +852,7 @@ fn test_bloom_filter_default_when_no_settings() { let (_temp_dir, filename) = get_temp_file_path("bloom_default.parquet"); let (_schema, schema_ptr) = create_test_ffi_schema(); - let result = NativeParquetWriter::create_writer( + let result = test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -1054,7 +1043,7 @@ fn test_chunked_writer_single_chunk_row_ids_sequential() { // Create writer with sort on "age" ascending let (_schema, schema_ptr) = create_row_id_schema_ptr(); - let result = NativeParquetWriter::create_writer( + let result = test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -1070,11 +1059,9 @@ fn test_chunked_writer_single_chunk_row_ids_sequential() { vec![Some("C"), Some("A"), Some("E"), Some("B"), Some("D")], vec![0, 1, 2, 3, 4], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); // Verify output is sorted by age ascending let ages = read_ages_from_parquet(&filename); @@ -1126,7 +1113,7 @@ fn test_chunked_writer_multi_chunk_row_ids_sequential() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_row_id_schema_ptr(); - let result = NativeParquetWriter::create_writer( + let result = test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -1153,11 +1140,9 @@ fn test_chunked_writer_multi_chunk_row_ids_sequential() { ], vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); // Verify output is sorted by age ascending let ages = read_ages_from_parquet(&filename); @@ -1213,7 +1198,7 @@ fn test_chunked_writer_multi_chunk_descending_sort() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -1229,11 +1214,9 @@ fn test_chunked_writer_multi_chunk_descending_sort() { vec![Some("A"), Some("E"), Some("C"), Some("B"), Some("D")], vec![0, 1, 2, 3, 4], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); // Verify output is sorted by age descending let ages = read_ages_from_parquet(&filename); @@ -1279,7 +1262,7 @@ fn test_chunked_writer_multiple_write_calls() { // First batch: row_ids 0..4 let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -1295,7 +1278,7 @@ fn test_chunked_writer_multiple_write_calls() { vec![Some("E"), Some("C"), Some("A"), Some("D"), Some("B")], vec![0, 1, 2, 3, 4], ); - NativeParquetWriter::write_data(filename.clone(), ap1, sp1).unwrap(); + test_write_data(filename.clone(), ap1, sp1).unwrap(); // Second batch: row_ids 5..9 let (ap2, sp2) = create_ffi_data_with_row_id( @@ -1303,11 +1286,9 @@ fn test_chunked_writer_multiple_write_calls() { vec![Some("F"), Some("G"), Some("H"), Some("I"), Some("J")], vec![5, 6, 7, 8, 9], ); - NativeParquetWriter::write_data(filename.clone(), ap2, sp2).unwrap(); + test_write_data(filename.clone(), ap2, sp2).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); // Verify sorted output let ages = read_ages_from_parquet(&filename); @@ -1360,7 +1341,7 @@ fn test_chunked_writer_empty_finalize() { // Just create the writer with the schema, don't write data let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -1371,9 +1352,7 @@ fn test_chunked_writer_empty_finalize() { ) .unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); assert_eq!( finalize_result.metadata.file_metadata().num_rows(), @@ -1407,7 +1386,7 @@ fn test_chunked_writer_permutation_is_invertible() { let names: Vec> = (0..10).map(|_| Some("x")).collect(); let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -1419,11 +1398,9 @@ fn test_chunked_writer_permutation_is_invertible() { .unwrap(); let (ap, sp) = create_ffi_data_with_row_id(original_ages.clone(), names, row_ids); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); let mapping = finalize_result.row_id_mapping.expect("Should have mapping"); assert_eq!(mapping.len(), 10); @@ -1474,7 +1451,7 @@ fn test_chunked_writer_large_dataset_multi_chunk() { let names: Vec> = (0..num_rows as usize).map(|_| Some("x")).collect(); let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -1486,11 +1463,9 @@ fn test_chunked_writer_large_dataset_multi_chunk() { .unwrap(); let (ap, sp) = create_ffi_data_with_row_id(ages.clone(), names, row_ids); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); // Verify output is sorted let sorted_ages = read_ages_from_parquet(&filename); @@ -1568,7 +1543,7 @@ fn test_chunked_writer_generation_in_metadata_single_chunk() { let writer_generation = 42i64; let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -1584,8 +1559,8 @@ fn test_chunked_writer_generation_in_metadata_single_chunk() { vec![Some("C"), Some("A"), Some("B")], vec![0, 1, 2], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); let gen = read_writer_generation_from_parquet(&filename); assert_eq!( @@ -1615,7 +1590,7 @@ fn test_chunked_writer_generation_in_metadata_multi_chunk() { let writer_generation = 7i64; let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -1642,8 +1617,8 @@ fn test_chunked_writer_generation_in_metadata_multi_chunk() { ], vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); // Previously writer_generation was lost in the k-way merge path. This has been fixed — // merge_sorted now propagates writer_generation into the output file metadata. @@ -1679,7 +1654,7 @@ fn test_chunked_writer_generation_zero() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -1692,8 +1667,8 @@ fn test_chunked_writer_generation_zero() { let (ap, sp) = create_ffi_data_with_row_id(vec![20, 10], vec![Some("B"), Some("A")], vec![0, 1]); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); let gen = read_writer_generation_from_parquet(&filename); assert_eq!( @@ -1720,7 +1695,7 @@ fn test_chunked_writer_generation_large_value() { let writer_generation = 999_999i64; let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -1746,11 +1721,9 @@ fn test_chunked_writer_generation_large_value() { ], vec![0, 1, 2, 3, 4, 5, 6, 7, 8], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); // Verify writer_generation in metadata let gen = read_writer_generation_from_parquet(&filename); @@ -1794,7 +1767,7 @@ fn test_unsorted_writer_generation_in_metadata() { let writer_generation = 17i64; let (_, schema_ptr) = create_row_id_schema_ptr(); // No sort columns — uses direct Parquet writer path - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -1810,8 +1783,8 @@ fn test_unsorted_writer_generation_in_metadata() { vec![Some("C"), Some("A"), Some("B")], vec![0, 1, 2], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); let gen = read_writer_generation_from_parquet(&filename); assert_eq!( @@ -1843,7 +1816,7 @@ fn test_chunked_writer_crc32_single_chunk() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -1859,11 +1832,9 @@ fn test_chunked_writer_crc32_single_chunk() { vec![Some("C"), Some("A"), Some("E"), Some("B"), Some("D")], vec![0, 1, 2, 3, 4], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); // CRC should be non-zero for a file with data assert_ne!( @@ -1895,7 +1866,7 @@ fn test_chunked_writer_crc32_multi_chunk() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -1922,11 +1893,9 @@ fn test_chunked_writer_crc32_multi_chunk() { ], vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); // Multi-chunk merge also produces a CRC (from merge_sorted's CrcWriter) // Verify it matches the actual file @@ -1957,7 +1926,7 @@ fn test_chunked_writer_crc32_differs_for_different_data() { // Writer 1 let (_schema, schema_ptr1) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename1.clone(), index_name1.to_string(), schema_ptr1, @@ -1972,14 +1941,12 @@ fn test_chunked_writer_crc32_differs_for_different_data() { vec![Some("A"), Some("B"), Some("C")], vec![0, 1, 2], ); - NativeParquetWriter::write_data(filename1.clone(), ap1, sp1).unwrap(); - let result1 = NativeParquetWriter::finalize_writer(filename1.clone()) - .unwrap() - .unwrap(); + test_write_data(filename1.clone(), ap1, sp1).unwrap(); + let result1 = test_finalize_writer(filename1.clone()).unwrap().unwrap(); // Writer 2 — different data let (_schema, schema_ptr2) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename2.clone(), index_name2.to_string(), schema_ptr2, @@ -1994,10 +1961,8 @@ fn test_chunked_writer_crc32_differs_for_different_data() { vec![Some("X"), Some("Y"), Some("Z")], vec![0, 1, 2], ); - NativeParquetWriter::write_data(filename2.clone(), ap2, sp2).unwrap(); - let result2 = NativeParquetWriter::finalize_writer(filename2.clone()) - .unwrap() - .unwrap(); + test_write_data(filename2.clone(), ap2, sp2).unwrap(); + let result2 = test_finalize_writer(filename2.clone()).unwrap().unwrap(); assert_ne!( result1.crc32, result2.crc32, @@ -2031,7 +1996,7 @@ fn test_chunked_writer_batch_slicing_large_batch() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -2057,11 +2022,9 @@ fn test_chunked_writer_batch_slicing_large_batch() { ], vec![0, 1, 2, 3, 4, 5, 6, 7], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); // Verify output is sorted let ages = read_ages_from_parquet(&filename); @@ -2116,7 +2079,7 @@ fn test_chunked_writer_batch_slicing_two_rows_per_slice() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -2144,11 +2107,9 @@ fn test_chunked_writer_batch_slicing_two_rows_per_slice() { ], vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); // Verify sorted output let ages = read_ages_from_parquet(&filename); @@ -2201,7 +2162,7 @@ fn test_chunked_writer_batch_slicing_descending() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -2217,11 +2178,9 @@ fn test_chunked_writer_batch_slicing_descending() { vec![Some("A"), Some("E"), Some("C"), Some("B"), Some("D")], vec![0, 1, 2, 3, 4], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); // Verify descending sort let ages = read_ages_from_parquet(&filename); @@ -2260,7 +2219,7 @@ fn test_chunked_writer_batch_slicing_multiple_writes() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -2277,7 +2236,7 @@ fn test_chunked_writer_batch_slicing_multiple_writes() { vec![Some("E"), Some("A"), Some("C")], vec![0, 1, 2], ); - NativeParquetWriter::write_data(filename.clone(), ap1, sp1).unwrap(); + test_write_data(filename.clone(), ap1, sp1).unwrap(); // Second batch: rows 3..6 let (ap2, sp2) = create_ffi_data_with_row_id( @@ -2285,11 +2244,9 @@ fn test_chunked_writer_batch_slicing_multiple_writes() { vec![Some("D"), Some("B"), Some("F")], vec![3, 4, 5], ); - NativeParquetWriter::write_data(filename.clone(), ap2, sp2).unwrap(); + test_write_data(filename.clone(), ap2, sp2).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); // Verify globally sorted let ages = read_ages_from_parquet(&filename); @@ -2369,7 +2326,7 @@ fn test_chunked_writer_no_row_id_single_chunk() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_no_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -2384,11 +2341,9 @@ fn test_chunked_writer_no_row_id_single_chunk() { vec![50, 10, 30, 20, 40], vec![Some("E"), Some("A"), Some("C"), Some("B"), Some("D")], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); // Verify sorted output let ages = read_ages_from_parquet(&filename); @@ -2422,7 +2377,7 @@ fn test_chunked_writer_no_row_id_multi_chunk() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_no_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -2446,11 +2401,9 @@ fn test_chunked_writer_no_row_id_multi_chunk() { Some("E"), ], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); // Verify sorted output let ages = read_ages_from_parquet(&filename); @@ -2484,7 +2437,7 @@ fn test_chunked_writer_no_row_id_batch_slicing() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_no_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -2506,11 +2459,9 @@ fn test_chunked_writer_no_row_id_batch_slicing() { Some("F"), ], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); // Verify descending sort let ages = read_ages_from_parquet(&filename); @@ -2620,7 +2571,7 @@ fn test_multi_column_sort_age_asc_score_desc() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_multi_sort_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -2645,11 +2596,9 @@ fn test_multi_column_sort_age_asc_score_desc() { ], vec![0, 1, 2, 3, 4, 5], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); let ages = read_ages_from_parquet(&filename); let scores = read_scores_from_parquet(&filename); @@ -2690,7 +2639,7 @@ fn test_multi_column_sort_age_desc_score_asc() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_multi_sort_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -2707,9 +2656,9 @@ fn test_multi_column_sort_age_desc_score_asc() { vec![Some("A"), Some("B"), Some("C"), Some("D"), Some("E")], vec![0, 1, 2, 3, 4], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); let ages = read_ages_from_parquet(&filename); let scores = read_scores_from_parquet(&filename); @@ -2739,7 +2688,7 @@ fn test_multi_column_sort_multi_chunk() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_multi_sort_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -2768,11 +2717,9 @@ fn test_multi_column_sort_multi_chunk() { ], vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); let ages = read_ages_from_parquet(&filename); let scores = read_scores_from_parquet(&filename); @@ -2816,7 +2763,7 @@ fn test_multi_column_sort_batch_slicing() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_multi_sort_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -2833,9 +2780,9 @@ fn test_multi_column_sort_batch_slicing() { vec![Some("A"), Some("B"), Some("C"), Some("D"), Some("E")], vec![0, 1, 2, 3, 4], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); let ages = read_ages_from_parquet(&filename); let scores = read_scores_from_parquet(&filename); @@ -2922,7 +2869,7 @@ fn test_nulls_first_true_ascending() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_nullable_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -2937,8 +2884,8 @@ fn test_nulls_first_true_ascending() { vec![Some(30), None, Some(10), None, Some(20)], vec![Some("C"), Some("X"), Some("A"), Some("Y"), Some("B")], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); let ages = read_nullable_ages_from_parquet(&filename); // NULLs first, then ascending non-nulls @@ -2961,7 +2908,7 @@ fn test_nulls_first_false_ascending() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_nullable_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -2976,8 +2923,8 @@ fn test_nulls_first_false_ascending() { vec![Some(30), None, Some(10), None, Some(20)], vec![Some("C"), Some("X"), Some("A"), Some("Y"), Some("B")], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); let ages = read_nullable_ages_from_parquet(&filename); // Non-nulls ascending first, then NULLs last @@ -3044,7 +2991,7 @@ fn test_nulls_first_with_row_id_and_permutation() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_nullable_with_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -3061,11 +3008,9 @@ fn test_nulls_first_with_row_id_and_permutation() { vec![Some("C"), Some("X"), Some("A"), Some("Y"), Some("B")], vec![0, 1, 2, 3, 4], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); // Verify sort: NULLs first, then ascending let ages = read_nullable_ages_from_parquet(&filename); @@ -3112,7 +3057,7 @@ fn test_nulls_last_with_row_id_and_permutation() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_nullable_with_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -3129,11 +3074,9 @@ fn test_nulls_last_with_row_id_and_permutation() { vec![Some("C"), Some("X"), Some("A"), Some("Y"), Some("B")], vec![0, 1, 2, 3, 4], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); // Verify sort: ascending non-nulls, then NULLs last let ages = read_nullable_ages_from_parquet(&filename); @@ -3178,7 +3121,7 @@ fn test_empty_batch_write_does_not_corrupt_sorted_writer() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -3191,7 +3134,7 @@ fn test_empty_batch_write_does_not_corrupt_sorted_writer() { // Write an empty batch (0 rows) let (ap_empty, sp_empty) = create_ffi_data_with_row_id(vec![], vec![], vec![]); - NativeParquetWriter::write_data(filename.clone(), ap_empty, sp_empty).unwrap(); + test_write_data(filename.clone(), ap_empty, sp_empty).unwrap(); // Write a real batch after the empty one let (ap, sp) = create_ffi_data_with_row_id( @@ -3199,11 +3142,9 @@ fn test_empty_batch_write_does_not_corrupt_sorted_writer() { vec![Some("C"), Some("A"), Some("B")], vec![0, 1, 2], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); // Verify sorted output — empty batch should have no effect let ages = read_ages_from_parquet(&filename); @@ -3232,7 +3173,7 @@ fn test_only_empty_batches_produces_empty_output() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -3246,12 +3187,10 @@ fn test_only_empty_batches_produces_empty_output() { // Write multiple empty batches for _ in 0..3 { let (ap, sp) = create_ffi_data_with_row_id(vec![], vec![], vec![]); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); } - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); assert_eq!(finalize_result.metadata.file_metadata().num_rows(), 0); assert!(finalize_result.row_id_mapping.is_none()); @@ -3275,7 +3214,7 @@ fn test_memory_usage_ipc_writer_reports_chunk_row_ids() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -3287,13 +3226,7 @@ fn test_memory_usage_ipc_writer_reports_chunk_row_ids() { .unwrap(); // Before writing, memory should be 0 (no chunks flushed yet) - let path_prefix = Path::new(&filename) - .parent() - .unwrap() - .to_string_lossy() - .to_string(); - let mem_before = - NativeParquetWriter::get_filtered_writer_memory_usage(path_prefix.clone()).unwrap(); + let mem_before = test_writer_memory_usage(&filename); assert_eq!(mem_before, 0, "Memory should be 0 before any chunk flush"); // Write enough data to trigger at least one chunk flush @@ -3313,11 +3246,10 @@ fn test_memory_usage_ipc_writer_reports_chunk_row_ids() { ], vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); // After writing (chunks flushed), memory should be > 0 due to chunk_row_ids - let mem_after = - NativeParquetWriter::get_filtered_writer_memory_usage(path_prefix.clone()).unwrap(); + let mem_after = test_writer_memory_usage(&filename); assert!( mem_after > 0, "Memory should be > 0 after chunk flushes (chunk_row_ids accumulated), got {}", @@ -3333,9 +3265,9 @@ fn test_memory_usage_ipc_writer_reports_chunk_row_ids() { ); // Finalize and verify writer is removed - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); - let mem_final = NativeParquetWriter::get_filtered_writer_memory_usage(path_prefix).unwrap(); + let mem_final = test_writer_memory_usage(&filename); assert_eq!(mem_final, 0, "Memory should be 0 after writer is finalized"); SETTINGS_STORE.remove(index_name); @@ -3358,7 +3290,7 @@ fn test_memory_usage_path_prefix_filtering() { // Create writer 1 let (_schema, schema_ptr1) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename1.clone(), index_name.to_string(), schema_ptr1, @@ -3371,7 +3303,7 @@ fn test_memory_usage_path_prefix_filtering() { // Create writer 2 let (_schema, schema_ptr2) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename2.clone(), index_name.to_string(), schema_ptr2, @@ -3388,30 +3320,19 @@ fn test_memory_usage_path_prefix_filtering() { vec![Some("E"), Some("C"), Some("A"), Some("D"), Some("B")], vec![0, 1, 2, 3, 4], ); - NativeParquetWriter::write_data(filename1.clone(), ap, sp).unwrap(); - - // Query with prefix that matches only writer 1's directory - let prefix1 = Path::new(&filename1) - .parent() - .unwrap() - .to_string_lossy() - .to_string(); - let prefix2 = Path::new(&filename2) - .parent() - .unwrap() - .to_string_lossy() - .to_string(); + test_write_data(filename1.clone(), ap, sp).unwrap(); - let mem1 = NativeParquetWriter::get_filtered_writer_memory_usage(prefix1).unwrap(); - let mem2 = NativeParquetWriter::get_filtered_writer_memory_usage(prefix2).unwrap(); + // Query per-writer memory by handle (the Java side sums these across a shard's writers). + let mem1 = test_writer_memory_usage(&filename1); + let mem2 = test_writer_memory_usage(&filename2); // Writer 1 had data written (and chunks flushed), writer 2 did not - assert!(mem1 >= 0, "Writer 1 memory should be reported"); + let _ = mem1; // writer 1 memory is queryable per handle assert_eq!(mem2, 0, "Writer 2 should have 0 memory (no data written)"); // Cleanup - NativeParquetWriter::finalize_writer(filename1).unwrap(); - NativeParquetWriter::finalize_writer(filename2).unwrap(); + test_finalize_writer(filename1).unwrap(); + test_finalize_writer(filename2).unwrap(); SETTINGS_STORE.remove(index_name); } @@ -3429,12 +3350,11 @@ fn test_write_data_no_ipc_writer() { vec![Some("A"), Some("B"), Some("C")], vec![0, 1, 2], ); - let result = NativeParquetWriter::write_data(filename.clone(), ap, sp); + let result = test_write_data(filename.clone(), ap, sp); - assert!(result.is_err()); assert!( - result.unwrap_err().to_string().contains("Writer not found"), - "Should get 'Writer not found' error for non-existent IPC writer" + result.is_err(), + "Should get an error writing to a non-existent (null-handle) IPC writer" ); } @@ -3460,7 +3380,7 @@ fn test_sort_all_identical_keys_produces_valid_permutation() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -3477,11 +3397,9 @@ fn test_sort_all_identical_keys_produces_valid_permutation() { vec![Some("E"), Some("D"), Some("C"), Some("B"), Some("A")], vec![0, 1, 2, 3, 4], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); - let finalize_result = NativeParquetWriter::finalize_writer(filename.clone()) - .unwrap() - .unwrap(); + let finalize_result = test_finalize_writer(filename.clone()).unwrap().unwrap(); // All ages should be 30 let ages = read_ages_from_parquet(&filename); @@ -3579,7 +3497,7 @@ fn test_writer_properties_honored_empty_path() { let writer_generation = 99i64; let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -3591,7 +3509,7 @@ fn test_writer_properties_honored_empty_path() { .unwrap(); // Don't write any data — triggers the empty path - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); // Verify format version is stamped let format_version = read_format_version_from_parquet(&filename); @@ -3633,7 +3551,7 @@ fn test_writer_properties_honored_single_chunk_snappy_no_bloom() { let writer_generation = 55i64; let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -3649,8 +3567,8 @@ fn test_writer_properties_honored_single_chunk_snappy_no_bloom() { vec![Some("C"), Some("A"), Some("E"), Some("B"), Some("D")], vec![0, 1, 2, 3, 4], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); // Verify compression is SNAPPY let compression = read_compression_from_parquet(&filename); @@ -3710,7 +3628,7 @@ fn test_writer_properties_honored_single_chunk_zstd_with_bloom() { let writer_generation = 12i64; let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -3726,8 +3644,8 @@ fn test_writer_properties_honored_single_chunk_zstd_with_bloom() { vec![Some("C"), Some("A"), Some("E"), Some("B"), Some("D")], vec![0, 1, 2, 3, 4], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); // Verify compression is ZSTD let compression = read_compression_from_parquet(&filename); @@ -3772,7 +3690,7 @@ fn test_writer_properties_honored_multi_chunk_snappy_no_bloom() { let writer_generation = 77i64; let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -3799,8 +3717,8 @@ fn test_writer_properties_honored_multi_chunk_snappy_no_bloom() { ], vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); // Verify compression is SNAPPY in the merged output let compression = read_compression_from_parquet(&filename); @@ -3852,7 +3770,7 @@ fn test_writer_properties_honored_multi_chunk_zstd_with_bloom() { let writer_generation = 33i64; let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -3879,8 +3797,8 @@ fn test_writer_properties_honored_multi_chunk_zstd_with_bloom() { ], vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); // Verify compression is ZSTD in the merged output let compression = read_compression_from_parquet(&filename); @@ -3924,7 +3842,7 @@ fn test_writer_properties_honored_single_chunk_uncompressed() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -3940,8 +3858,8 @@ fn test_writer_properties_honored_single_chunk_uncompressed() { vec![Some("C"), Some("A"), Some("E"), Some("B"), Some("D")], vec![0, 1, 2, 3, 4], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); // Verify compression is UNCOMPRESSED let compression = read_compression_from_parquet(&filename); @@ -3976,7 +3894,7 @@ fn test_writer_properties_honored_multi_chunk_uncompressed() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -4003,8 +3921,8 @@ fn test_writer_properties_honored_multi_chunk_uncompressed() { ], vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); // Verify compression is UNCOMPRESSED let compression = read_compression_from_parquet(&filename); @@ -4043,7 +3961,7 @@ fn test_writer_properties_defaults_single_chunk() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -4059,8 +3977,8 @@ fn test_writer_properties_defaults_single_chunk() { vec![Some("C"), Some("A"), Some("E"), Some("B"), Some("D")], vec![0, 1, 2, 3, 4], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); // Default compression is LZ4_RAW let compression = read_compression_from_parquet(&filename); @@ -4097,7 +4015,7 @@ fn test_writer_properties_defaults_multi_chunk() { SETTINGS_STORE.insert(index_name.to_string(), settings); let (_schema, schema_ptr) = create_row_id_schema_ptr(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), index_name.to_string(), schema_ptr, @@ -4124,8 +4042,8 @@ fn test_writer_properties_defaults_multi_chunk() { ], vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9], ); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); - NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); + test_finalize_writer(filename.clone()).unwrap(); // Default compression is LZ4_RAW let compression = read_compression_from_parquet(&filename); diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/writer.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/writer.rs index 0167362c30c4f..32e5859b83367 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/writer.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/writer.rs @@ -362,34 +362,32 @@ impl SortingChunkedWriter { } } -/// Bundles all per-writer resources so a single `DashMap::remove` atomically -/// drops the writer, closes the file handle, and cleans up sort config. -struct WriterState { +/// Bundles all per-writer resources. The writer's lifetime is owned by Java via an opaque +/// handle (`Box::into_raw`); dropping the box (on finalize/free) closes the file handle, +/// releases the memory reservation, and cleans up sort config. Stores its own file paths +/// (previously the registry key) so the handle-based FFI entry points don't need them passed in. +pub struct WriterState { variant: WriterVariant, settings: NativeSettings, crc_handle: Option, writer_generation: i64, reservation: MemoryReservation, + /// Final output path (temp file is renamed to this on finalize). + filename: String, + /// Temporary path written to before finalize (`temp-`); also the IPC staging base. + temp_filename: String, } /// Path suffix for the intermediate Arrow IPC file used during sort-on-close. const IPC_STAGING_SUFFIX: &str = ".arrow_ipc_staging"; lazy_static! { - /// Unified per-writer registry. Keyed by temp filename. - /// Holds both Parquet and IPC writers via the `WriterVariant` enum. - static ref WRITERS: DashMap = DashMap::new(); pub static ref SETTINGS_STORE: DashMap = DashMap::new(); } pub struct NativeParquetWriter; impl NativeParquetWriter { - /// Returns true if a writer is currently open for the given filename. - pub fn has_writer(filename: &str) -> bool { - let temp_filename = Self::temp_filename(filename); - WRITERS.contains_key(&temp_filename) - } /// Build the temp filename by prepending "temp-" to the basename. fn temp_filename(filename: &str) -> String { let path = Path::new(filename); @@ -411,7 +409,7 @@ impl NativeParquetWriter { reverse_sorts: Vec, nulls_first: Vec, writer_generation: i64, - ) -> Result<(), Box> { + ) -> Result<*mut WriterState, Box> { log_debug!( "create_writer called for file: {}, index: {}, schema_address: {}, sort_columns: {:?}, reverse_sorts: {:?}, nulls_first: {:?}, writer_generation: {}", filename, index_name, schema_address, sort_columns, reverse_sorts, nulls_first, writer_generation @@ -426,12 +424,6 @@ impl NativeParquetWriter { } let temp_filename = Self::temp_filename(&filename); - - if WRITERS.contains_key(&temp_filename) { - log_error!("ERROR: Writer already exists for file: {}", temp_filename); - return Err("Writer already exists for this file".into()); - } - let arrow_schema = unsafe { FFI_ArrowSchema::from_raw(schema_address as *mut _) }; let schema = Arc::new(arrow::datatypes::Schema::try_from(&arrow_schema)?); log_debug!("Schema created with {} fields", schema.fields().len()); @@ -483,38 +475,43 @@ impl NativeParquetWriter { ) }; - WRITERS.insert( + let state = Box::new(WriterState { + variant, + settings, + crc_handle, + writer_generation, + reservation: MemoryReservation::new( + write_pool(), + "parquet_writer", + PoolBehavior::IgnoreLimit, + ), + filename, temp_filename, - WriterState { - variant, - settings, - crc_handle, - writer_generation, - reservation: MemoryReservation::new( - write_pool(), - "parquet_writer", - PoolBehavior::IgnoreLimit, - ), - }, - ); - - Ok(()) + }); + let handle = Box::into_raw(state); + Ok(handle) } pub fn write_data( - filename: String, + handle: *mut WriterState, array_address: i64, schema_address: i64, ) -> Result<(), Box> { - let temp_filename = Self::temp_filename(&filename); - log_debug!( - "write_data called for file: {} (temp: {})", - filename, - temp_filename - ); + if handle.is_null() { + return Err("Invalid writer handle (null)".into()); + } + // Safety: `handle` is a live pointer minted by `create_writer` and owned by Java, which + // serializes all handle-touching native calls (write/finalize/free/memory) under a + // per-writer lock, so there is no concurrent `&mut`/`&` aliasing on this allocation. + // Reborrow only (not Box::from_raw): does NOT take ownership, so the writer is not dropped here and lives on for subsequent writes. + let state = unsafe { &mut *handle }; + log_debug!("write_data called for temp file: {}", state.temp_filename); if (array_address as *mut u8).is_null() || (schema_address as *mut u8).is_null() { - log_error!("ERROR: Invalid FFI addresses for file: {}", temp_filename); + log_error!( + "ERROR: Invalid FFI addresses for file: {}", + state.temp_filename + ); return Err("Invalid FFI addresses (null pointers)".into()); } @@ -533,36 +530,31 @@ impl NativeParquetWriter { record_batch.num_columns() ); - if let Some(mut state) = WRITERS.get_mut(&temp_filename) { - match &state.variant { - WriterVariant::Ipc(writer_arc) => { - log_debug!("Writing RecordBatch to IPC staging file"); - let writer_arc = writer_arc.clone(); - let mut writer = writer_arc.lock().unwrap(); - writer.write(&record_batch, &mut state.reservation)?; - } - WriterVariant::Parquet(writer_arc) => { - log_debug!("Writing RecordBatch to Parquet file"); - let batch_bytes = record_batch.get_array_memory_size(); - // Reserve 3× batch as estimate — ArrowWriter encoding may temporarily - // hold dictionary, compressed pages, and data page buffers. - let estimated = batch_bytes * 3; - let writer_arc = writer_arc.clone(); - state.reservation.reserve_estimated(estimated)?; - let mut writer = writer_arc.lock().unwrap(); - let before = writer.memory_size(); - writer.write(&record_batch)?; - // Reconcile: adjust reservation to actual delta reported by ArrowWriter - let actual = writer.memory_size().saturating_sub(before); - drop(writer); - state.reservation.reconcile(estimated, actual); - } + match &state.variant { + WriterVariant::Ipc(writer_arc) => { + log_debug!("Writing RecordBatch to IPC staging file"); + let writer_arc = writer_arc.clone(); + let mut writer = writer_arc.lock().unwrap(); + writer.write(&record_batch, &mut state.reservation)?; + } + WriterVariant::Parquet(writer_arc) => { + log_debug!("Writing RecordBatch to Parquet file"); + let batch_bytes = record_batch.get_array_memory_size(); + // Reserve 3× batch as estimate — ArrowWriter encoding may temporarily + // hold dictionary, compressed pages, and data page buffers. + let estimated = batch_bytes * 3; + let writer_arc = writer_arc.clone(); + state.reservation.reserve_estimated(estimated)?; + let mut writer = writer_arc.lock().unwrap(); + let before = writer.memory_size(); + writer.write(&record_batch)?; + // Reconcile: adjust reservation to actual delta reported by ArrowWriter + let actual = writer.memory_size().saturating_sub(before); + drop(writer); + state.reservation.reconcile(estimated, actual); } - Ok(()) - } else { - log_error!("ERROR: No writer found for temp file: {}", temp_filename); - Err("Writer not found".into()) } + Ok(()) } else { log_error!( "ERROR: Array is not a StructArray, type: {:?}", @@ -574,135 +566,133 @@ impl NativeParquetWriter { } pub fn finalize_writer( - filename: String, + handle: *mut WriterState, ) -> Result, Box> { - let temp_filename = Self::temp_filename(&filename); + if handle.is_null() { + return Ok(None); + } + // Safety: Java hands back a live handle exactly once for finalize; reclaim ownership so the + // Box (and its reservation + open file handles) is dropped when this function returns. + // Box::from_raw TAKES ownership: the writer is dropped exactly once here (success path). + let WriterState { + variant, + settings, + crc_handle, + writer_generation, + mut reservation, + filename, + temp_filename, + } = unsafe { *Box::from_raw(handle) }; log_debug!( "finalize_writer called for file: {} (temp: {})", filename, temp_filename ); - - if let Some((_, state)) = WRITERS.remove(&temp_filename) { - let WriterState { - variant, - settings, - crc_handle, - writer_generation, - mut reservation, - } = state; - let index_name = settings.index_name.as_deref().unwrap_or(""); - - match variant { - WriterVariant::Ipc(writer_arc) => { - match Arc::try_unwrap(writer_arc) { - Ok(mutex) => { - let chunked_writer = mutex.into_inner().unwrap(); - let total_rows = chunked_writer.total_rows(); - let schema = chunked_writer.schema.clone(); - let (chunk_paths, chunk_row_ids, chunk_crcs) = - chunked_writer.finish(&mut reservation)?; - log_info!( + let index_name = settings.index_name.as_deref().unwrap_or(""); + + match variant { + WriterVariant::Ipc(writer_arc) => { + match Arc::try_unwrap(writer_arc) { + Ok(mutex) => { + let chunked_writer = mutex.into_inner().unwrap(); + let total_rows = chunked_writer.total_rows(); + let schema = chunked_writer.schema.clone(); + let (chunk_paths, chunk_row_ids, chunk_crcs) = + chunked_writer.finish(&mut reservation)?; + log_info!( "Successfully closed sorting chunked writer for: {}, total_rows={}, chunks={}", temp_filename, total_rows, chunk_paths.len() ); - let (crc32, row_id_mapping) = Self::finalize_sorted_chunks( - &chunk_paths, - &chunk_row_ids, - &chunk_crcs, - &filename, - index_name, - &settings.sort_columns, - &settings.reverse_sorts, - &settings.nulls_first, - writer_generation, - schema.clone(), - &mut reservation, - )?; - - // Clean up sorted chunk files only after successful finalization. - // On failure, chunks are preserved as they may be the only copy of the data. - for path in &chunk_paths { - let _ = std::fs::remove_file(path); - } - - log_debug!("CRC32 for file {}: {:#010x}", filename, crc32); + let (crc32, row_id_mapping) = Self::finalize_sorted_chunks( + &chunk_paths, + &chunk_row_ids, + &chunk_crcs, + &filename, + index_name, + &settings.sort_columns, + &settings.reverse_sorts, + &settings.nulls_first, + writer_generation, + schema.clone(), + &mut reservation, + )?; + + // Clean up sorted chunk files only after successful finalization. + // On failure, chunks are preserved as they may be the only copy of the data. + for path in &chunk_paths { + let _ = std::fs::remove_file(path); + } - let file = File::open(&filename)?; - let reader = SerializedFileReader::new(file)?; - let parquet_metadata = reader.metadata().clone(); + log_debug!("CRC32 for file {}: {:#010x}", filename, crc32); - // Detach mapping from reservation before handing to FFI/Java. - // FFI layer will track it via write_pool().grow/shrink. - if let Some(ref mapping) = row_id_mapping { - reservation.shrink(mapping.len() * std::mem::size_of::()); - } + let file = File::open(&filename)?; + let reader = SerializedFileReader::new(file)?; + let parquet_metadata = reader.metadata().clone(); - Ok(Some(FinalizeResult { - metadata: parquet_metadata, - crc32, - row_id_mapping, - })) - } - Err(_) => { - log_error!( - "ERROR: IPC Writer still in use for temp file: {}", - temp_filename - ); - Err("IPC Writer still in use".into()) + // Detach mapping from reservation before handing to FFI/Java. + // FFI layer will track it via write_pool().grow/shrink. + if let Some(ref mapping) = row_id_mapping { + reservation.shrink(mapping.len() * std::mem::size_of::()); } + + Ok(Some(FinalizeResult { + metadata: parquet_metadata, + crc32, + row_id_mapping, + })) + } + Err(_) => { + log_error!( + "ERROR: IPC Writer still in use for temp file: {}", + temp_filename + ); + Err("IPC Writer still in use".into()) } } - WriterVariant::Parquet(writer_arc) => { - match Arc::try_unwrap(writer_arc) { - Ok(mutex) => { - let writer = mutex.into_inner().unwrap(); - match writer.close() { - Ok(_) => { - let crc32 = crc_handle.map(|h| h.crc32()).unwrap_or(0); - log_info!( - "Successfully closed temp writer for: {}", - temp_filename - ); - - // Parquet variant is used for non-sorted data; just rename. - std::fs::rename(&temp_filename, &filename)?; - - log_debug!("CRC32 for file {}: {:#010x}", filename, crc32); - - let file = File::open(&filename)?; - let reader = SerializedFileReader::new(file)?; - let parquet_metadata = reader.metadata().clone(); - - Ok(Some(FinalizeResult { - metadata: parquet_metadata, - crc32, - row_id_mapping: None, - })) - } - Err(e) => { - log_error!( - "ERROR: Failed to close writer for temp file: {}", - temp_filename - ); - Err(e.into()) - } + } + WriterVariant::Parquet(writer_arc) => { + match Arc::try_unwrap(writer_arc) { + Ok(mutex) => { + let writer = mutex.into_inner().unwrap(); + match writer.close() { + Ok(_) => { + let crc32 = crc_handle.map(|h| h.crc32()).unwrap_or(0); + log_info!("Successfully closed temp writer for: {}", temp_filename); + + // Parquet variant is used for non-sorted data; just rename. + std::fs::rename(&temp_filename, &filename)?; + + log_debug!("CRC32 for file {}: {:#010x}", filename, crc32); + + let file = File::open(&filename)?; + let reader = SerializedFileReader::new(file)?; + let parquet_metadata = reader.metadata().clone(); + + Ok(Some(FinalizeResult { + metadata: parquet_metadata, + crc32, + row_id_mapping: None, + })) + } + Err(e) => { + log_error!( + "ERROR: Failed to close writer for temp file: {}", + temp_filename + ); + Err(e.into()) } } - Err(_) => { - log_error!( - "ERROR: Writer still in use for temp file: {}", - temp_filename - ); - Err("Writer still in use".into()) - } + } + Err(_) => { + log_error!( + "ERROR: Writer still in use for temp file: {}", + temp_filename + ); + Err("Writer still in use".into()) } } } - } else { - log_error!("ERROR: Writer not found for temp file: {}", temp_filename); - Err("Writer not found".into()) } } @@ -924,16 +914,58 @@ impl NativeParquetWriter { Ok(crc32) } - pub fn get_filtered_writer_memory_usage( - path_prefix: String, - ) -> Result> { - let mut total_memory = 0; - for entry in WRITERS.iter() { - if entry.key().starts_with(&path_prefix) { - total_memory += entry.value().reservation.size(); + /// Returns the native memory reserved by the writer identified by `handle`, or 0 if the handle + /// is null. Access is serialized by the Java-side monitor, so the shared borrow never aliases a + /// concurrent `&mut`/reclaim. Never panics: any unexpected panic is caught and reported as 0, so + /// this can never unwind across the FFI boundary or fail the Java caller. + pub fn get_writer_memory_usage(handle: *const WriterState) -> usize { + if handle.is_null() { + return 0; + } + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + // Safety: live handle, access serialized by the Java-side monitor. + // Shared reborrow only (not Box::from_raw): does NOT take ownership, so the writer is not dropped here. + let state = unsafe { &*handle }; + state.reservation.size() + })) + .unwrap_or_else(|_| { + log_error!("get_writer_memory_usage: swallowed panic; reporting 0"); + 0 + }) + } + + /// Best-effort teardown of a writer that Java is abandoning without finalizing (idempotent on + /// the Java side via a released-flag). Reclaims the Box (dropping the reservation and closing + /// file handles) and deletes any leftover temp/IPC/sorted-chunk artifacts. Never fails and + /// never panics — filesystem errors are ignored and any unexpected panic is caught, so this can + /// never unwind across the FFI boundary or fail the Java caller. + pub fn free_writer(handle: *mut WriterState) { + if handle.is_null() { + return; + } + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + // Safety: Java guarantees free is called at most once per handle (released CAS + Cleaner). + // Box::from_raw TAKES ownership: the writer is dropped exactly once here (abandonment path). + let state = unsafe { *Box::from_raw(handle) }; + let temp_filename = state.temp_filename.clone(); + let ipc_base = format!("{}{}", temp_filename, IPC_STAGING_SUFFIX); + // Drop first so the underlying File handles are closed before we unlink. + drop(state); + let _ = std::fs::remove_file(&temp_filename); + let _ = std::fs::remove_file(&ipc_base); + // Sorted-chunk files are named "{ipc_base}.sorted_chunk_{i}.parquet"; remove any remaining. + let mut idx = 0usize; + loop { + let chunk = format!("{}.sorted_chunk_{}.parquet", ipc_base, idx); + if std::fs::remove_file(&chunk).is_err() { + break; + } + idx += 1; } + })); + if result.is_err() { + log_error!("free_writer: swallowed panic while freeing writer handle"); } - Ok(total_memory) } pub fn get_file_metadata( diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/tests/writer_integration_tests.rs b/sandbox/plugins/parquet-data-format/src/main/rust/tests/writer_integration_tests.rs index cc6b41fe94df4..9e8456b0f6d84 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/tests/writer_integration_tests.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/tests/writer_integration_tests.rs @@ -55,7 +55,7 @@ fn test_concurrent_writer_creation() { let file_path = temp_dir_path.join(format!("concurrent_{}.parquet", i)); let filename = file_path.to_string_lossy().to_string(); let (_schema, schema_ptr) = create_test_ffi_schema(); - if NativeParquetWriter::create_writer( + if test_create_writer( filename.clone(), "test-index".to_string(), schema_ptr, @@ -67,7 +67,7 @@ fn test_concurrent_writer_creation() { .is_ok() { success_count.fetch_add(1, Ordering::SeqCst); - let _ = NativeParquetWriter::finalize_writer(filename); + let _ = test_finalize_writer(filename); } cleanup_ffi_schema(schema_ptr); }); @@ -91,7 +91,9 @@ fn test_concurrent_close_operations_same_file() { let filename = filename.clone(); let success_count = Arc::clone(&success_count); let handle = thread::spawn(move || { - if NativeParquetWriter::finalize_writer(filename).is_ok() { + // Exactly one thread wins the handle (the test map's remove is atomic) and performs the + // real finalize (Ok(Some)); the others observe a null handle and no-op (Ok(None)). + if matches!(test_finalize_writer(filename), Ok(Some(_))) { success_count.fetch_add(1, Ordering::SeqCst); } }); @@ -117,7 +119,7 @@ fn test_concurrent_writes_same_file() { let success_count = Arc::clone(&success_count); let handle = thread::spawn(move || { let (array_ptr, data_schema_ptr) = create_test_ffi_data().unwrap(); - if NativeParquetWriter::write_data(filename, array_ptr, data_schema_ptr).is_ok() { + if test_write_data(filename, array_ptr, data_schema_ptr).is_ok() { success_count.fetch_add(1, Ordering::SeqCst); } cleanup_ffi_data(array_ptr, data_schema_ptr); @@ -157,9 +159,7 @@ fn test_concurrent_writes_different_files() { let handle = thread::spawn(move || { for _ in 0..2 { let (array_ptr, data_schema_ptr) = create_test_ffi_data().unwrap(); - if NativeParquetWriter::write_data(filename.clone(), array_ptr, data_schema_ptr) - .is_ok() - { + if test_write_data(filename.clone(), array_ptr, data_schema_ptr).is_ok() { success_count.fetch_add(1, Ordering::SeqCst); } cleanup_ffi_data(array_ptr, data_schema_ptr); @@ -192,7 +192,7 @@ fn test_concurrent_complete_writer_lifecycle() { let filename = file_path.to_string_lossy().to_string(); let (_schema, schema_ptr) = create_test_ffi_schema(); - if NativeParquetWriter::create_writer( + if test_create_writer( filename.clone(), "test-index".to_string(), schema_ptr, @@ -205,14 +205,11 @@ fn test_concurrent_complete_writer_lifecycle() { { let (array_ptr, data_schema_ptr) = create_test_ffi_data().unwrap(); let write_ok = - NativeParquetWriter::write_data(filename.clone(), array_ptr, data_schema_ptr) - .is_ok(); + test_write_data(filename.clone(), array_ptr, data_schema_ptr).is_ok(); cleanup_ffi_data(array_ptr, data_schema_ptr); if write_ok { - if let Ok(Some(metadata)) = - NativeParquetWriter::finalize_writer(filename.clone()) - { + if let Ok(Some(metadata)) = test_finalize_writer(filename.clone()) { if metadata.metadata.file_metadata().num_rows() == 3 && file_path.exists() { success_count.fetch_add(1, Ordering::SeqCst); } @@ -237,7 +234,7 @@ fn test_ipc_staging_sorted_writer_integration() { let (_temp_dir, filename) = get_temp_file_path("ipc_integ_sorted.parquet"); let (_schema, schema_ptr) = create_test_ffi_schema(); - NativeParquetWriter::create_writer( + test_create_writer( filename.clone(), "test-index".to_string(), schema_ptr, @@ -252,11 +249,11 @@ fn test_ipc_staging_sorted_writer_integration() { for batch_ids in [vec![50, 30, 10], vec![40, 20, 60]] { let names: Vec> = batch_ids.iter().map(|_| Some("x")).collect(); let (ap, sp) = create_test_ffi_data_with_ids(batch_ids, names).unwrap(); - NativeParquetWriter::write_data(filename.clone(), ap, sp).unwrap(); + test_write_data(filename.clone(), ap, sp).unwrap(); cleanup_ffi_data(ap, sp); } - let result = NativeParquetWriter::finalize_writer(filename.clone()); + let result = test_finalize_writer(filename.clone()); assert!(result.is_ok()); let metadata = result.unwrap().unwrap(); assert_eq!(metadata.metadata.file_metadata().num_rows(), 6); @@ -285,7 +282,7 @@ fn test_ipc_staging_concurrent_sorted_lifecycle() { let filename = file_path.to_string_lossy().to_string(); let (_schema, schema_ptr) = create_test_ffi_schema(); - if NativeParquetWriter::create_writer( + if test_create_writer( filename.clone(), "test-index".to_string(), schema_ptr, @@ -301,13 +298,11 @@ fn test_ipc_staging_concurrent_sorted_lifecycle() { vec![Some("C"), Some("A"), Some("B")], ) .unwrap(); - let write_ok = NativeParquetWriter::write_data(filename.clone(), ap, sp).is_ok(); + let write_ok = test_write_data(filename.clone(), ap, sp).is_ok(); cleanup_ffi_data(ap, sp); if write_ok { - if let Ok(Some(metadata)) = - NativeParquetWriter::finalize_writer(filename.clone()) - { + if let Ok(Some(metadata)) = test_finalize_writer(filename.clone()) { if metadata.metadata.file_metadata().num_rows() == 3 && file_path.exists() { let ids = read_parquet_file_sorted_ids(&filename); if ids == vec![10, 20, 30] { @@ -353,7 +348,7 @@ fn test_ipc_and_parquet_mixed_concurrent_lifecycle() { let reverse = if use_sort { vec![false] } else { vec![] }; let nulls = if use_sort { vec![false] } else { vec![] }; - if NativeParquetWriter::create_writer( + if test_create_writer( filename.clone(), "test-index".to_string(), schema_ptr, @@ -369,13 +364,11 @@ fn test_ipc_and_parquet_mixed_concurrent_lifecycle() { vec![Some("C"), Some("A"), Some("B")], ) .unwrap(); - let write_ok = NativeParquetWriter::write_data(filename.clone(), ap, sp).is_ok(); + let write_ok = test_write_data(filename.clone(), ap, sp).is_ok(); cleanup_ffi_data(ap, sp); if write_ok { - if let Ok(Some(metadata)) = - NativeParquetWriter::finalize_writer(filename.clone()) - { + if let Ok(Some(metadata)) = test_finalize_writer(filename.clone()) { if metadata.metadata.file_metadata().num_rows() == 3 && file_path.exists() { let ids = read_parquet_file_sorted_ids(&filename); let expected = if use_sort { @@ -414,7 +407,6 @@ use std::fs::File; fn read_column_compression(path: &str, col_name: &str) -> Compression { let reader = SerializedFileReader::new(File::open(path).unwrap()).unwrap(); let meta = reader.metadata(); - let schema = meta.file_metadata().schema_descr(); let rg = meta.row_group(0); for i in 0..rg.num_columns() { let col = rg.column(i); @@ -462,9 +454,9 @@ fn test_index_level_field_compression_applied() { let (_tmp, path) = get_temp_file_path("idx_field_comp.parquet"); let (_schema, schema_ptr) = create_writer_and_assert_success_for_index(&path, index); let (ap, sp) = create_test_ffi_data().unwrap(); - NativeParquetWriter::write_data(path.clone(), ap, sp).unwrap(); + test_write_data(path.clone(), ap, sp).unwrap(); cleanup_ffi_data(ap, sp); - NativeParquetWriter::finalize_writer(path.clone()).unwrap(); + test_finalize_writer(path.clone()).unwrap(); cleanup_ffi_schema(schema_ptr); assert!(matches!( @@ -492,9 +484,9 @@ fn test_cluster_level_type_compression_fallback() { let (_tmp, path) = get_temp_file_path("cluster_type_comp.parquet"); let (_schema, schema_ptr) = create_writer_and_assert_success_for_index(&path, index); let (ap, sp) = create_test_ffi_data().unwrap(); - NativeParquetWriter::write_data(path.clone(), ap, sp).unwrap(); + test_write_data(path.clone(), ap, sp).unwrap(); cleanup_ffi_data(ap, sp); - NativeParquetWriter::finalize_writer(path.clone()).unwrap(); + test_finalize_writer(path.clone()).unwrap(); cleanup_ffi_schema(schema_ptr); // "id" is Int32 → should get SNAPPY from cluster type config @@ -532,9 +524,9 @@ fn test_index_level_overrides_cluster_level_compression() { let (_tmp, path) = get_temp_file_path("idx_overrides_cluster_comp.parquet"); let (_schema, schema_ptr) = create_writer_and_assert_success_for_index(&path, index); let (ap, sp) = create_test_ffi_data().unwrap(); - NativeParquetWriter::write_data(path.clone(), ap, sp).unwrap(); + test_write_data(path.clone(), ap, sp).unwrap(); cleanup_ffi_data(ap, sp); - NativeParquetWriter::finalize_writer(path.clone()).unwrap(); + test_finalize_writer(path.clone()).unwrap(); cleanup_ffi_schema(schema_ptr); // Index says UNCOMPRESSED for "id", cluster says SNAPPY for int32 — index wins @@ -563,9 +555,9 @@ fn test_cluster_level_type_encoding_fallback() { let (_tmp, path) = get_temp_file_path("cluster_type_enc.parquet"); let (_schema, schema_ptr) = create_writer_and_assert_success_for_index(&path, index); let (ap, sp) = create_test_ffi_data().unwrap(); - NativeParquetWriter::write_data(path.clone(), ap, sp).unwrap(); + test_write_data(path.clone(), ap, sp).unwrap(); cleanup_ffi_data(ap, sp); - NativeParquetWriter::finalize_writer(path.clone()).unwrap(); + test_finalize_writer(path.clone()).unwrap(); cleanup_ffi_schema(schema_ptr); // "id" is Int32 → should get DELTA_BINARY_PACKED from cluster type config @@ -605,9 +597,9 @@ fn test_index_level_overrides_cluster_level_encoding() { let (_tmp, path) = get_temp_file_path("idx_overrides_cluster_enc.parquet"); let (_schema, schema_ptr) = create_writer_and_assert_success_for_index(&path, index); let (ap, sp) = create_test_ffi_data().unwrap(); - NativeParquetWriter::write_data(path.clone(), ap, sp).unwrap(); + test_write_data(path.clone(), ap, sp).unwrap(); cleanup_ffi_data(ap, sp); - NativeParquetWriter::finalize_writer(path.clone()).unwrap(); + test_finalize_writer(path.clone()).unwrap(); cleanup_ffi_schema(schema_ptr); // Index says PLAIN for "id", cluster says DELTA_BINARY_PACKED for int32 — index wins @@ -649,9 +641,9 @@ fn test_mixed_index_and_cluster_level_configs() { let (_tmp, path) = get_temp_file_path("mixed_idx_cluster.parquet"); let (_schema, schema_ptr) = create_writer_and_assert_success_for_index(&path, index); let (ap, sp) = create_test_ffi_data().unwrap(); - NativeParquetWriter::write_data(path.clone(), ap, sp).unwrap(); + test_write_data(path.clone(), ap, sp).unwrap(); cleanup_ffi_data(ap, sp); - NativeParquetWriter::finalize_writer(path.clone()).unwrap(); + test_finalize_writer(path.clone()).unwrap(); cleanup_ffi_schema(schema_ptr); assert!(matches!( @@ -671,7 +663,7 @@ fn create_writer_and_assert_success_for_index( index: &str, ) -> (std::sync::Arc, i64) { let (schema, schema_ptr) = create_test_ffi_schema(); - NativeParquetWriter::create_writer( + test_create_writer( filename.to_string(), index.to_string(), schema_ptr, @@ -698,7 +690,7 @@ fn create_three_col_writer( ])); let ffi_schema = FFI_ArrowSchema::try_from(schema.as_ref()).unwrap(); let schema_ptr = Box::into_raw(Box::new(ffi_schema)) as i64; - NativeParquetWriter::create_writer( + test_create_writer( filename.to_string(), index.to_string(), schema_ptr, @@ -737,7 +729,7 @@ fn write_three_col_data(filename: &str) { let ffi_schema = FFI_ArrowSchema::try_from(schema.as_ref()).unwrap(); let array_ptr = Box::into_raw(Box::new(ffi_array)) as i64; let schema_ptr = Box::into_raw(Box::new(ffi_schema)) as i64; - NativeParquetWriter::write_data(filename.to_string(), array_ptr, schema_ptr).unwrap(); + test_write_data(filename.to_string(), array_ptr, schema_ptr).unwrap(); unsafe { let _ = Box::from_raw(array_ptr as *mut FFI_ArrowArray); let _ = Box::from_raw(schema_ptr as *mut FFI_ArrowSchema); @@ -776,7 +768,7 @@ fn test_three_tier_compression_fallback() { let (_tmp, path) = get_temp_file_path("three_tier_comp.parquet"); let (_schema, schema_ptr) = create_three_col_writer(&path, index); write_three_col_data(&path); - NativeParquetWriter::finalize_writer(path.clone()).unwrap(); + test_finalize_writer(path.clone()).unwrap(); cleanup_ffi_schema(schema_ptr); // index level wins for "id" diff --git a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/bridge/NativeParquetWriterRecoveryTests.java b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/bridge/NativeParquetWriterRecoveryTests.java new file mode 100644 index 0000000000000..6c3275dd07abf --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/bridge/NativeParquetWriterRecoveryTests.java @@ -0,0 +1,99 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.parquet.bridge; + +import org.apache.arrow.c.ArrowSchema; +import org.apache.arrow.c.Data; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.opensearch.nativebridge.spi.ArrowExport; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.List; + +/** + * Regression test for the "Writer already exists" recovery failure. + * + *

When the native Parquet writer is left un-finalized/un-closed on the Rust side (e.g. a + * Rust-side failure aborts a flush before {@code finalize}, or the shard fails without a clean + * writer close), the native registry keeps the entry. On shard recovery the engine re-creates a + * writer for the same generation — and therefore the same underlying file/temp name. + * With the filename-keyed native registry this second creation throws + * {@code "Writer already exists for this file"}, so the shard can never come back up. + * + *

This test reproduces exactly that sequence at the Java↔Rust boundary: initialize a writer for + * a file and deliberately never close/finalize it (the abandoned/dangling writer), then initialize + * a second writer for the same file (the recovery re-create). Recovery must succeed. + * + *

Expected: FAILS on the current (filename-keyed) implementation — the second + * {@code initialize} throws {@code IOException("Writer already exists for this file")}. Once the + * native writer lifecycle is owned by Java via opaque handles (no filename collision), it PASSES. + * This test is intentionally left unchanged across that fix so it acts as the regression guard. + * + *

Note: it does not close the writers (no {@code close()} exists pre-fix); post-fix the handles + * are reclaimed by the {@code Cleaner} backstop when the writers become unreachable. + */ +public class NativeParquetWriterRecoveryTests extends OpenSearchTestCase { + + private BufferAllocator allocator; + private Schema schema; + + @Override + public void setUp() throws Exception { + super.setUp(); + RustBridge.initLogger(); + allocator = new RootAllocator(); + schema = new Schema( + List.of( + new Field("id", FieldType.nullable(new ArrowType.Int(32, true)), null), + new Field("name", FieldType.nullable(new ArrowType.Utf8()), null) + ) + ); + } + + @Override + public void tearDown() throws Exception { + allocator.close(); + super.tearDown(); + } + + public void testRecoveryReCreatesWriterForSameFileAfterUnclosedWriter() throws Exception { + String filePath = createTempDir().resolve("recovery.parquet").toString(); + + // Writer 1: initialize the native writer, then abandon it WITHOUT finalize/close — + // simulating a Rust-side failure where cleanup never ran. + NativeParquetWriter writer1 = new NativeParquetWriter(filePath); + try (ArrowExport export = exportSchema()) { + writer1.initialize("test-index", export.getSchemaAddress(), ParquetSortConfig.empty(), 0L); + } + assertTrue("writer1 should initialize", writer1.isInitialized()); + + // Recovery: a fresh writer for the SAME file (same generation -> same temp file name). + // On the current filename-keyed native registry this throws "Writer already exists", + // which is what leaves a recovering shard stuck red. It must be allowed to succeed. + NativeParquetWriter writer2 = new NativeParquetWriter(filePath); + try (ArrowExport export = exportSchema()) { + writer2.initialize("test-index", export.getSchemaAddress(), ParquetSortConfig.empty(), 0L); + } + assertTrue( + "recovery must be able to re-create a native writer for the same file after an unclosed writer", + writer2.isInitialized() + ); + } + + private ArrowExport exportSchema() { + ArrowSchema arrowSchema = ArrowSchema.allocateNew(allocator); + Data.exportSchema(allocator, schema, null, arrowSchema); + return new ArrowExport(null, arrowSchema); + } +} diff --git a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/bridge/NativeParquetWriterTests.java b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/bridge/NativeParquetWriterTests.java index acf76402bae9c..52372b916d773 100644 --- a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/bridge/NativeParquetWriterTests.java +++ b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/bridge/NativeParquetWriterTests.java @@ -185,25 +185,26 @@ public void testWriteWithSchemaMismatch() throws Exception { expectThrows(IOException.class, writer::flush); } - public void testCreateDuplicateWriterForSameFile() throws Exception { + public void testRecreateWriterForSameFileAfterAbandonedWriter() throws Exception { String filePath = createTempDir().resolve("duplicate.parquet").toString(); + + // writer1 is initialized then abandoned WITHOUT flush/close — a dangling native writer, + // exactly what a Rust-side failure leaves behind. NativeParquetWriter writer1 = createWriter(filePath); - // Initialize writer1 by writing data + // With Java-owned handles there is no filename-keyed native registry, so a second writer + // for the same file can be created (this is what lets a shard recover) and can complete. + NativeParquetWriter writer2 = createWriter(filePath); try (ArrowExport export = exportData(new int[] { 1 }, new String[] { "a" }, new long[] { 1L })) { - writer1.write(export.getArrayAddress(), export.getSchemaAddress()); - } - - // Native side rejects creating a second writer for the same file - NativeParquetWriter writer2 = new NativeParquetWriter(filePath); - try (ArrowExport export = exportSchema()) { - expectThrows( - IOException.class, - () -> writer2.initialize("test-index", export.getSchemaAddress(), ParquetSortConfig.empty(), 0L) - ); + writer2.write(export.getArrayAddress(), export.getSchemaAddress()); } + writer2.flush(); + assertNotNull(writer2.getMetadata()); + assertEquals(1, writer2.getMetadata().numRows()); + assertTrue("Parquet file should exist after writer2 flush", Files.exists(Path.of(filePath))); - writer1.flush(); + // Release the abandoned writer's native handle (best-effort cleanup). + writer1.close(); } public void testWriteEmptyBatch() throws Exception { From c45a08c13752047706395bfb1bd99904e0eb2d33 Mon Sep 17 00:00:00 2001 From: rayshrey Date: Thu, 23 Jul 2026 17:20:27 +0530 Subject: [PATCH 2/2] Basic working ObjectStore implementation on write and merge side Signed-off-by: rayshrey --- .../tiered-storage/src/main/rust/src/ffm.rs | 50 ++ .../composite/CompositeRefreshSortedIT.java | 67 +++ .../parquet/bridge/NativeParquetWriter.java | 21 +- .../opensearch/parquet/bridge/RustBridge.java | 25 +- .../parquet/engine/ParquetIndexingEngine.java | 33 +- .../merge/NativeParquetMergeStrategy.java | 26 +- .../parquet/store/TieredStorageBridge.java | 46 ++ .../opensearch/parquet/vsr/VSRManager.java | 43 +- .../parquet/writer/ParquetWriter.java | 9 +- .../src/main/rust/Cargo.toml | 9 +- .../src/main/rust/src/crc_writer.rs | 13 + .../src/main/rust/src/ffm.rs | 13 + .../src/main/rust/src/lib.rs | 1 + .../src/main/rust/src/merge/context.rs | 45 +- .../src/main/rust/src/merge/cursor.rs | 64 +-- .../src/main/rust/src/merge/io_task.rs | 70 ++- .../src/main/rust/src/merge/mod.rs | 1 + .../src/main/rust/src/merge/reader.rs | 130 +++++ .../src/main/rust/src/merge/sorted.rs | 9 + .../src/main/rust/src/merge/unsorted.rs | 49 +- .../src/main/rust/src/store_io.rs | 72 +++ .../src/main/rust/src/test_utils.rs | 1 + .../src/main/rust/src/writer.rs | 493 +++++++++++++++--- .../rust/tests/merge_integration_tests.rs | 450 +++++++--------- .../src/main/rust/tests/sort_types_tests.rs | 208 ++++---- .../bridge/ParquetMergeIntegrationTests.java | 6 +- 26 files changed, 1401 insertions(+), 553 deletions(-) create mode 100644 sandbox/plugins/parquet-data-format/src/main/rust/src/merge/reader.rs create mode 100644 sandbox/plugins/parquet-data-format/src/main/rust/src/store_io.rs diff --git a/sandbox/libs/tiered-storage/src/main/rust/src/ffm.rs b/sandbox/libs/tiered-storage/src/main/rust/src/ffm.rs index 2d65e61cdb3d2..553dbbe6b462d 100644 --- a/sandbox/libs/tiered-storage/src/main/rust/src/ffm.rs +++ b/sandbox/libs/tiered-storage/src/main/rust/src/ffm.rs @@ -58,6 +58,56 @@ unsafe fn arc_from_ptr(ptr: i64) -> Result, String> { Ok(Arc::from_raw(raw)) } +// --------------------------------------------------------------------------- +// Generic ObjectStore handle lifecycle (write-path, shard-scoped) +// --------------------------------------------------------------------------- +// +// These entry points mint a plain `Box>` handle that the +// parquet write/merge paths use as a shard-scoped sink/source. For hot indices +// this is a `LocalFileSystem` rooted at the shard directory; the warm path can +// later hand back the same-shaped handle backed by a `TieredObjectStore`, so +// consumers (writer/merge/search) never change. +// +// Ownership model: the engine owns the master handle and is the ONLY caller of +// `os_destroy_store` (exactly once, at shard close). Consumers (parquet +// writer/merge) clone the `Arc` inline from the raw box pointer so the store +// outlives any in-flight op. + +/// Create a local filesystem `ObjectStore` rooted at `path`, returned as a +/// `Box>` raw pointer (i64). Object paths are then plain +/// filenames relative to `path`. Free with [`os_destroy_store`]. +#[ffm_safe] +#[no_mangle] +pub unsafe extern "C" fn os_create_local_store(path_ptr: *const u8, path_len: i64) -> i64 { + let path = str_from_raw(path_ptr, path_len) + .map_err(|e| format!("os_create_local_store path: {}", e))?; + let store: Arc = Arc::new( + object_store::local::LocalFileSystem::new_with_prefix(path).map_err(|e| { + format!( + "os_create_local_store: failed to open LocalFileSystem at '{}': {}", + path, e + ) + })?, + ); + let ptr = Box::into_raw(Box::new(store)) as i64; + native_bridge_common::log_info!("ffm: os_create_local_store root='{}' ok", path); + Ok(ptr) +} + +/// Drop a `Box>` created by [`os_create_local_store`], +/// decrementing the `Arc` strong count. Must be called exactly once, by the +/// engine, at shard close. Outstanding consumer clones keep the store alive. +#[ffm_safe] +#[no_mangle] +pub unsafe extern "C" fn os_destroy_store(handle: i64) -> i64 { + if handle == NULL_PTR { + return Err("os_destroy_store: null pointer (0)".to_string()); + } + let _ = Box::from_raw(handle as *mut Arc); + native_bridge_common::log_info!("ffm: os_destroy_store ok"); + Ok(0) +} + // --------------------------------------------------------------------------- // Public FFM exports // --------------------------------------------------------------------------- diff --git a/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/CompositeRefreshSortedIT.java b/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/CompositeRefreshSortedIT.java index 3e2dab4583668..b4ad7a437c7da 100644 --- a/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/CompositeRefreshSortedIT.java +++ b/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/CompositeRefreshSortedIT.java @@ -189,6 +189,73 @@ public void testSortedRefreshWithRandomizedData() throws Exception { verifyLuceneRowIdSequential(); } + /** + * Forces many small sorted chunks via a tiny {@code index.parquet.sort_in_memory_threshold}, + * so that finalize runs the store-backed k-way merge over N>1 chunks. This is the only test + * that exercises the multi-chunk merge path: sorted chunks written to / read from / deleted in + * the ObjectStore plus the merge output streamed back through the store. (With a 16kb threshold + * and 10k docs this flushes dozens of chunks — see the "merging N pre-sorted chunks" native + * log.) Correctness of the merged output is asserted via full row count and global sort order. + */ + public void testSortedRefreshMultiChunkThroughStore() throws Exception { + int numFields = 5; + String[] fieldNames = new String[numFields]; + String[] sortFields = new String[numFields]; + String[] sortOrders = new String[numFields]; + String[] sortMissing = new String[numFields]; + for (int f = 0; f < numFields; f++) { + fieldNames[f] = "field_" + f; + sortFields[f] = fieldNames[f]; + sortOrders[f] = (f % 2 == 0) ? "asc" : "desc"; + sortMissing[f] = sortOrders[f].equals("asc") ? "_last" : "_first"; + } + + Settings settings = Settings.builder() + .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1) + .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0) + .put("index.refresh_interval", "-1") + .put("index.pluggable.dataformat.enabled", true) + .put("index.pluggable.dataformat", "composite") + .put("index.composite.primary_data_format", "parquet") + .putList("index.composite.secondary_data_formats", "lucene") + // Tiny threshold -> many sorted chunks -> multi-chunk k-way merge through the store. + .put("index.parquet.sort_in_memory_threshold", "16kb") + .putList("index.sort.field", sortFields) + .putList("index.sort.order", sortOrders) + .putList("index.sort.missing", sortMissing) + .build(); + + String[] mappingArgs = new String[numFields * 2]; + for (int f = 0; f < numFields; f++) { + mappingArgs[f * 2] = fieldNames[f]; + mappingArgs[f * 2 + 1] = "type=integer"; + } + client().admin().indices().prepareCreate(INDEX_NAME).setSettings(settings).setMapping(mappingArgs).get(); + ensureGreen(INDEX_NAME); + + int totalDocs = 10000; + for (int i = 0; i < totalDocs; i++) { + Map source = new HashMap<>(); + for (int f = 0; f < numFields; f++) { + // Occasionally emit null values to exercise null handling across chunk boundaries + if (randomIntBetween(0, 9) == 0) { + continue; // skip field → null + } + source.put(fieldNames[f], randomIntBetween(0, 1000)); + } + IndexResponse response = client().prepareIndex().setIndex(INDEX_NAME).setSource(source).get(); + assertEquals(RestStatus.CREATED, response.status()); + } + + flushAndRefresh(); + + DataformatAwareCatalogSnapshot snapshot = getCatalogSnapshot(); + verifyParquetRowCount(snapshot, totalDocs); + verifyParquetSortOrderMultiField(snapshot, fieldNames, sortOrders, sortMissing); + verifyLuceneDocCount(totalDocs); + verifyLuceneRowIdSequential(); + } + /** * Parquet primary + Lucene secondary without sort. Without a sort * permutation, Parquet does not produce a {@code RowIdMapping} and the diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeParquetWriter.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeParquetWriter.java index a438865c26e56..cbad4473faeb4 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeParquetWriter.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeParquetWriter.java @@ -38,6 +38,12 @@ public class NativeParquetWriter { private final SetOnce metadata = new SetOnce<>(); private final SetOnce rowIdMapping = new SetOnce<>(); private final ParquetShardStatsTracker stats; + /** + * Shard-scoped native ObjectStore handle ({@code Box>} pointer) minted by + * {@code os_create_local_store} and owned by the engine, or 0 for the legacy local-file path. + * Passed to the native writer at {@link #initialize} time so finalize publishes through the store. + */ + private final long storePtr; private volatile boolean initialized = false; /** Reclaims leaked native writers if this object is GC'd without an explicit close/flush. */ @@ -62,8 +68,21 @@ public class NativeParquetWriter { * @param stats shard-level stats tracker */ public NativeParquetWriter(String filePath, ParquetShardStatsTracker stats) { + this(filePath, stats, 0L); + } + + /** + * Creates a new NativeParquetWriter handle bound to a shard-scoped ObjectStore. + * + * @param filePath the path to the Parquet file to write + * @param stats shard-level stats tracker + * @param storePtr shard ObjectStore handle ({@code os_create_local_store} pointer), or 0 for + * the legacy local-file path + */ + public NativeParquetWriter(String filePath, ParquetShardStatsTracker stats, long storePtr) { this.filePath = filePath; this.stats = stats; + this.storePtr = storePtr; } /** @@ -91,7 +110,7 @@ public synchronized void initialize(String indexName, long schemaAddress, Parque if (initialized) { throw new IllegalStateException("Writer already initialized: " + filePath); } - handle = RustBridge.createWriter(filePath, indexName, schemaAddress, sortConfig, writerGeneration); + handle = RustBridge.createWriter(filePath, indexName, schemaAddress, sortConfig, writerGeneration, storePtr); initialized = true; // GC backstop: if this writer is dropped without close()/flush(), free the native handle. // HandleCleanup holds only the primitive handle + shared released flag (never `this`). diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/RustBridge.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/RustBridge.java index 43718cab40eb6..73db1c8e090f2 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/RustBridge.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/RustBridge.java @@ -77,7 +77,8 @@ public class RustBridge { ValueLayout.JAVA_LONG, // reverse_sorts (vals, count) ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, // nulls_first (vals, count) - ValueLayout.JAVA_LONG // writer_generation + ValueLayout.JAVA_LONG, // writer_generation + ValueLayout.JAVA_LONG // store_handle (0 = legacy local file path) ) ); WRITE = linker.downcallHandle( @@ -231,7 +232,8 @@ public class RustBridge { ValueLayout.ADDRESS, // out_gen_count ValueLayout.ADDRESS, // out_flush_and_sort_chunk_count ValueLayout.ADDRESS, // out_flush_and_sort_chunk_time_millis - ValueLayout.ADDRESS // out_row_id_mapping_max + ValueLayout.ADDRESS, // out_row_id_mapping_max + ValueLayout.JAVA_LONG // store_handle (0 = local segment merge) ) ); FREE_MERGE_RESULT = linker.downcallHandle( @@ -295,8 +297,14 @@ public class RustBridge { public static void initLogger() {} - static long createWriter(String file, String indexName, long schemaAddress, ParquetSortConfig sortConfig, long writerGeneration) - throws IOException { + static long createWriter( + String file, + String indexName, + long schemaAddress, + ParquetSortConfig sortConfig, + long writerGeneration, + long storeHandle + ) throws IOException { try (var call = new NativeCall()) { var f = call.str(file); var idx = call.str(indexName); @@ -317,7 +325,8 @@ static long createWriter(String file, String indexName, long schemaAddress, Parq (long) sortConfig.reverseSorts().size(), nullsFirstArray, (long) sortConfig.nullsFirst().size(), - writerGeneration + writerGeneration, + storeHandle ); } } @@ -548,7 +557,8 @@ public static MergeFilesResult mergeParquetFilesInRust( List inputFiles, String outputFile, String indexName, - long outputWriterGeneration + long outputWriterGeneration, + long storePtr ) { String[] paths = inputFiles.stream().map(Path::toString).toArray(String[]::new); try (var call = new NativeCall()) { @@ -598,7 +608,8 @@ public static MergeFilesResult mergeParquetFilesInRust( outGenCount, outFlushChunkCount, outFlushChunkTimeMillis, - outRowIdMappingMax + outRowIdMappingMax, + storePtr ); int createdByLen = (int) createdByOut.lenOut().get(ValueLayout.JAVA_LONG, 0); diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/engine/ParquetIndexingEngine.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/engine/ParquetIndexingEngine.java index 95b6cb512ad59..6694d93fd8149 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/engine/ParquetIndexingEngine.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/engine/ParquetIndexingEngine.java @@ -28,6 +28,7 @@ import org.opensearch.parquet.ParquetSettings; import org.opensearch.parquet.bridge.NativeSettings; import org.opensearch.parquet.bridge.RustBridge; +import org.opensearch.parquet.store.TieredStorageBridge; import org.opensearch.parquet.memory.ArrowBufferPool; import org.opensearch.parquet.merge.NativeParquetMergeStrategy; import org.opensearch.parquet.merge.ParquetMergeExecutor; @@ -89,6 +90,12 @@ public class ParquetIndexingEngine implements IndexingExecutionEngine createWriter(WriterConfig config) { threadPool, checksumStrategy, statsTracker, - activeWriters::remove + activeWriters::remove, + storePtr ); activeWriters.add(writer); return writer; @@ -365,6 +385,15 @@ public void close() throws IOException { ); } bufferPool.close(); + // Free the shard-scoped native ObjectStore. Writers are drained before engine close, so + // their Arc clones are already dropped; this releases the engine's master reference. + if (storePtr > 0) { + try { + TieredStorageBridge.destroyStore(storePtr); + } catch (Exception e) { + logger.warn("Failed to destroy shard ObjectStore (ptr={})", storePtr, e); + } + } } /** diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/NativeParquetMergeStrategy.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/NativeParquetMergeStrategy.java index 06acac2f164e9..f0cf3007191d9 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/NativeParquetMergeStrategy.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/NativeParquetMergeStrategy.java @@ -46,19 +46,27 @@ public class NativeParquetMergeStrategy implements ParquetMergeStrategy { private final ShardPath shardPath; private final TriConsumer checksumUpdater; private final ParquetShardStatsTracker stats; + /** + * Shard-scoped native ObjectStore handle (owned by the engine), or 0 for local segment merges. + * When set, merge inputs/output are passed to native as object paths (basenames relative to the + * shard store root) so the k-way merge reads inputs from / streams output to the store. + */ + private final long storePtr; public NativeParquetMergeStrategy( DataFormat dataFormat, String indexName, ShardPath shardPath, TriConsumer checksumUpdater, - ParquetShardStatsTracker stats + ParquetShardStatsTracker stats, + long storePtr ) { this.dataFormat = dataFormat; this.indexName = indexName; this.shardPath = shardPath; this.checksumUpdater = checksumUpdater; this.stats = stats; + this.storePtr = storePtr; } @Override @@ -83,10 +91,24 @@ public MergeResult mergeParquetFiles(MergeInput mergeInput) { Path mergedFilePath = ParquetIndexingEngine.buildParquetFilePath(shardPath, writerGeneration, "merged"); String mergedFileName = mergedFilePath.getFileName().toString(); + // When a shard ObjectStore is present, pass object paths (basenames relative to the shard + // store root) to native; the merge reads inputs from / streams output to the store. Inputs + // live at `/` (store root == the shard parquet dir). Otherwise pass + // absolute local paths. + boolean useStore = storePtr > 0; + List nativeInputPaths = useStore ? files.stream().map(mono -> Path.of(mono.file())).toList() : filePaths; + String nativeOutputPath = useStore ? mergedFileName : mergedFilePath.toString(); + long startNanos = System.nanoTime(); try { // Merge files in Rust - MergeFilesResult merged = RustBridge.mergeParquetFilesInRust(filePaths, mergedFilePath.toString(), indexName, writerGeneration); + MergeFilesResult merged = RustBridge.mergeParquetFilesInRust( + nativeInputPaths, + nativeOutputPath, + indexName, + writerGeneration, + storePtr + ); ParquetFileMetadata mergeMetadata = merged.metadata(); RowIdMapping rowIdMapping = merged.rowIdMapping(); diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/store/TieredStorageBridge.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/store/TieredStorageBridge.java index 4a0d56261f6d8..4d410850cd6c1 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/store/TieredStorageBridge.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/store/TieredStorageBridge.java @@ -36,6 +36,8 @@ public final class TieredStorageBridge { private static final MethodHandle REMOVE_FILE; private static final MethodHandle GET_OBJECT_STORE_BOX_PTR; private static final MethodHandle DESTROY_OBJECT_STORE_BOX_PTR; + private static final MethodHandle OS_CREATE_LOCAL_STORE; + private static final MethodHandle OS_DESTROY_STORE; static { SymbolLookup lib = NativeLibraryLoader.symbolLookup(); @@ -79,10 +81,54 @@ public final class TieredStorageBridge { .map(sym -> linker.downcallHandle(sym, FunctionDescriptor.of(ValueLayout.JAVA_LONG, ValueLayout.JAVA_LONG))) .orElse(null); + // i64 os_create_local_store(*const u8 path, i64 path_len) -> Box> ptr + OS_CREATE_LOCAL_STORE = linker.downcallHandle( + lib.find("os_create_local_store").orElseThrow(), + FunctionDescriptor.of(ValueLayout.JAVA_LONG, ValueLayout.ADDRESS, ValueLayout.JAVA_LONG) + ); + // i64 os_destroy_store(i64 handle) + OS_DESTROY_STORE = linker.downcallHandle( + lib.find("os_destroy_store").orElseThrow(), + FunctionDescriptor.of(ValueLayout.JAVA_LONG, ValueLayout.JAVA_LONG) + ); + } private TieredStorageBridge() {} + /** + * Create a shard-scoped local filesystem ObjectStore rooted at {@code rootPath}, returned as a + * {@code Box>} pointer. Object paths are then plain filenames relative to + * {@code rootPath}. The engine owns this handle and frees it with {@link #destroyStore(long)}. + * + * @param rootPath absolute directory the store is rooted at (must already exist) + * @return native ObjectStore handle ({@code > 0}) + */ + public static long createLocalStore(String rootPath) { + byte[] pathBytes = rootPath.getBytes(java.nio.charset.StandardCharsets.UTF_8); + try (Arena arena = Arena.ofConfined()) { + MemorySegment seg = arena.allocateFrom(ValueLayout.JAVA_BYTE, pathBytes); + return NativeLibraryLoader.checkResult( + (long) OS_CREATE_LOCAL_STORE.invokeExact(seg, (long) pathBytes.length) + ); + } catch (Throwable t) { + throw new RuntimeException("Failed to create local ObjectStore at: " + rootPath, t); + } + } + + /** + * Destroy an ObjectStore handle created by {@link #createLocalStore(String)}, decrementing the + * Arc strong count. No-op for a non-positive handle. + */ + public static void destroyStore(long handle) { + if (handle <= 0) return; + try { + NativeLibraryLoader.checkResult((long) OS_DESTROY_STORE.invokeExact(handle)); + } catch (Throwable t) { + throw new RuntimeException("Failed to destroy ObjectStore (ptr=" + handle + ")", t); + } + } + /** * Create a TieredObjectStore with optional local store, remote store, and block cache. * diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/vsr/VSRManager.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/vsr/VSRManager.java index 9b9185d031210..18132c937c3df 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/vsr/VSRManager.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/vsr/VSRManager.java @@ -93,6 +93,26 @@ public VSRManager( this(fileName, indexSettings, schema, bufferPool, maxRowsPerVSR, threadPool, true, writerGeneration, stats); } + /** + * Creates a new VSRManager (async writes) bound to a shard-scoped ObjectStore. + * + * @param storePtr shard ObjectStore handle ({@code os_create_local_store} pointer), or 0 for + * the legacy local-file path + */ + public VSRManager( + String fileName, + IndexSettings indexSettings, + Schema schema, + ArrowBufferPool bufferPool, + int maxRowsPerVSR, + ThreadPool threadPool, + long writerGeneration, + ParquetShardStatsTracker stats, + long storePtr + ) { + this(fileName, indexSettings, schema, bufferPool, maxRowsPerVSR, threadPool, true, writerGeneration, stats, storePtr); + } + /** * Creates a new VSRManager with asynchronous background writes and no stats collection. */ @@ -168,6 +188,27 @@ public VSRManager( boolean runAsync, long writerGeneration, ParquetShardStatsTracker stats + ) { + this(fileName, indexSettings, schema, bufferPool, maxRowsPerVSR, threadPool, runAsync, writerGeneration, stats, 0L); + } + + /** + * Creates a new VSRManager bound to a shard-scoped ObjectStore. + * + * @param storePtr shard ObjectStore handle ({@code os_create_local_store} pointer), or 0 for + * the legacy local-file path + */ + public VSRManager( + String fileName, + IndexSettings indexSettings, + Schema schema, + ArrowBufferPool bufferPool, + int maxRowsPerVSR, + ThreadPool threadPool, + boolean runAsync, + long writerGeneration, + ParquetShardStatsTracker stats, + long storePtr ) { this.fileName = fileName; this.indexSettings = indexSettings; @@ -177,7 +218,7 @@ public VSRManager( this.threadPool = threadPool; this.vsrRotationThread = runAsync ? ParquetDataFormatPlugin.PARQUET_THREAD_POOL_NAME : ThreadPool.Names.SAME; this.managedVSR.set(vsrPool.getActiveVSR()); - this.writer = new NativeParquetWriter(fileName, stats); + this.writer = new NativeParquetWriter(fileName, stats, storePtr); } /** diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/writer/ParquetWriter.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/writer/ParquetWriter.java index c7ad13da8ffa4..693b2d473298a 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/writer/ParquetWriter.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/writer/ParquetWriter.java @@ -93,7 +93,8 @@ public ParquetWriter( ThreadPool threadPool, FormatChecksumStrategy checksumStrategy, ParquetShardStatsTracker stats, - Consumer onClose + Consumer onClose, + long storePtr ) { this.file = file; this.writerGeneration = writerGeneration; @@ -111,7 +112,8 @@ public ParquetWriter( ParquetSettings.MAX_ROWS_PER_VSR.get(indexSettings.getSettings()), threadPool, writerGeneration, - stats + stats, + storePtr ); } @@ -142,7 +144,8 @@ public ParquetWriter( threadPool, checksumStrategy, new ParquetShardStatsTracker(), - null + null, + 0L ); } diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/Cargo.toml b/sandbox/plugins/parquet-data-format/src/main/rust/Cargo.toml index 3834490ca6c35..7849321d26269 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/Cargo.toml +++ b/sandbox/plugins/parquet-data-format/src/main/rust/Cargo.toml @@ -14,7 +14,9 @@ crate-type = ["rlib"] [dependencies] arrow = { workspace = true } -parquet = { workspace = true } +# `object_store` feature pulls in parquet's async writer stack (AsyncArrowWriter + +# ParquetObjectWriter), which streams the Parquet output directly into an ObjectStore. +parquet = { workspace = true, features = ["object_store"] } arrow-ipc = { workspace = true } lazy_static = { workspace = true } dashmap = { workspace = true } @@ -24,6 +26,11 @@ rayon = { workspace = true } tokio = { workspace = true } crc32fast = { workspace = true } serde_json = { workspace = true } +object_store = { workspace = true } +bytes = { workspace = true } +futures = { workspace = true } +# SyncIoBridge (io-util) bridges the sync merge/IPC writers onto object_store's async BufWriter. +tokio-util = { workspace = true, features = ["io-util"] } [dev-dependencies] opensearch-parquet-format = { path = ".", features = ["test-utils"] } diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/crc_writer.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/crc_writer.rs index 0be16dd89377e..11ccce6617338 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/crc_writer.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/crc_writer.rs @@ -19,6 +19,19 @@ impl CrcHandle { pub fn crc32(&self) -> u32 { self.hasher.lock().unwrap().clone().finalize() } + + /// Create a fresh handle together with the shared hasher it reads from, for writers that + /// update the CRC out-of-band (e.g. the async `ObjectStore` sink, which cannot wrap a + /// synchronous [`Write`]). The caller feeds the returned hasher; the handle reads the CRC. + pub fn new_shared() -> (CrcHandle, Arc>) { + let hasher = Arc::new(Mutex::new(crc32fast::Hasher::new())); + ( + CrcHandle { + hasher: hasher.clone(), + }, + hasher, + ) + } } /// A writer wrapper that computes CRC32 incrementally on every write. diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/ffm.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/ffm.rs index db92a137864cc..688d5285dc600 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/ffm.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/ffm.rs @@ -83,6 +83,7 @@ pub unsafe extern "C" fn parquet_create_writer( nulls_first_vals: *const i64, nulls_first_count: i64, writer_generation: i64, + store_handle: i64, ) -> i64 { let filename = str_from_raw(file_ptr, file_len) .map_err(|e| format!("parquet_create_writer file: {}", e))? @@ -103,6 +104,7 @@ pub unsafe extern "C" fn parquet_create_writer( reverse_sorts, nulls_first, writer_generation, + store_handle, ) .map(|ptr| ptr as i64) .map_err(|e| e.to_string()) @@ -656,6 +658,7 @@ pub unsafe extern "C" fn parquet_merge_files( out_flush_and_sort_chunk_count: *mut i64, out_flush_and_sort_chunk_time_millis: *mut i64, out_row_id_mapping_max: *mut i64, + store_handle: i64, ) -> i64 { let input_files = str_array_from_raw(input_ptrs, input_lens, input_count) .map_err(|e| format!("parquet_merge_files inputs: {}", e))?; @@ -683,12 +686,21 @@ pub unsafe extern "C" fn parquet_merge_files( } }; + // Clone the engine-owned store Arc (never consume the box), or None for local segment merges + // and native unit tests. When set, `input_files`/`output_path` are object paths (basenames). + let store: Option> = if store_handle > 0 { + Some((*(store_handle as *const std::sync::Arc)).clone()) + } else { + None + }; + let result = if sort_cols.is_empty() { merge::merge_unsorted( &input_files, output_path, index_name, output_writer_generation, + store, ) } else { merge::merge_sorted( @@ -699,6 +711,7 @@ pub unsafe extern "C" fn parquet_merge_files( &reverse_flags, &nulls_first_flags, output_writer_generation, + store, ) } .map_err(|e| format!("{}", e))?; diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/lib.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/lib.rs index effface637f64..131883a7f4327 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/lib.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/lib.rs @@ -19,6 +19,7 @@ pub mod memory; pub mod merge; pub mod native_settings; pub mod rate_limited_writer; +pub mod store_io; pub mod writer; pub mod writer_properties_builder; diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/context.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/context.rs index 03d9dec407be4..0ecfaaf56908f 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/context.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/context.rs @@ -24,9 +24,12 @@ use crate::writer_properties_builder::WriterPropertiesBuilder; use crate::{log_debug, log_error, SETTINGS_STORE}; use native_bridge_common::memory_pool::MemoryReservation; +use object_store::ObjectStore; use super::error::{MergeError, MergeResult}; -use super::io_task::{get_merge_pool, spawn_io_task, IoCommand, RATE_LIMIT_MB_PER_SEC}; +use super::io_task::{ + get_merge_pool, spawn_io_task, IoCommand, MergeSink, ShutdownCell, RATE_LIMIT_MB_PER_SEC, +}; use super::schema::{append_row_id, build_parquet_root_schema, ROW_ID_COLUMN_NAME}; /// Owns all shared state for a merge operation: schemas, writer factory, @@ -65,13 +68,18 @@ impl MergeContext { io_threads: Option, output_writer_generation: i64, reservation: MemoryReservation, + store: Option>, ) -> MergeResult { - if let Some(parent) = Path::new(output_path).parent() { - if !parent.exists() { - return Err(MergeError::Logic(format!( - "Output directory '{}' does not exist.", - parent.display() - ))); + // For the local path, validate the parent directory exists. For the store path, + // `output_path` is an object key (no local parent), so skip this check. + if store.is_none() { + if let Some(parent) = Path::new(output_path).parent() { + if !parent.exists() { + return Err(MergeError::Logic(format!( + "Output directory '{}' does not exist.", + parent.display() + ))); + } } } @@ -97,9 +105,26 @@ impl MergeContext { let parquet_root = build_parquet_root_schema(parquet_descriptors)?; - let output_file = File::create(output_path)?; + // Build the output sink: store-backed multipart upload, or a local file. + let (sink, shutdown_cell): (MergeSink, Option) = match &store { + Some(s) => { + let cell: ShutdownCell = Arc::new(std::sync::Mutex::new(None)); + let bridge = crate::store_io::store_sync_writer( + s.clone(), + object_store::path::Path::from(output_path), + ); + ( + MergeSink::Store { + bridge: Some(bridge), + result: cell.clone(), + }, + Some(cell), + ) + } + None => (MergeSink::Local(File::create(output_path)?), None), + }; let throttled_writer = - RateLimitedWriter::new(output_file, RATE_LIMIT_MB_PER_SEC).map_err(MergeError::Io)?; + RateLimitedWriter::new(sink, RATE_LIMIT_MB_PER_SEC).map_err(MergeError::Io)?; let (crc_writer, crc_handle) = CrcWriter::new(throttled_writer); @@ -120,7 +145,7 @@ impl MergeContext { let writer = SerializedFileWriter::new(crc_writer, parquet_root, writer_props)?; let rg_writer_factory = ArrowRowGroupWriterFactory::new(&writer, output_schema.clone()); - let io_tx = spawn_io_task(writer, crc_handle, io_threads); + let io_tx = spawn_io_task(writer, crc_handle, io_threads, shutdown_cell); let col_writers = rg_writer_factory.create_column_writers(0)?; diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/cursor.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/cursor.rs index bcb881b534378..bcbc39af92d64 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/cursor.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/cursor.rs @@ -6,17 +6,18 @@ * compatible open source license. */ -use std::fs::File; use std::sync::{Arc, Mutex}; use arrow::array::RecordBatch; use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema}; -use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use object_store::ObjectStore; +use parquet::arrow::ProjectionMask; use parquet::schema::types::SchemaDescriptor; use super::error::{MergeError, MergeResult}; use super::heap::{get_sort_values, SortKey}; use super::io_task::get_merge_pool; +use super::reader::{build_source, read_source_meta, BatchSource}; use super::schema::projection_indices_excluding_row_id; use native_bridge_common::memory_pool::MemoryReservation; @@ -28,13 +29,13 @@ use native_bridge_common::memory_pool::MemoryReservation; /// uses two readers: a sort-only reader for the merge heap and a data reader /// loaded on demand. Otherwise uses a single all-column reader. pub struct FileCursor { - sort_reader: Arc>, + sort_reader: Arc>, sort_prefetch_rx: std::sync::mpsc::Receiver>>, sort_prefetch_tx: std::sync::mpsc::SyncSender>>, sort_prefetch_pending: bool, pub sort_batch: Option, - data_reader: Option, + data_reader: Option, data_batch: Option, sort_batch_index: usize, data_batch_index: usize, @@ -58,17 +59,11 @@ impl FileCursor { batch_size: usize, deferred_threshold: usize, reservation: &mut MemoryReservation, + store: &Option>, ) -> MergeResult<(Self, Arc, SchemaDescriptor, i64, usize)> { - // Open file and read metadata - let file = File::open(path)?; - let builder = ParquetRecordBatchReaderBuilder::try_new(file)?; - let schema = builder.schema().clone(); - let writer_generation = crate::writer_properties_builder::read_writer_generation( - builder.metadata().file_metadata(), - file_id, - ); - let total_row_count = builder.metadata().file_metadata().num_rows() as usize; - let parquet_schema_descr = builder.parquet_schema().clone(); + // Read metadata (local file or store object) + let (schema, parquet_schema_descr, writer_generation, total_row_count) = + read_source_meta(store, path, file_id)?; // Resolve sort column types let sort_col_types: Vec = sort_columns @@ -112,39 +107,22 @@ impl FileCursor { let data_projection_indices = projection_indices_excluding_row_id(&schema); // Build sort reader (sort-only in deferred, all-columns in eager) - let file1 = File::open(path)?; - let builder1 = ParquetRecordBatchReaderBuilder::try_new(file1)?; let sort_projection = if deferred { let sort_indices: Vec = sort_columns .iter() .filter_map(|c| schema.fields().iter().position(|f| f.name() == c.as_str())) .collect(); - parquet::arrow::ProjectionMask::roots(builder1.parquet_schema(), sort_indices) + ProjectionMask::roots(&parquet_schema_descr, sort_indices) } else { - parquet::arrow::ProjectionMask::roots( - builder1.parquet_schema(), - data_projection_indices.clone(), - ) + ProjectionMask::roots(&parquet_schema_descr, data_projection_indices.clone()) }; - let mut sort_reader = builder1 - .with_batch_size(batch_size) - .with_projection(sort_projection) - .build()?; + let mut sort_reader = build_source(store, path, batch_size, sort_projection)?; // Build data reader (only in deferred mode) let data_reader = if deferred { - let file2 = File::open(path)?; - let builder2 = ParquetRecordBatchReaderBuilder::try_new(file2)?; - let data_proj = parquet::arrow::ProjectionMask::roots( - builder2.parquet_schema(), - data_projection_indices.clone(), - ); - Some( - builder2 - .with_batch_size(batch_size) - .with_projection(data_proj) - .build()?, - ) + let data_proj = + ProjectionMask::roots(&parquet_schema_descr, data_projection_indices.clone()); + Some(build_source(store, path, batch_size, data_proj)?) } else { None }; @@ -158,9 +136,9 @@ impl FileCursor { )); // Read first sort batch - let first_sort_batch = match sort_reader.next() { + let first_sort_batch = match sort_reader.next_batch() { Some(Ok(b)) if b.num_rows() > 0 => b, - Some(Err(e)) => return Err(e.into()), + Some(Err(e)) => return Err(e), _ => { return Err(MergeError::Logic(format!( "File '{}' (cursor {}) yielded no rows", @@ -235,9 +213,9 @@ impl FileCursor { let tx = self.sort_prefetch_tx.clone(); get_merge_pool(None).spawn(move || { let mut reader = reader.lock().unwrap(); - let result = match reader.next() { + let result = match reader.next_batch() { Some(Ok(batch)) if batch.num_rows() > 0 => Some(Ok(batch)), - Some(Err(e)) => Some(Err(MergeError::Arrow(e))), + Some(Err(e)) => Some(Err(e)), _ => None, }; let _ = tx.send(result); @@ -330,7 +308,7 @@ impl FileCursor { self.data_batch = None; while self.data_batch_index <= self.sort_batch_index { - match reader.next() { + match reader.next_batch() { Some(Ok(batch)) => { if batch.num_rows() == 0 { return Err(MergeError::Logic(format!( @@ -359,7 +337,7 @@ impl FileCursor { // Skipped batch — discard self.data_batch_index += 1; } - Some(Err(e)) => return Err(e.into()), + Some(Err(e)) => return Err(e), None => { return Err(MergeError::Logic(format!( "Data reader exhausted at position {}, needed sort_batch_index={}", diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/io_task.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/io_task.rs index e2383aba53237..23bfd3d2a826c 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/io_task.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/io_task.rs @@ -7,7 +7,8 @@ */ use std::fs::File; -use std::sync::OnceLock; +use std::io::Write; +use std::sync::{Arc, Mutex, OnceLock}; use parquet::file::metadata::ParquetMetaData; use parquet::file::writer::SerializedFileWriter; @@ -17,6 +18,7 @@ use rayon::ThreadPool; use crate::crc_writer::CrcWriter; use crate::log_error; use crate::rate_limited_writer::RateLimitedWriter; +use crate::store_io::StoreSyncWriter; use native_bridge_common::log_info; use tokio::runtime::Runtime; use tokio::sync::{mpsc as tokio_mpsc, oneshot}; @@ -80,8 +82,61 @@ fn get_io_runtime(num_threads: Option) -> &'static Runtime { // IO task protocol // ============================================================================= -/// Writer type used by the IO task: CRC → rate-limit → file. -pub type MergeWriter = CrcWriter>; +/// Shared cell used to surface the store multipart-completion (`shutdown`) result out of +/// [`MergeSink`]'s `Drop`, so the IO task can propagate an upload failure instead of silently +/// losing it. +pub type ShutdownCell = Arc>>>; + +/// Final output sink for the merge: a local file, or a store-backed multipart upload. +/// +/// The store variant streams through `object_store`'s `BufWriter` (bridged to sync I/O). The +/// multipart upload is only finalized on `shutdown()`, which is invoked from `Drop` — this fires +/// when `SerializedFileWriter::close()` consumes and drops the writer chain. Any shutdown error is +/// recorded into the shared `result` cell for the IO task to propagate. +pub enum MergeSink { + Local(File), + Store { + bridge: Option, + result: ShutdownCell, + }, +} + +impl Write for MergeSink { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + match self { + MergeSink::Local(f) => f.write(buf), + MergeSink::Store { bridge, .. } => bridge + .as_mut() + .expect("MergeSink store writer used after finalize") + .write(buf), + } + } + fn flush(&mut self) -> std::io::Result<()> { + match self { + MergeSink::Local(f) => f.flush(), + MergeSink::Store { bridge, .. } => bridge + .as_mut() + .expect("MergeSink store writer used after finalize") + .flush(), + } + } +} + +impl Drop for MergeSink { + fn drop(&mut self) { + if let MergeSink::Store { bridge, result } = self { + if let Some(mut b) = bridge.take() { + // Complete the multipart upload (for LocalFileSystem this renames the temp into + // place). Record the result so the IO task can surface a failure. + let r = b.shutdown(); + *result.lock().unwrap() = Some(r); + } + } + } +} + +/// Writer type used by the IO task: CRC → rate-limit → sink (local file or store multipart). +pub type MergeWriter = CrcWriter>; /// Commands sent from the merge loop to the background IO task. pub enum IoCommand { @@ -109,6 +164,7 @@ pub fn spawn_io_task( writer: SerializedFileWriter, crc_handle: crate::crc_writer::CrcHandle, io_threads: Option, + store_shutdown: Option, ) -> tokio_mpsc::Sender { let (tx, mut rx) = tokio_mpsc::channel::(IO_CHANNEL_BUFFER); @@ -170,8 +226,16 @@ pub fn spawn_io_task( let w = writer.take().unwrap(); let crc = crc_handle.clone(); + let store_shutdown = store_shutdown.clone(); let result = tokio::task::spawn_blocking(move || { let metadata = w.close().map_err(MergeError::from)?; + // `w` is dropped by close(); for the store sink that ran `shutdown()` to + // finalize the multipart upload — propagate any failure. + if let Some(cell) = &store_shutdown { + if let Some(res) = cell.lock().unwrap().take() { + res.map_err(MergeError::Io)?; + } + } let crc32 = crc.crc32(); log_info!( "[RUST] IO task close: version={}, num_rows={}, created_by={:?}, crc32={:#010x}", diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/mod.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/mod.rs index 545337729e297..a987ad632498a 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/mod.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/mod.rs @@ -12,6 +12,7 @@ pub mod error; pub mod heap; pub mod io_task; pub mod metrics; +mod reader; pub mod schema; mod sorted; mod unsorted; diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/reader.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/reader.rs new file mode 100644 index 0000000000000..b668958c22234 --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/reader.rs @@ -0,0 +1,130 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +//! Shared input-reader abstraction for the merge paths (sorted `FileCursor` and unsorted merge). +//! +//! A merge input is either a local Parquet file (segment merges when no store is configured, and +//! native unit tests) or a store object read via the async `ParquetObjectReader`. The async stream +//! is driven to completion with `block_on` on the store runtime — safe here because merges run on +//! rayon / finalize / `spawn_blocking` threads, never Tokio workers. + +use std::fs::File; +use std::sync::Arc; + +use arrow::array::RecordBatch; +use arrow::datatypes::Schema as ArrowSchema; +use futures::StreamExt; +use object_store::path::Path as ObjectPath; +use object_store::ObjectStore; +use parquet::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder}; +use parquet::arrow::async_reader::{ + ParquetObjectReader, ParquetRecordBatchStream, ParquetRecordBatchStreamBuilder, +}; +use parquet::arrow::ProjectionMask; +use parquet::schema::types::SchemaDescriptor; + +use super::error::{MergeError, MergeResult}; +use crate::store_io::os_store_runtime; + +/// A source of `RecordBatch`es for one merge input — a synchronous local Parquet reader, or an +/// async store-backed Parquet stream driven to completion via `block_on` on the store runtime. +pub(crate) enum BatchSource { + Local(ParquetRecordBatchReader), + Store(ParquetRecordBatchStream), +} + +impl BatchSource { + /// Pull the next batch, or `None` at end of input. + pub(crate) fn next_batch(&mut self) -> Option> { + match self { + BatchSource::Local(reader) => match reader.next() { + Some(Ok(b)) => Some(Ok(b)), + Some(Err(e)) => Some(Err(MergeError::Arrow(e))), + None => None, + }, + BatchSource::Store(stream) => os_store_runtime().block_on(async { + match stream.next().await { + Some(Ok(b)) => Some(Ok(b)), + Some(Err(e)) => Some(Err(MergeError::from(e))), + None => None, + } + }), + } + } +} + +/// Read just the metadata/schema of an input (local or store) to drive projection/mode decisions +/// before building the batch readers. +pub(crate) fn read_source_meta( + store: &Option>, + path: &str, + file_id: usize, +) -> MergeResult<(Arc, SchemaDescriptor, i64, usize)> { + match store { + Some(s) => { + let reader = ParquetObjectReader::new(s.clone(), ObjectPath::from(path)); + os_store_runtime() + .block_on(async move { + let builder = ParquetRecordBatchStreamBuilder::new(reader).await?; + let schema = builder.schema().clone(); + let wg = crate::writer_properties_builder::read_writer_generation( + builder.metadata().file_metadata(), + file_id, + ); + let trc = builder.metadata().file_metadata().num_rows() as usize; + let descr = builder.parquet_schema().clone(); + Ok::<_, parquet::errors::ParquetError>((schema, descr, wg, trc)) + }) + .map_err(MergeError::from) + } + None => { + let file = File::open(path)?; + let builder = ParquetRecordBatchReaderBuilder::try_new(file)?; + let schema = builder.schema().clone(); + let wg = crate::writer_properties_builder::read_writer_generation( + builder.metadata().file_metadata(), + file_id, + ); + let trc = builder.metadata().file_metadata().num_rows() as usize; + let descr = builder.parquet_schema().clone(); + Ok((schema, descr, wg, trc)) + } + } +} + +/// Build a batch source (local reader or store stream) for `path` with the given projection. +pub(crate) fn build_source( + store: &Option>, + path: &str, + batch_size: usize, + projection: ProjectionMask, +) -> MergeResult { + match store { + Some(s) => { + let reader = ParquetObjectReader::new(s.clone(), ObjectPath::from(path)); + let stream = os_store_runtime() + .block_on(async move { + ParquetRecordBatchStreamBuilder::new(reader) + .await? + .with_batch_size(batch_size) + .with_projection(projection) + .build() + }) + .map_err(MergeError::from)?; + Ok(BatchSource::Store(stream)) + } + None => { + let file = File::open(path)?; + let reader = ParquetRecordBatchReaderBuilder::try_new(file)? + .with_batch_size(batch_size) + .with_projection(projection) + .build()?; + Ok(BatchSource::Local(reader)) + } + } +} diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/sorted.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/sorted.rs index 7e1dbca8586e0..4ee5eafd91dd6 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/sorted.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/sorted.rs @@ -23,8 +23,12 @@ use super::schema::ColumnMapping; use crate::memory::merge_pool; use native_bridge_common::memory_pool::{MemoryReservation, PoolBehavior}; +use object_store::ObjectStore; /// Performs a streaming k-way merge with an explicit sort direction per column. +/// +/// `store` selects the I/O backend: `Some` reads inputs from / writes output to the object store +/// (the sorted-refresh write path); `None` uses the local filesystem (segment merges, unit tests). pub fn merge_sorted( input_files: &[String], output_path: &str, @@ -33,6 +37,7 @@ pub fn merge_sorted( reverse_sorts: &[bool], nulls_first: &[bool], output_writer_generation: i64, + store: Option>, ) -> super::MergeResult { let mut reservation = MemoryReservation::new(merge_pool(), "merge_sorted", PoolBehavior::Reject); @@ -45,6 +50,7 @@ pub fn merge_sorted( nulls_first, output_writer_generation, &mut reservation, + store, ) } @@ -58,6 +64,7 @@ pub fn merge_sorted_with_pool( nulls_first: &[bool], output_writer_generation: i64, reservation: &mut MemoryReservation, + store: Option>, ) -> super::MergeResult { let config = crate::writer::SETTINGS_STORE .get(index_name) @@ -118,6 +125,7 @@ pub fn merge_sorted_with_pool( batch_size, deferred_threshold, reservation, + &store, )?; cursors.push(cursor); arrow_schemas.push(projected_schema.as_ref().clone()); @@ -140,6 +148,7 @@ pub fn merge_sorted_with_pool( io_threads, output_writer_generation, ctx_reservation, + store.clone(), )?; // Precompute column mappings per cursor (avoids per-batch name lookups) diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/unsorted.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/unsorted.rs index 8f1ef0260666d..c3627de9f8f66 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/unsorted.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/unsorted.rs @@ -6,17 +6,18 @@ * compatible open source license. */ -use std::fs::File; +use std::sync::Arc; -use arrow::array::RecordBatchReader; use arrow::datatypes::Schema as ArrowSchema; -use parquet::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder}; +use object_store::ObjectStore; +use parquet::arrow::ProjectionMask; use parquet::schema::types::SchemaDescriptor; use crate::log_debug; use super::context::MergeContext; use super::error::MergeResult; +use super::reader::{build_source, read_source_meta, BatchSource}; use super::schema::{projection_indices_excluding_row_id, ColumnMapping}; use crate::memory::merge_pool; @@ -29,6 +30,7 @@ pub fn merge_unsorted( output_path: &str, index_name: &str, output_writer_generation: i64, + store: Option>, ) -> MergeResult { let mut reservation = MemoryReservation::new(merge_pool(), "merge_unsorted", PoolBehavior::Reject); @@ -38,6 +40,7 @@ pub fn merge_unsorted( index_name, output_writer_generation, &mut reservation, + store, ) } @@ -48,6 +51,7 @@ pub fn merge_unsorted_with_pool( index_name: &str, output_writer_generation: i64, reservation: &mut MemoryReservation, + store: Option>, ) -> MergeResult { let config = crate::writer::SETTINGS_STORE .get(index_name) @@ -63,33 +67,29 @@ pub fn merge_unsorted_with_pool( output_path ); - // Single pass: collect schemas and build readers. + // Single pass: collect schemas and build readers (local files or store objects). let mut arrow_schemas: Vec = Vec::with_capacity(input_files.len()); let mut parquet_descriptors: Vec = Vec::with_capacity(input_files.len()); - let mut readers: Vec = Vec::with_capacity(input_files.len()); + let mut readers: Vec = Vec::with_capacity(input_files.len()); let mut file_row_counts: Vec = Vec::with_capacity(input_files.len()); let mut file_generations: Vec = Vec::with_capacity(input_files.len()); for (file_idx, path) in input_files.iter().enumerate() { - let file = File::open(path)?; - let builder = ParquetRecordBatchReaderBuilder::try_new(file)?; - let schema = builder.schema().clone(); - let parquet_descr = builder.parquet_schema().clone(); - let num_rows = builder.metadata().file_metadata().num_rows() as usize; - let generation = crate::writer_properties_builder::read_writer_generation( - builder.metadata().file_metadata(), - file_idx, - ); + let (schema, parquet_descr, generation, num_rows) = + read_source_meta(&store, path, file_idx)?; let projection_indices = projection_indices_excluding_row_id(&schema); - let projection = parquet::arrow::ProjectionMask::roots(&parquet_descr, projection_indices); - let reader = builder - .with_batch_size(batch_size) - .with_projection(projection) - .build()?; - - // The reader's schema is the projected schema (__row_id__ excluded). - arrow_schemas.push(reader.schema().as_ref().clone()); + // Projected schema (__row_id__ excluded) — matches what the reader yields. + let projected_schema = ArrowSchema::new( + projection_indices + .iter() + .map(|&i| schema.field(i).clone()) + .collect::>(), + ); + let projection = ProjectionMask::roots(&parquet_descr, projection_indices); + let reader = build_source(&store, path, batch_size, projection)?; + + arrow_schemas.push(projected_schema); parquet_descriptors.push(parquet_descr); readers.push(reader); file_row_counts.push(num_rows); @@ -107,6 +107,7 @@ pub fn merge_unsorted_with_pool( io_threads, output_writer_generation, ctx_reservation, + store.clone(), )?; // Precompute column mappings per reader @@ -131,7 +132,7 @@ pub fn merge_unsorted_with_pool( let mut new_row_id: i64 = 0; // Iterate readers for data. - for (file_idx, reader) in readers.into_iter().enumerate() { + for (file_idx, mut reader) in readers.into_iter().enumerate() { log_debug!( "[RUST] Unsorted merge: processing file {} of {}", file_idx + 1, @@ -144,7 +145,7 @@ pub fn merge_unsorted_with_pool( let col_mapping = &col_mappings[file_idx]; let mut batch_tracked: usize = 0; - for batch_result in reader { + while let Some(batch_result) = reader.next_batch() { let batch = batch_result?; let num_rows = batch.num_rows(); let batch_bytes = batch.get_array_memory_size(); diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/store_io.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/store_io.rs new file mode 100644 index 0000000000000..7bcc6f2ca06c0 --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/store_io.rs @@ -0,0 +1,72 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +//! Shared helpers for driving the shard-scoped `ObjectStore` from the otherwise-synchronous +//! Parquet write/merge paths. +//! +//! The Parquet writer, the sorting chunked writer, and the k-way merge are all synchronous and +//! serialized per-writer by Java (or run on rayon/finalize threads). None of those threads are +//! Tokio runtime workers, so it is safe to drive the async `ObjectStore` operations to completion +//! with `block_on` on the dedicated runtime here. `object_store` I/O (multipart PUT parts, ranged +//! GETs) runs on this runtime's worker threads while the caller blocks. + +use bytes::Bytes; +use object_store::buffered::BufWriter as ObjectBufWriter; +use object_store::path::Path as ObjectPath; +use object_store::{ObjectStore, ObjectStoreExt}; +use std::sync::{Arc, OnceLock}; +use tokio_util::io::SyncIoBridge; + +static OS_STORE_RUNTIME: OnceLock = OnceLock::new(); + +/// Dedicated multi-thread runtime for all shard `ObjectStore` I/O on the write path +/// (streaming Parquet/IPC uploads, chunk reads during the sort-finalize merge, deletes/renames). +pub fn os_store_runtime() -> &'static tokio::runtime::Runtime { + OS_STORE_RUNTIME.get_or_init(|| { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .thread_name("parquet-os-store") + .enable_all() + .build() + .expect("Failed to build parquet ObjectStore runtime") + }) +} + +/// A synchronous [`std::io::Write`] that streams into an `ObjectStore` multipart upload. +/// +/// Backed by `object_store::buffered::BufWriter` (which chunks writes into multipart parts) bridged +/// to sync I/O via [`SyncIoBridge`]. The upload is **only finalized** when [`SyncIoBridge::shutdown`] +/// is called — callers MUST call `shutdown()` after finishing their writes (drop alone will NOT +/// complete the multipart upload). +pub type StoreSyncWriter = SyncIoBridge; + +/// Build a sync streaming writer targeting `path` in `store`. +pub fn store_sync_writer(store: Arc, path: ObjectPath) -> StoreSyncWriter { + let buf = ObjectBufWriter::new(store, path); + SyncIoBridge::new_with_handle(buf, os_store_runtime().handle().clone()) +} + +/// Fully read an object into memory (used to read back an IPC staging object for sort/merge, which +/// the local path also loads fully into memory at flush time — so this adds no extra peak). +pub fn read_object(store: &Arc, path: &ObjectPath) -> object_store::Result { + os_store_runtime().block_on(async move { store.get(path).await?.bytes().await }) +} + +/// Delete an object (best-effort semantics are the caller's; this surfaces the error). +pub fn delete_object(store: &Arc, path: &ObjectPath) -> object_store::Result<()> { + os_store_runtime().block_on(async move { store.delete(path).await }) +} + +/// Rename `from` to `to` within the store (metadata-only for `LocalFileSystem`). +pub fn rename_object( + store: &Arc, + from: &ObjectPath, + to: &ObjectPath, +) -> object_store::Result<()> { + os_store_runtime().block_on(async move { store.rename(from, to).await }) +} diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/test_utils.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/test_utils.rs index e368f43c582f5..4d1a309508693 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/test_utils.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/test_utils.rs @@ -47,6 +47,7 @@ pub fn test_create_writer( reverse_sorts, nulls_first, writer_generation, + 0, // store_handle=0 -> legacy local-file path (no ObjectStore in native unit tests) )?; TEST_HANDLES.lock().unwrap().insert(filename, handle as i64); Ok(()) diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/writer.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/writer.rs index 32e5859b83367..79cb076ab9288 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/writer.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/writer.rs @@ -20,7 +20,7 @@ use std::fs::File; use std::path::Path; use std::sync::{Arc, Mutex}; -use crate::crc_writer::CrcWriter; +use crate::crc_writer::{CrcHandle, CrcWriter}; use crate::memory::write_pool; use crate::merge::{merge_sorted_with_pool, schema::ROW_ID_COLUMN_NAME}; use crate::native_settings::NativeSettings; @@ -28,6 +28,93 @@ use crate::writer_properties_builder::WriterPropertiesBuilder; use crate::{log_debug, log_error, log_info}; use native_bridge_common::memory_pool::{MemoryReservation, PoolBehavior}; +use object_store::ObjectStore; +use parquet::arrow::async_writer::{AsyncFileWriter, ParquetObjectWriter}; +use parquet::arrow::AsyncArrowWriter; +use bytes::Bytes; +use futures::future::BoxFuture; +use crate::store_io::os_store_runtime; + +/// An [`AsyncFileWriter`] that tees the Parquet byte stream into a CRC32 hasher before handing it +/// to the underlying [`ParquetObjectWriter`] sink. This is the async counterpart to +/// [`crate::crc_writer::CrcWriter`]: the same whole-file CRC32 that the local path computes, but +/// over the exact bytes uploaded to the `ObjectStore` (multipart, streamed — never staged to a +/// local temp file). The CRC is read out-of-band via a [`CrcHandle`] sharing `hasher`. +struct CrcObjectWriter { + inner: ParquetObjectWriter, + hasher: Arc>, +} + +impl AsyncFileWriter for CrcObjectWriter { + fn write(&mut self, bs: Bytes) -> BoxFuture<'_, parquet::errors::Result<()>> { + // Hash synchronously (in write order) before delegating the async upload. + self.hasher.lock().unwrap().update(&bs); + self.inner.write(bs) + } + + fn complete(&mut self) -> BoxFuture<'_, parquet::errors::Result<()>> { + self.inner.complete() + } +} + +/// Build a store-backed streaming Parquet sink (`AsyncArrowWriter` over `ParquetObjectWriter`), +/// teeing every uploaded byte into a returned CRC32 handle. Shared by the non-sorted writer and +/// by sorted-chunk writes. +fn new_store_parquet_sink( + store: Arc, + object_path: &str, + schema: Arc, + props: parquet::file::properties::WriterProperties, +) -> Result<(AsyncArrowWriter, CrcHandle), Box> { + let object_writer = + ParquetObjectWriter::new(store, object_store::path::Path::from(object_path)); + let (crc_handle, hasher) = CrcHandle::new_shared(); + let crc_writer = CrcObjectWriter { + inner: object_writer, + hasher, + }; + let writer = AsyncArrowWriter::try_new(crc_writer, schema, Some(props))?; + Ok((writer, crc_handle)) +} + +/// Write a single `RecordBatch` as a complete Parquet object into the store (used for sorted +/// chunks). Streams via `AsyncArrowWriter` — no local temp file — and returns the whole-object +/// CRC32. Mirrors [`NativeParquetWriter::write_final_file`] for the store path. +fn write_batch_to_store( + store: Arc, + object_path: &str, + index_name: &str, + batch: &RecordBatch, + schema: Arc, + writer_generation: Option, +) -> Result> { + let config = SETTINGS_STORE + .get(index_name) + .map(|r| r.clone()) + .unwrap_or_default(); + let props = WriterPropertiesBuilder::build_with_generation(&config, writer_generation, &schema) + .map_err(|e| format!("Invalid encoding/compression config: {}", e))?; + let (mut writer, crc_handle) = new_store_parquet_sink(store, object_path, schema, props)?; + os_store_runtime().block_on(writer.write(batch))?; + os_store_runtime().block_on(writer.close())?; + Ok(crc_handle.crc32()) +} + +/// Collect all non-empty batches from an Arrow IPC file reader (local `File` or in-memory +/// `Cursor` over a store object — both are `Read + Seek`). +fn collect_ipc_batches( + reader: IpcFileReader, +) -> Result, Box> { + let mut batches: Vec = Vec::new(); + for batch_result in reader { + let batch = batch_result?; + if batch.num_rows() > 0 { + batches.push(batch); + } + } + Ok(batches) +} + /// Result from finalizing a writer: Parquet metadata + whole-file CRC32 + optional sort permutation. #[derive(Debug)] pub struct FinalizeResult { @@ -53,6 +140,45 @@ enum WriterVariant { /// sorted Parquet chunk, then starts a new IPC file. At finalize, only /// a k-way merge is needed. Ipc(Arc>), + /// Store-backed Parquet writer — used for the non-sorted path when a shard `ObjectStore` + /// handle is present (the production path). The Parquet output is streamed directly into the + /// `ObjectStore` (multipart upload) via `AsyncArrowWriter`; there is no local temp file and no + /// post-hoc copy. For hot indices the store is a `LocalFileSystem` rooted at the shard dir, so + /// the object lands at `filename`. + ParquetStore(Arc>>), +} + +/// IPC staging sink — local file or a store-backed multipart upload. Arrow IPC batches are written +/// incrementally as they arrive; on flush the object/file is read back fully, sorted, and written +/// out as a sorted Parquet chunk. +enum IpcStaging { + Local(IpcFileWriter), + Store(IpcFileWriter), +} + +impl IpcStaging { + fn write(&mut self, batch: &RecordBatch) -> Result<(), arrow::error::ArrowError> { + match self { + IpcStaging::Local(w) => w.write(batch), + IpcStaging::Store(w) => w.write(batch), + } + } + + /// Write the IPC footer and, for the store sink, `shutdown()` the bridge to finalize the + /// multipart upload (a plain drop would NOT complete it). + fn finish_and_close(self) -> Result<(), Box> { + match self { + IpcStaging::Local(mut w) => { + w.finish()?; + } + IpcStaging::Store(mut w) => { + w.finish()?; + let mut bridge = w.into_inner()?; + bridge.shutdown()?; + } + } + Ok(()) + } } /// Hybrid IPC staging + eager sort writer. @@ -64,7 +190,7 @@ enum WriterVariant { /// - At finalize: flushes remaining IPC data (sort + write), returns sorted /// Parquet chunk paths for k-way merge. struct SortingChunkedWriter { - /// Base path for staging/chunk files. + /// Base path for local staging/chunk files (used when `store` is `None`). base_path: String, /// Arrow schema shared across all chunks. schema: Arc, @@ -79,7 +205,7 @@ struct SortingChunkedWriter { reverse_sorts: Vec, nulls_first: Vec, /// Current IPC writer for staging incoming batches. - current_ipc_writer: Option>, + current_ipc_writer: Option, /// Tracked byte size of the current IPC staging file (approximated from /// the Arrow array memory sizes of batches written so far). current_chunk_bytes: u64, @@ -87,7 +213,7 @@ struct SortingChunkedWriter { current_rows: usize, /// Index of the next chunk (0-based). chunk_idx: usize, - /// Paths of all completed sorted Parquet chunk files. + /// Paths of all completed sorted Parquet chunk files (local paths or store object paths). completed_chunks: Vec, /// Row IDs captured from each sorted chunk (for permutation building). chunk_row_ids: Vec>, @@ -97,6 +223,11 @@ struct SortingChunkedWriter { total_rows: usize, /// Writer generation propagated into Parquet file metadata for each chunk. writer_generation: i64, + /// Shard store, or `None` for the local path. When set, IPC staging and sorted chunks are + /// written to / read from / deleted in the store instead of the local filesystem. + store: Option>, + /// Basename of the final output; IPC + chunk object paths are derived from it. + object_base: String, } impl SortingChunkedWriter { @@ -109,6 +240,8 @@ impl SortingChunkedWriter { reverse_sorts: Vec, nulls_first: Vec, writer_generation: i64, + store: Option>, + object_base: String, ) -> Result> { let mut writer = Self { base_path, @@ -127,23 +260,45 @@ impl SortingChunkedWriter { chunk_crcs: Vec::new(), total_rows: 0, writer_generation, + store, + object_base, }; writer.open_new_ipc()?; Ok(writer) } + /// IPC staging path (local temp file when `store` is `None`). fn ipc_staging_path(&self) -> String { self.base_path.clone() } + /// IPC staging object path within the store root. + fn ipc_object_path(&self) -> String { + format!("{}{}", self.object_base, IPC_STAGING_SUFFIX) + } + fn sorted_chunk_path(&self, idx: usize) -> String { - format!("{}.sorted_chunk_{}.parquet", self.base_path, idx) + if self.store.is_some() { + format!("{}.sorted_chunk_{}.parquet", self.object_base, idx) + } else { + format!("{}.sorted_chunk_{}.parquet", self.base_path, idx) + } } fn open_new_ipc(&mut self) -> Result<(), Box> { - let path = self.ipc_staging_path(); - let file = File::create(&path)?; - let ipc_writer = IpcFileWriter::try_new(file, &self.schema)?; + let ipc_writer = match &self.store { + Some(store) => { + let sink = crate::store_io::store_sync_writer( + store.clone(), + object_store::path::Path::from(self.ipc_object_path()), + ); + IpcStaging::Store(IpcFileWriter::try_new(sink, &self.schema)?) + } + None => { + let file = File::create(self.ipc_staging_path())?; + IpcStaging::Local(IpcFileWriter::try_new(file, &self.schema)?) + } + }; self.current_ipc_writer = Some(ipc_writer); self.current_chunk_bytes = 0; self.current_rows = 0; @@ -233,28 +388,31 @@ impl SortingChunkedWriter { let sort_reserve = self.current_chunk_bytes as usize * 2; reservation.request(sort_reserve)?; - // Close the IPC writer - if let Some(mut writer) = self.current_ipc_writer.take() { - writer.finish()?; + // Close the IPC writer (finalizes the multipart upload for the store sink). + if let Some(writer) = self.current_ipc_writer.take() { + writer.finish_and_close()?; } - let ipc_path = self.ipc_staging_path(); - - // Read back the IPC file (still hot in page cache since we just wrote it) - let file = File::open(&ipc_path)?; - let reader = IpcFileReader::try_new(file, None)?; - let mut batches: Vec = Vec::new(); - for batch_result in reader { - let batch = batch_result?; - if batch.num_rows() > 0 { - batches.push(batch); + // Read the staged IPC back fully, then sort. The local path also loads the whole staging + // file into memory here, so reading a store object into memory adds no extra peak. + let batches: Vec = match &self.store { + Some(store) => { + let data = crate::store_io::read_object( + store, + &object_store::path::Path::from(self.ipc_object_path()), + )?; + collect_ipc_batches(IpcFileReader::try_new(std::io::Cursor::new(data), None)?)? } - } + None => { + let file = File::open(self.ipc_staging_path())?; + collect_ipc_batches(IpcFileReader::try_new(file, None)?)? + } + }; if batches.is_empty() { // Nothing to sort, just reopen reservation.shrink(sort_reserve); - let _ = std::fs::remove_file(&ipc_path); + self.delete_ipc_staging(); self.open_new_ipc()?; return Ok(()); } @@ -297,15 +455,25 @@ impl SortingChunkedWriter { sorted_batch }; - // Write sorted chunk as Parquet + // Write sorted chunk as Parquet (to the store, or a local file when store is None) let chunk_path = self.sorted_chunk_path(self.chunk_idx); - let crc32 = NativeParquetWriter::write_final_file( - &chunk_path, - &self.index_name, - &final_batch, - self.schema.clone(), - Some(self.writer_generation), - )?; + let crc32 = match &self.store { + Some(store) => write_batch_to_store( + store.clone(), + &chunk_path, + &self.index_name, + &final_batch, + self.schema.clone(), + Some(self.writer_generation), + )?, + None => NativeParquetWriter::write_final_file( + &chunk_path, + &self.index_name, + &final_batch, + self.schema.clone(), + Some(self.writer_generation), + )?, + }; self.completed_chunks.push(chunk_path); self.chunk_crcs.push(crc32); @@ -318,12 +486,27 @@ impl SortingChunkedWriter { let row_ids_bytes = self.current_rows * std::mem::size_of::(); reservation.grow(row_ids_bytes); - // Delete the IPC staging file and open a fresh one - let _ = std::fs::remove_file(&ipc_path); + // Delete the IPC staging object/file and open a fresh one + self.delete_ipc_staging(); self.open_new_ipc()?; Ok(()) } + /// Delete the current IPC staging artifact (store object or local file); best-effort. + fn delete_ipc_staging(&self) { + match &self.store { + Some(store) => { + let _ = crate::store_io::delete_object( + store, + &object_store::path::Path::from(self.ipc_object_path()), + ); + } + None => { + let _ = std::fs::remove_file(self.ipc_staging_path()); + } + } + } + /// Finalize: flush remaining IPC data (sort + write) and return chunk paths + row IDs + CRCs. fn finish( mut self, @@ -332,11 +515,11 @@ impl SortingChunkedWriter { if self.current_rows > 0 { self.flush_and_sort_chunk(reservation)?; } - // Close and remove the trailing IPC staging file - if let Some(mut writer) = self.current_ipc_writer.take() { - writer.finish()?; + // Close and remove the trailing IPC staging artifact + if let Some(writer) = self.current_ipc_writer.take() { + writer.finish_and_close()?; } - let _ = std::fs::remove_file(&self.ipc_staging_path()); + self.delete_ipc_staging(); Ok((self.completed_chunks, self.chunk_row_ids, self.chunk_crcs)) } @@ -374,8 +557,15 @@ pub struct WriterState { reservation: MemoryReservation, /// Final output path (temp file is renamed to this on finalize). filename: String, - /// Temporary path written to before finalize (`temp-`); also the IPC staging base. + /// Temporary path written to before finalize (`temp-`); used only by the local + /// (`store == None`) non-sorted path. Store-backed variants stream to the `ObjectStore`. temp_filename: String, + /// Shard-scoped store (cloned from the engine handle), or `None` for the local path (native + /// unit tests). Used by the sorted (IPC) variant at finalize to run the k-way merge through + /// the store, and carried so the store outlives all in-flight writer ops. + store: Option>, + /// Object path (basename of `filename`) for the final output within the store root. + object_path: String, } /// Path suffix for the intermediate Arrow IPC file used during sort-on-close. @@ -409,10 +599,11 @@ impl NativeParquetWriter { reverse_sorts: Vec, nulls_first: Vec, writer_generation: i64, + store_handle: i64, ) -> Result<*mut WriterState, Box> { log_debug!( - "create_writer called for file: {}, index: {}, schema_address: {}, sort_columns: {:?}, reverse_sorts: {:?}, nulls_first: {:?}, writer_generation: {}", - filename, index_name, schema_address, sort_columns, reverse_sorts, nulls_first, writer_generation + "create_writer called for file: {}, index: {}, schema_address: {}, sort_columns: {:?}, reverse_sorts: {:?}, nulls_first: {:?}, writer_generation: {}, store_handle: {}", + filename, index_name, schema_address, sort_columns, reverse_sorts, nulls_first, writer_generation, store_handle ); if (schema_address as *mut u8).is_null() { @@ -424,6 +615,18 @@ impl NativeParquetWriter { } let temp_filename = Self::temp_filename(&filename); + let object_path = Path::new(&filename) + .file_name() + .and_then(|s| s.to_str()) + .unwrap_or(&filename) + .to_string(); + // Clone the engine-owned store Arc once (never consume the box), or None for the local + // path (native unit tests). Shared by all store-backed sinks and the sorted finalize. + let store: Option> = if store_handle > 0 { + Some(unsafe { (*(store_handle as *const Arc)).clone() }) + } else { + None + }; let arrow_schema = unsafe { FFI_ArrowSchema::from_raw(schema_address as *mut _) }; let schema = Arc::new(arrow::datatypes::Schema::try_from(&arrow_schema)?); log_debug!("Schema created with {} fields", schema.fields().len()); @@ -454,25 +657,40 @@ impl NativeParquetWriter { settings.reverse_sorts.clone(), settings.nulls_first.clone(), writer_generation, + store.clone(), + object_path.clone(), )?; ( WriterVariant::Ipc(Arc::new(Mutex::new(chunked_writer))), None, ) } else { - let file = File::create(&temp_filename)?; - let (crc_file, crc_handle) = CrcWriter::new(file); let props = WriterPropertiesBuilder::build_with_generation( &settings, Some(writer_generation), &schema, ) .map_err(|e| format!("Invalid encoding/compression config: {}", e))?; - let writer = ArrowWriter::try_new(crc_file, schema, Some(props))?; - ( - WriterVariant::Parquet(Arc::new(Mutex::new(writer))), - Some(crc_handle), - ) + + if let Some(store_ref) = store.clone() { + // Production path: stream the Parquet output straight into the shard ObjectStore. + let (writer, crc_handle) = + new_store_parquet_sink(store_ref, &object_path, schema, props)?; + ( + WriterVariant::ParquetStore(Arc::new(Mutex::new(writer))), + Some(crc_handle), + ) + } else { + // Legacy local path (native unit tests / no store): synchronous ArrowWriter to a + // local temp file, renamed into place at finalize. + let file = File::create(&temp_filename)?; + let (crc_file, crc_handle) = CrcWriter::new(file); + let writer = ArrowWriter::try_new(crc_file, schema, Some(props))?; + ( + WriterVariant::Parquet(Arc::new(Mutex::new(writer))), + Some(crc_handle), + ) + } }; let state = Box::new(WriterState { @@ -487,6 +705,8 @@ impl NativeParquetWriter { ), filename, temp_filename, + store, + object_path, }); let handle = Box::into_raw(state); Ok(handle) @@ -553,6 +773,23 @@ impl NativeParquetWriter { drop(writer); state.reservation.reconcile(estimated, actual); } + WriterVariant::ParquetStore(writer_arc) => { + log_debug!("Writing RecordBatch to Parquet ObjectStore sink"); + let batch_bytes = record_batch.get_array_memory_size(); + // Same 3× estimate as the local path; AsyncArrowWriter exposes the same + // `memory_size()` for the in-progress buffered encoding. + let estimated = batch_bytes * 3; + let writer_arc = writer_arc.clone(); + state.reservation.reserve_estimated(estimated)?; + let mut writer = writer_arc.lock().unwrap(); + let before = writer.memory_size(); + // Drive the async write to completion. Java serializes handle-touching + // calls per writer, so no other thread contends this lock/runtime slot. + os_store_runtime().block_on(writer.write(&record_batch))?; + let actual = writer.memory_size().saturating_sub(before); + drop(writer); + state.reservation.reconcile(estimated, actual); + } } Ok(()) } else { @@ -582,6 +819,8 @@ impl NativeParquetWriter { mut reservation, filename, temp_filename, + store, + object_path, } = unsafe { *Box::from_raw(handle) }; log_debug!( "finalize_writer called for file: {} (temp: {})", @@ -604,31 +843,52 @@ impl NativeParquetWriter { temp_filename, total_rows, chunk_paths.len() ); - let (crc32, row_id_mapping) = Self::finalize_sorted_chunks( - &chunk_paths, - &chunk_row_ids, - &chunk_crcs, - &filename, - index_name, - &settings.sort_columns, - &settings.reverse_sorts, - &settings.nulls_first, - writer_generation, - schema.clone(), - &mut reservation, - )?; - - // Clean up sorted chunk files only after successful finalization. + let (crc32, row_id_mapping, merged_metadata) = + Self::finalize_sorted_chunks( + &chunk_paths, + &chunk_row_ids, + &chunk_crcs, + &filename, + &object_path, + store.as_ref(), + index_name, + &settings.sort_columns, + &settings.reverse_sorts, + &settings.nulls_first, + writer_generation, + schema.clone(), + &mut reservation, + )?; + + // Clean up sorted chunk artifacts only after successful finalization. // On failure, chunks are preserved as they may be the only copy of the data. for path in &chunk_paths { - let _ = std::fs::remove_file(path); + match &store { + Some(s) => { + let _ = crate::store_io::delete_object( + s, + &object_store::path::Path::from(path.as_str()), + ); + } + None => { + let _ = std::fs::remove_file(path); + } + } } log_debug!("CRC32 for file {}: {:#010x}", filename, crc32); - let file = File::open(&filename)?; - let reader = SerializedFileReader::new(file)?; - let parquet_metadata = reader.metadata().clone(); + // Prefer the metadata produced by the merge (backend-agnostic). Only the + // single-chunk/empty store cases and the local path fall back to reading + // the final file (for hot indices it is a real local file at `filename`). + let parquet_metadata = match merged_metadata { + Some(md) => md, + None => { + let file = File::open(&filename)?; + let reader = SerializedFileReader::new(file)?; + reader.metadata().clone() + } + }; // Detach mapping from reservation before handing to FFI/Java. // FFI layer will track it via write_pool().grow/shrink. @@ -660,7 +920,8 @@ impl NativeParquetWriter { let crc32 = crc_handle.map(|h| h.crc32()).unwrap_or(0); log_info!("Successfully closed temp writer for: {}", temp_filename); - // Parquet variant is used for non-sorted data; just rename. + // Legacy local path (no ObjectStore): atomically move the finished + // temp file into place. std::fs::rename(&temp_filename, &filename)?; log_debug!("CRC32 for file {}: {:#010x}", filename, crc32); @@ -693,6 +954,37 @@ impl NativeParquetWriter { } } } + WriterVariant::ParquetStore(writer_arc) => { + match Arc::try_unwrap(writer_arc) { + Ok(mutex) => { + let writer = mutex.into_inner().unwrap(); + // `close()` force-flushes the buffered Parquet, completes the ObjectStore + // multipart upload, and returns the full ParquetMetaData — so, unlike the + // local path, there is no file re-open (works for any store backend). The + // CrcObjectWriter has hashed every uploaded byte into `crc_handle`. + let parquet_metadata = + os_store_runtime().block_on(writer.close())?; + let crc32 = crc_handle.map(|h| h.crc32()).unwrap_or(0); + log_info!( + "Successfully closed store-backed writer for: {} (crc32={:#010x})", + filename, + crc32 + ); + Ok(Some(FinalizeResult { + metadata: parquet_metadata, + crc32, + row_id_mapping: None, + })) + } + Err(_) => { + log_error!( + "ERROR: Store-backed writer still in use for file: {}", + filename + ); + Err("Store-backed writer still in use".into()) + } + } + } } } @@ -704,6 +996,8 @@ impl NativeParquetWriter { chunk_row_ids: &[Vec], chunk_crcs: &[u32], output_filename: &str, + output_object_path: &str, + store: Option<&Arc>, index_name: &str, sort_columns: &[String], reverse_sorts: &[bool], @@ -711,7 +1005,8 @@ impl NativeParquetWriter { writer_generation: i64, schema: Arc, reservation: &mut MemoryReservation, - ) -> Result<(u32, Option>), Box> { + ) -> Result<(u32, Option>, Option), Box> + { if chunk_paths.is_empty() { log_info!("finalize_sorted_chunks: no chunks, writing empty Parquet file"); let config = SETTINGS_STORE @@ -724,16 +1019,38 @@ impl NativeParquetWriter { &schema, ) .map_err(|e| format!("Invalid encoding/compression config: {}", e))?; - let file = File::create(output_filename)?; - let writer = ArrowWriter::try_new(file, schema, Some(props))?; - writer.close()?; - return Ok((0, None)); + match store { + Some(s) => { + // Empty Parquet straight to the store. + let (writer, _crc) = + new_store_parquet_sink(s.clone(), output_object_path, schema, props)?; + os_store_runtime().block_on(writer.close())?; + } + None => { + let file = File::create(output_filename)?; + let writer = ArrowWriter::try_new(file, schema, Some(props))?; + writer.close()?; + } + } + return Ok((0, None, None)); } if chunk_paths.len() == 1 { - // Single chunk: just rename to final output (already sorted Parquet) + // Single chunk is already sorted Parquet — move it to the final output. + // Store: metadata-only rename within the store. Local: filesystem rename. log_info!("finalize_sorted_chunks: single chunk, renaming to final output"); - std::fs::rename(&chunk_paths[0], output_filename)?; + match store { + Some(s) => { + crate::store_io::rename_object( + s, + &object_store::path::Path::from(chunk_paths[0].as_str()), + &object_store::path::Path::from(output_object_path), + )?; + } + None => { + std::fs::rename(&chunk_paths[0], output_filename)?; + } + } // Use the CRC computed when the chunk was written let crc32 = chunk_crcs.first().copied().unwrap_or(0); @@ -757,7 +1074,7 @@ impl NativeParquetWriter { None }; - return Ok((crc32, row_id_mapping)); + return Ok((crc32, row_id_mapping, None)); } // Multiple chunks: k-way merge @@ -775,13 +1092,18 @@ impl NativeParquetWriter { ); let merge_output = merge_sorted_with_pool( chunk_paths, - output_filename, + if store.is_some() { + output_object_path + } else { + output_filename + }, index_name, sort_columns, reverse_sorts, nulls_first, writer_generation, &mut merge_reservation, + store.cloned(), ) .map_err(|e| -> Box { format!("Streaming merge failed: {}", e).into() @@ -793,12 +1115,15 @@ impl NativeParquetWriter { merge_duration ); - // Build the flat permutation: result[original_row_id] = new_row_id + // Take the merged output's metadata (backend-agnostic) and mapping; destructure so we can + // free the merge mapping without holding the whole struct. let crc32 = merge_output.crc32; - let row_id_mapping = if !merge_output.mapping.is_empty() && !chunk_row_ids.is_empty() { - let total = merge_output.mapping.len(); + let merged_metadata = merge_output.metadata; + let merge_mapping = merge_output.mapping; + let row_id_mapping = if !merge_mapping.is_empty() && !chunk_row_ids.is_empty() { + let total = merge_mapping.len(); let mapping_bytes = total * std::mem::size_of::(); - // Reserve 2× mapping: merge_output.mapping (alive) + flat_mapping (about to allocate) + // Reserve 2× mapping: merge mapping (alive) + flat_mapping (about to allocate) reservation.request(mapping_bytes * 2)?; let mut flat_mapping = vec![0i64; total]; for i in 0..total { @@ -809,13 +1134,13 @@ impl NativeParquetWriter { for &original_row_id in chunk_ids { let orig_idx = original_row_id as usize; if orig_idx < total && pos < total { - flat_mapping[orig_idx] = merge_output.mapping[pos]; + flat_mapping[orig_idx] = merge_mapping[pos]; } pos += 1; } } - drop(merge_output); - // merge_output.mapping freed — release its share, flat_mapping remains tracked + drop(merge_mapping); + // merge mapping freed — release its share, flat_mapping remains tracked reservation.shrink(mapping_bytes); log_info!( "finalize_sorted_chunks: produced {} permutation entries for {}", @@ -833,7 +1158,7 @@ impl NativeParquetWriter { chunk_paths.len(), merge_duration ); - Ok((crc32, row_id_mapping)) + Ok((crc32, row_id_mapping, Some(merged_metadata))) } /// Sort a batch using RowConverter: converts sort columns into compact diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/tests/merge_integration_tests.rs b/sandbox/plugins/parquet-data-format/src/main/rust/tests/merge_integration_tests.rs index 2c6390dc65a24..85e81b2c33d25 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/tests/merge_integration_tests.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/tests/merge_integration_tests.rs @@ -124,7 +124,7 @@ fn test_unsorted_merge_real_files() { let output_str = output.to_string_lossy().to_string(); // Empty sort columns → unsorted merge - merge_unsorted(&files, &output_str, "test-index", 0).unwrap(); + merge_unsorted(&files, &output_str, "test-index", 0, None).unwrap(); assert!(output.exists(), "Output file was not created"); let actual_rows = count_rows(&output_str); @@ -185,15 +185,13 @@ fn test_sorted_merge_real_files() { let reverse = vec![false]; let nulls_first = vec![false]; - merge_sorted( - &files, - &output_str, - "test-index", - &sort_cols, - &reverse, - &nulls_first, - 0, - ) + merge_sorted(&files, + &output_str, + "test-index", + &sort_cols, + &reverse, + &nulls_first, + 0, None) .unwrap(); assert!(output.exists(), "Output file was not created"); @@ -297,15 +295,13 @@ fn test_tier2_yield_after_batch_boundary() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &[file_a, file_b], - &output, - index, - &["v".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_a, file_b], + &output, + index, + &["v".into()], + &[false], + &[false], + 0, None) .unwrap(); let vals = read_all_int64(&output, "v"); @@ -354,15 +350,13 @@ fn test_tier2_yield_multiple_cursors() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &[file_a, file_b, file_c], - &output, - index, - &["v".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_a, file_b, file_c], + &output, + index, + &["v".into()], + &[false], + &[false], + 0, None) .unwrap(); let vals = read_all_int64(&output, "v"); @@ -405,15 +399,13 @@ fn test_tier2_yield_descending() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &[file_a, file_b], - &output, - index, - &["v".into()], - &[true], - &[false], - 0, - ) + merge_sorted(&[file_a, file_b], + &output, + index, + &["v".into()], + &[true], + &[false], + 0, None) .unwrap(); let vals = read_all_int64(&output, "v"); @@ -456,15 +448,13 @@ fn test_tier2_no_yield_when_equal_to_heap_top() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &[file_a, file_b], - &output, - index, - &["v".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_a, file_b], + &output, + index, + &["v".into()], + &[false], + &[false], + 0, None) .unwrap(); let vals = read_all_int64(&output, "v"); @@ -509,15 +499,13 @@ fn test_tier2_yield_many_small_batches() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &[file_a, file_b], - &output, - index, - &["v".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_a, file_b], + &output, + index, + &["v".into()], + &[false], + &[false], + 0, None) .unwrap(); let vals = read_all_int64(&output, "v"); @@ -626,15 +614,13 @@ fn test_default_settings_ascending_nulls_last() { ); let output = tmp.path().join("out.parquet").to_string_lossy().to_string(); - merge_sorted( - &[file_a, file_b], - &output, - "test_default_asc_nulls_last", - &["v".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_a, file_b], + &output, + "test_default_asc_nulls_last", + &["v".into()], + &[false], + &[false], + 0, None) .unwrap(); // ── Row group structure ────────────────────────────────────────────── @@ -685,15 +671,13 @@ fn test_default_settings_descending_nulls_last() { ); let output = tmp.path().join("out.parquet").to_string_lossy().to_string(); - merge_sorted( - &[file_a, file_b], - &output, - "test_default_desc_nulls_last", - &["v".into()], - &[true], - &[false], - 0, - ) + merge_sorted(&[file_a, file_b], + &output, + "test_default_desc_nulls_last", + &["v".into()], + &[true], + &[false], + 0, None) .unwrap(); // ── Row group structure ────────────────────────────────────────────── @@ -756,15 +740,13 @@ fn test_default_settings_ascending_nulls_first() { ); let output = tmp.path().join("out.parquet").to_string_lossy().to_string(); - merge_sorted( - &[file_a, file_b], - &output, - "test_default_asc_nulls_first", - &["v".into()], - &[false], - &[true], - 0, - ) + merge_sorted(&[file_a, file_b], + &output, + "test_default_asc_nulls_first", + &["v".into()], + &[false], + &[true], + 0, None) .unwrap(); // ── Row group structure ────────────────────────────────────────────── @@ -818,15 +800,13 @@ fn test_single_large_file_passthrough() { ); let output = tmp.path().join("out.parquet").to_string_lossy().to_string(); - merge_sorted( - &[file_a], - &output, - "test_single_large_file", - &["v".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_a], + &output, + "test_single_large_file", + &["v".into()], + &[false], + &[false], + 0, None) .unwrap(); let (rg_sizes, rg_firsts, rg_lasts) = inspect_row_groups(&output, "v"); @@ -873,15 +853,13 @@ fn test_skewed_file_sizes_large_small() { ); let output = tmp.path().join("out.parquet").to_string_lossy().to_string(); - merge_sorted( - &[file_a, file_b], - &output, - "test_skewed_large_small", - &["v".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_a, file_b], + &output, + "test_skewed_large_small", + &["v".into()], + &[false], + &[false], + 0, None) .unwrap(); let total = large + small; @@ -949,15 +927,13 @@ fn test_three_files_middle_exhausts_first() { ); let output = tmp.path().join("out.parquet").to_string_lossy().to_string(); - merge_sorted( - &[file_a, file_b, file_c], - &output, - "test_three_files_middle_exhausts", - &["v".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_a, file_b, file_c], + &output, + "test_three_files_middle_exhausts", + &["v".into()], + &[false], + &[false], + 0, None) .unwrap(); let total = a_count + b_count + c_count; @@ -1007,15 +983,13 @@ fn test_all_duplicate_sort_keys_large() { .collect(); let output = tmp.path().join("out.parquet").to_string_lossy().to_string(); - merge_sorted( - &files, - &output, - "test_all_dupes_large", - &["v".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&files, + &output, + "test_all_dupes_large", + &["v".into()], + &[false], + &[false], + 0, None) .unwrap(); let total = n * 3; @@ -1065,15 +1039,13 @@ fn test_non_multiple_of_batch_size() { ); let output = tmp.path().join("out.parquet").to_string_lossy().to_string(); - merge_sorted( - &[file_a, file_b], - &output, - "test_non_multiple_batch", - &["v".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_a, file_b], + &output, + "test_non_multiple_batch", + &["v".into()], + &[false], + &[false], + 0, None) .unwrap(); let total = a_count + b_count; @@ -1125,15 +1097,13 @@ fn test_rg_size_overshoots_when_batch_straddles_threshold() { ); let output = tmp.path().join("out.parquet").to_string_lossy().to_string(); - merge_sorted( - &[file_a], - &output, - index, - &["v".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_a], + &output, + index, + &["v".into()], + &[false], + &[false], + 0, None) .unwrap(); let (rg_sizes, rg_firsts, rg_lasts) = inspect_row_groups(&output, "v"); @@ -1252,15 +1222,13 @@ fn test_deferred_wide_schema_correctness() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &[file_a, file_b], - &output, - index, - &["ts".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_a, file_b], + &output, + index, + &["ts".into()], + &[false], + &[false], + 0, None) .unwrap(); let ts_vals = read_all_int64(&output, "ts"); @@ -1322,15 +1290,13 @@ fn test_eager_forced_by_high_threshold() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &[file_a, file_b], - &output, - index, - &["ts".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_a, file_b], + &output, + index, + &["ts".into()], + &[false], + &[false], + 0, None) .unwrap(); let ts_vals = read_all_int64(&output, "ts"); @@ -1383,15 +1349,13 @@ fn test_deferred_multi_batch_sync() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &[file_a, file_b], - &output, - index, - &["ts".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_a, file_b], + &output, + index, + &["ts".into()], + &[false], + &[false], + 0, None) .unwrap(); let ts_vals = read_all_int64(&output, "ts"); @@ -1444,15 +1408,13 @@ fn test_deferred_tier3_interleaved() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &[file_a, file_b], - &output, - index, - &["ts".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_a, file_b], + &output, + index, + &["ts".into()], + &[false], + &[false], + 0, None) .unwrap(); let ts_vals = read_all_int64(&output, "ts"); @@ -1510,15 +1472,13 @@ fn test_deferred_vs_eager_identical_output() { .join("merged_deferred.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &[file_a.clone(), file_b.clone()], - &output_deferred, - index_deferred, - &["ts".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_a.clone(), file_b.clone()], + &output_deferred, + index_deferred, + &["ts".into()], + &[false], + &[false], + 0, None) .unwrap(); // Run with eager (threshold=9999) @@ -1529,15 +1489,13 @@ fn test_deferred_vs_eager_identical_output() { .join("merged_eager.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &[file_a, file_b], - &output_eager, - index_eager, - &["ts".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_a, file_b], + &output_eager, + index_eager, + &["ts".into()], + &[false], + &[false], + 0, None) .unwrap(); // Compare outputs — must be identical @@ -1601,15 +1559,13 @@ fn test_deferred_tier1_single_cursor_drain() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &[file_a, file_b], - &output, - index, - &["ts".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_a, file_b], + &output, + index, + &["ts".into()], + &[false], + &[false], + 0, None) .unwrap(); let ts_vals = read_all_int64(&output, "ts"); @@ -1665,15 +1621,13 @@ fn test_deferred_tier1_multi_batch_drain() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &[file_a, file_b], - &output, - index, - &["ts".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_a, file_b], + &output, + index, + &["ts".into()], + &[false], + &[false], + 0, None) .unwrap(); let ts_vals = read_all_int64(&output, "ts"); @@ -1728,15 +1682,13 @@ fn test_deferred_tier2_full_batch_emit() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &[file_a, file_b], - &output, - index, - &["ts".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_a, file_b], + &output, + index, + &["ts".into()], + &[false], + &[false], + 0, None) .unwrap(); let ts_vals = read_all_int64(&output, "ts"); @@ -1792,14 +1744,14 @@ fn test_deferred_tier2_descending() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &[file_a, file_b], - &output, - index, - &["ts".into()], - &[true], - &[false], - 0, // reverse=true (descending) + merge_sorted(&[file_a, file_b], + &output, + index, + &["ts".into()], + &[true], + &[false], + 0, // reverse=true (descending) + None, ) .unwrap(); @@ -1863,15 +1815,13 @@ fn test_deferred_tier3_many_cursors() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &files, - &output, - index, - &["ts".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&files, + &output, + index, + &["ts".into()], + &[false], + &[false], + 0, None) .unwrap(); let ts_vals = read_all_int64(&output, "ts"); @@ -1937,15 +1887,13 @@ fn test_deferred_different_schemas() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &[file_a, file_b], - &output, - index, - &["ts".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_a, file_b], + &output, + index, + &["ts".into()], + &[false], + &[false], + 0, None) .unwrap(); let ts_vals = read_all_int64(&output, "ts"); @@ -2067,15 +2015,13 @@ fn test_deferred_three_files_different_schemas() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &[file_1, file_2, file_3], - &output, - index, - &["ts".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&[file_1, file_2, file_3], + &output, + index, + &["ts".into()], + &[false], + &[false], + 0, None) .unwrap(); let ts_vals = read_all_int64(&output, "ts"); diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/tests/sort_types_tests.rs b/sandbox/plugins/parquet-data-format/src/main/rust/tests/sort_types_tests.rs index 02832ba32f3d8..d564af9f35a76 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/tests/sort_types_tests.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/tests/sort_types_tests.rs @@ -128,15 +128,13 @@ fn test_merge_sort_by_int64() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &files, - &output, - "test", - &["val".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&files, + &output, + "test", + &["val".into()], + &[false], + &[false], + 0, None) .unwrap(); let vals = read_primitive_col::(&output, "val"); @@ -182,15 +180,13 @@ fn test_merge_sort_by_int64_with_nulls() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &files, - &output, - "test", - &["val".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&files, + &output, + "test", + &["val".into()], + &[false], + &[false], + 0, None) .unwrap(); let vals = read_primitive_col::(&output, "val"); @@ -233,15 +229,13 @@ fn test_merge_sort_by_int32() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &files, - &output, - "test", - &["val".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&files, + &output, + "test", + &["val".into()], + &[false], + &[false], + 0, None) .unwrap(); let vals = read_primitive_col::(&output, "val"); @@ -289,15 +283,13 @@ fn test_merge_sort_by_float64() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &files, - &output, - "test", - &["val".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&files, + &output, + "test", + &["val".into()], + &[false], + &[false], + 0, None) .unwrap(); let vals = read_primitive_col::(&output, "val"); @@ -353,15 +345,13 @@ fn test_merge_sort_by_float64_with_nulls() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &files, - &output, - "test", - &["val".into()], - &[false], - &[true], - 0, - ) + merge_sorted(&files, + &output, + "test", + &["val".into()], + &[false], + &[true], + 0, None) .unwrap(); let vals = read_primitive_col::(&output, "val"); @@ -411,15 +401,13 @@ fn test_merge_sort_by_float32() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &files, - &output, - "test", - &["val".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&files, + &output, + "test", + &["val".into()], + &[false], + &[false], + 0, None) .unwrap(); let vals = read_primitive_col::(&output, "val"); @@ -471,15 +459,13 @@ fn test_merge_sort_by_float32_with_nulls() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &files, - &output, - "test", - &["val".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&files, + &output, + "test", + &["val".into()], + &[false], + &[false], + 0, None) .unwrap(); let vals = read_primitive_col::(&output, "val"); @@ -525,15 +511,13 @@ fn test_merge_sort_by_string() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &files, - &output, - "test", - &["val".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&files, + &output, + "test", + &["val".into()], + &[false], + &[false], + 0, None) .unwrap(); let vals = read_string_col(&output, "val"); @@ -588,15 +572,13 @@ fn test_merge_sort_by_string_with_nulls() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &files, - &output, - "test", - &["val".into()], - &[false], - &[true], - 0, - ) + merge_sorted(&files, + &output, + "test", + &["val".into()], + &[false], + &[true], + 0, None) .unwrap(); let vals = read_string_col(&output, "val"); @@ -655,15 +637,13 @@ fn test_merge_sort_descending() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &files, - &output, - "test", - &["val".into()], - &[true], - &[false], - 0, - ) + merge_sorted(&files, + &output, + "test", + &["val".into()], + &[true], + &[false], + 0, None) .unwrap(); let vals = read_primitive_col::(&output, "val"); @@ -719,15 +699,13 @@ fn test_merge_sort_multi_column_string_and_int() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &files, - &output, - "test", - &["category".into(), "priority".into()], - &[false, false], - &[false, false], - 0, - ) + merge_sorted(&files, + &output, + "test", + &["category".into(), "priority".into()], + &[false, false], + &[false, false], + 0, None) .unwrap(); let cats = read_string_col(&output, "category"); @@ -780,15 +758,13 @@ fn test_merge_sort_with_nulls_first() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &files, - &output, - "test", - &["val".into()], - &[false], - &[true], - 0, - ) + merge_sorted(&files, + &output, + "test", + &["val".into()], + &[false], + &[true], + 0, None) .unwrap(); let vals = read_primitive_col::(&output, "val"); @@ -830,15 +806,13 @@ fn test_merge_sort_with_nulls_last() { .join("merged.parquet") .to_string_lossy() .to_string(); - merge_sorted( - &files, - &output, - "test", - &["val".into()], - &[false], - &[false], - 0, - ) + merge_sorted(&files, + &output, + "test", + &["val".into()], + &[false], + &[false], + 0, None) .unwrap(); let vals = read_primitive_col::(&output, "val"); diff --git a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/bridge/ParquetMergeIntegrationTests.java b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/bridge/ParquetMergeIntegrationTests.java index 492f4f0895938..ccabd8826cfd7 100644 --- a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/bridge/ParquetMergeIntegrationTests.java +++ b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/bridge/ParquetMergeIntegrationTests.java @@ -76,7 +76,7 @@ public void testMergeSortedFiles() throws Exception { // 3. Merge String mergedFile = tempDir.resolve("merged.parquet").toString(); - RustBridge.mergeParquetFilesInRust(List.of(Path.of(file1), Path.of(file2), Path.of(file3)), mergedFile, INDEX_NAME, 0L); + RustBridge.mergeParquetFilesInRust(List.of(Path.of(file1), Path.of(file2), Path.of(file3)), mergedFile, INDEX_NAME, 0L, 0L); // 4. Verify merged output ParquetFileMetadata mergedMeta = RustBridge.getFileMetadata(mergedFile); @@ -97,7 +97,7 @@ public void testMergeWithInterleavedTimestamps() throws Exception { String file2 = createSortedFile(tempDir, "f2.parquet", new long[] { 200, 400, 600 }, new String[] { "b", "d", "f" }); String mergedFile = tempDir.resolve("merged.parquet").toString(); - RustBridge.mergeParquetFilesInRust(List.of(Path.of(file1), Path.of(file2)), mergedFile, INDEX_NAME, 0L); + RustBridge.mergeParquetFilesInRust(List.of(Path.of(file1), Path.of(file2)), mergedFile, INDEX_NAME, 0L, 0L); assertEquals(6, RustBridge.getFileMetadata(mergedFile).numRows()); @@ -112,7 +112,7 @@ public void testMergeSingleFile() throws Exception { String file1 = createSortedFile(tempDir, "f1.parquet", new long[] { 10, 20, 30 }, new String[] { "x", "y", "z" }); String mergedFile = tempDir.resolve("merged.parquet").toString(); - RustBridge.mergeParquetFilesInRust(List.of(Path.of(file1)), mergedFile, INDEX_NAME, 0L); + RustBridge.mergeParquetFilesInRust(List.of(Path.of(file1)), mergedFile, INDEX_NAME, 0L, 0L); assertEquals(3, RustBridge.getFileMetadata(mergedFile).numRows());