From d5286a605f6675b774dfa3e6e4fe285769b1af05 Mon Sep 17 00:00:00 2001 From: Shailesh-Kumar-Singh Date: Sun, 19 Apr 2026 03:28:24 +0530 Subject: [PATCH 01/10] add parquet merge support through a K way merge streaming merge sort Signed-off-by: Shailesh-Kumar-Singh --- .../libs/dataformat-native/rust/Cargo.toml | 1 + .../benchmark/VSRRotationBenchmark.java | 9 +- .../opensearch/parquet/ParquetSettings.java | 80 ++- .../parquet/bridge/NativeParquetWriter.java | 18 +- .../parquet/bridge/NativeSettings.java | 80 +++ .../opensearch/parquet/bridge/RustBridge.java | 163 +++++- .../parquet/engine/ParquetIndexingEngine.java | 36 +- .../parquet/merge/ParquetMergeExecutor.java | 31 ++ .../parquet/merge/ParquetMergeStrategy.java | 25 + .../merge/StreamingParquetMergeStrategy.java | 99 ++++ .../opensearch/parquet/vsr/VSRManager.java | 33 +- .../parquet/writer/ParquetWriter.java | 14 +- .../src/main/rust/Cargo.toml | 8 +- .../src/main/rust/src/ffm.rs | 187 ++++++- .../src/main/rust/src/field_config.rs | 45 ++ .../src/main/rust/src/lib.rs | 9 + .../src/main/rust/src/merge/context.rs | 208 ++++++++ .../src/main/rust/src/merge/cursor.rs | 216 ++++++++ .../src/main/rust/src/merge/error.rs | 48 ++ .../src/main/rust/src/merge/heap.rs | 179 +++++++ .../src/main/rust/src/merge/io_task.rs | 175 ++++++ .../src/main/rust/src/merge/mod.rs | 12 + .../src/main/rust/src/merge/schema.rs | 101 ++++ .../src/main/rust/src/merge/sorted.rs | 204 +++++++ .../src/main/rust/src/merge/unsorted.rs | 82 +++ .../src/main/rust/src/native_settings.rs | 125 +++++ .../src/main/rust/src/rate_limited_writer.rs | 213 ++++++++ .../src/main/rust/src/test_utils.rs | 39 +- .../src/main/rust/src/tests/mod.rs | 380 +++++++++---- .../src/main/rust/src/writer.rs | 438 ++++++++++++--- .../rust/src/writer_properties_builder.rs | 198 +++++++ .../rust/tests/merge_integration_tests.rs | 175 ++++++ .../src/main/rust/tests/sort_types_tests.rs | 497 ++++++++++++++++++ .../rust/tests/writer_integration_tests.rs | 9 +- .../bridge/NativeParquetWriterTests.java | 7 +- .../parquet/vsr/VSRManagerTests.java | 33 +- .../parquet/writer/ParquetWriterTests.java | 93 +--- .../index/engine/dataformat/MergeInput.java | 16 +- 38 files changed, 3919 insertions(+), 367 deletions(-) create mode 100644 sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeSettings.java create mode 100644 sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/ParquetMergeExecutor.java create mode 100644 sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/ParquetMergeStrategy.java create mode 100644 sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/StreamingParquetMergeStrategy.java create mode 100644 sandbox/plugins/parquet-data-format/src/main/rust/src/field_config.rs create mode 100644 sandbox/plugins/parquet-data-format/src/main/rust/src/merge/context.rs create mode 100644 sandbox/plugins/parquet-data-format/src/main/rust/src/merge/cursor.rs create mode 100644 sandbox/plugins/parquet-data-format/src/main/rust/src/merge/error.rs create mode 100644 sandbox/plugins/parquet-data-format/src/main/rust/src/merge/heap.rs create mode 100644 sandbox/plugins/parquet-data-format/src/main/rust/src/merge/io_task.rs create mode 100644 sandbox/plugins/parquet-data-format/src/main/rust/src/merge/mod.rs create mode 100644 sandbox/plugins/parquet-data-format/src/main/rust/src/merge/schema.rs create mode 100644 sandbox/plugins/parquet-data-format/src/main/rust/src/merge/sorted.rs create mode 100644 sandbox/plugins/parquet-data-format/src/main/rust/src/merge/unsorted.rs create mode 100644 sandbox/plugins/parquet-data-format/src/main/rust/src/native_settings.rs create mode 100644 sandbox/plugins/parquet-data-format/src/main/rust/src/rate_limited_writer.rs create mode 100644 sandbox/plugins/parquet-data-format/src/main/rust/src/writer_properties_builder.rs create mode 100644 sandbox/plugins/parquet-data-format/src/main/rust/tests/merge_integration_tests.rs create mode 100644 sandbox/plugins/parquet-data-format/src/main/rust/tests/sort_types_tests.rs diff --git a/sandbox/libs/dataformat-native/rust/Cargo.toml b/sandbox/libs/dataformat-native/rust/Cargo.toml index c69ed2fa6c9b5..e353cc71f91ff 100644 --- a/sandbox/libs/dataformat-native/rust/Cargo.toml +++ b/sandbox/libs/dataformat-native/rust/Cargo.toml @@ -51,6 +51,7 @@ once_cell = "1.21.3" crc32fast = "1.4" parking_lot = "0.12.5" lazy_static = "1.4.0" +rayon = "1.10" criterion = { version = "0.5", features = ["async_tokio"] } # Internal diff --git a/sandbox/plugins/parquet-data-format/benchmarks/src/main/java/org/opensearch/parquet/benchmark/VSRRotationBenchmark.java b/sandbox/plugins/parquet-data-format/benchmarks/src/main/java/org/opensearch/parquet/benchmark/VSRRotationBenchmark.java index aa47e2f44b287..aee6483e81efd 100644 --- a/sandbox/plugins/parquet-data-format/benchmarks/src/main/java/org/opensearch/parquet/benchmark/VSRRotationBenchmark.java +++ b/sandbox/plugins/parquet-data-format/benchmarks/src/main/java/org/opensearch/parquet/benchmark/VSRRotationBenchmark.java @@ -10,7 +10,10 @@ import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; +import org.opensearch.Version; +import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.common.settings.Settings; +import org.opensearch.index.IndexSettings; import org.opensearch.index.mapper.KeywordFieldMapper; import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.index.mapper.NumberFieldMapper; @@ -80,6 +83,7 @@ public class VSRRotationBenchmark { private List fieldTypes; private VSRManager vsrManager; private String filePath; + private IndexSettings indexSettings; @Setup(Level.Trial) public void setupTrial() { @@ -123,7 +127,10 @@ public void setupTrial() { public void setup() throws IOException { bufferPool = new ArrowBufferPool(Settings.EMPTY); filePath = Path.of(System.getProperty("java.io.tmpdir"), "benchmark_vsr_" + System.nanoTime() + ".parquet").toString(); - vsrManager = new VSRManager(filePath, schema, bufferPool, maxRowsPerVSR, threadPool, runAsync); + Settings idxSettings = Settings.builder().put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT).build(); + IndexMetadata indexMetadata = IndexMetadata.builder("benchmark-index").settings(idxSettings).build(); + indexSettings = new IndexSettings(indexMetadata, Settings.EMPTY); + vsrManager = new VSRManager(filePath, indexSettings, schema, bufferPool, maxRowsPerVSR, threadPool, runAsync); } @Benchmark diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetSettings.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetSettings.java index 71e57fb0542fa..4f1b8dfb7d4e6 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetSettings.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetSettings.java @@ -9,40 +9,81 @@ package org.opensearch.parquet; import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; +import org.opensearch.core.common.unit.ByteSizeUnit; +import org.opensearch.core.common.unit.ByteSizeValue; import java.util.List; /** - * Node-scoped settings for the Parquet data format plugin. - * - *

All settings are registered with OpenSearch via - * {@link ParquetDataFormatPlugin#getSettings()} and can be configured in - * {@code opensearch.yml} or via cluster settings API. - * - *

    - *
  • {@link #MAX_NATIVE_ALLOCATION} — Maximum native memory allocation for Arrow buffers, - * expressed as a percentage of available non-heap system memory (default {@code "10%"}).
    • - *
  • {@link #MAX_ROWS_PER_VSR} — Row count threshold that triggers VectorSchemaRoot rotation - * during document ingestion (default {@code 50000}).
    • - *
+ * Settings for Parquet data format. */ public final class ParquetSettings { private ParquetSettings() {} - /** Default maximum native memory allocation as a percentage of available non-heap memory. */ public static final String DEFAULT_MAX_NATIVE_ALLOCATION = "10%"; - /** Default maximum number of rows per VectorSchemaRoot before rotation. */ public static final int DEFAULT_MAX_ROWS_PER_VSR = 50000; - /** Maximum native memory allocation for Arrow buffers, as a percentage of non-heap memory. */ + /** Group setting prefix for all Parquet settings. */ + public static final Setting PARQUET_SETTINGS = Setting.groupSetting( + "parquet.", + Setting.Property.IndexScope + ); + + /** Maximum row group size in bytes (default 128MB). */ + public static final Setting ROW_GROUP_SIZE_BYTES = Setting.byteSizeSetting( + "parquet.row_group_size_bytes", + new ByteSizeValue(128, ByteSizeUnit.MB), + Setting.Property.IndexScope + ); + + /** Data page size limit in bytes (default 1MB). */ + public static final Setting PAGE_SIZE_BYTES = Setting.byteSizeSetting( + "parquet.page_size_bytes", + new ByteSizeValue(1, ByteSizeUnit.MB), + Setting.Property.IndexScope + ); + + /** Maximum number of rows per data page (default 20000). */ + public static final Setting PAGE_ROW_LIMIT = Setting.intSetting( + "parquet.page_row_limit", + 20000, + 1, + Setting.Property.IndexScope + ); + + /** Dictionary page size limit in bytes (default 2MB). */ + public static final Setting DICT_SIZE_BYTES = Setting.byteSizeSetting( + "parquet.dict_size_bytes", + new ByteSizeValue(2, ByteSizeUnit.MB), + Setting.Property.IndexScope + ); + + /** Compression codec for Parquet files, e.g. ZSTD, SNAPPY, LZ4_RAW (default LZ4_RAW). */ + public static final Setting COMPRESSION_TYPE = Setting.simpleString( + "parquet.compression_type", + "LZ4_RAW", + Setting.Property.IndexScope + ); + + /** Compression level for the chosen codec (default 2, range 1–9). */ + public static final Setting COMPRESSION_LEVEL = Setting.intSetting( + "parquet.compression_level", + 2, + 1, + 9, + Setting.Property.IndexScope + ); + + /** Maximum native memory allocation for Arrow buffers, as a percentage of non-heap memory (default 10%). */ public static final Setting MAX_NATIVE_ALLOCATION = Setting.simpleString( "parquet.max_native_allocation", DEFAULT_MAX_NATIVE_ALLOCATION, Setting.Property.NodeScope ); - /** Maximum number of rows per VectorSchemaRoot before rotation is triggered. */ + /** Maximum rows per VectorSchemaRoot before rotation is triggered (default 50000). */ public static final Setting MAX_ROWS_PER_VSR = Setting.intSetting( "parquet.max_rows_per_vsr", DEFAULT_MAX_ROWS_PER_VSR, @@ -52,6 +93,11 @@ private ParquetSettings() {} /** Returns all settings defined by the Parquet plugin. */ public static List> getSettings() { - return List.of(MAX_NATIVE_ALLOCATION, MAX_ROWS_PER_VSR); + return List.of( + PARQUET_SETTINGS, + ROW_GROUP_SIZE_BYTES, PAGE_SIZE_BYTES, PAGE_ROW_LIMIT, DICT_SIZE_BYTES, + COMPRESSION_TYPE, COMPRESSION_LEVEL, + MAX_NATIVE_ALLOCATION, MAX_ROWS_PER_VSR + ); } } 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 34b41d635d41a..6db3727e499ae 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 @@ -11,6 +11,7 @@ import org.opensearch.common.SetOnce; import java.io.IOException; +import java.util.List; import java.util.concurrent.atomic.AtomicBoolean; /** @@ -18,7 +19,7 @@ * *

Wraps the stateless JNI methods in {@link RustBridge} with a file-scoped lifecycle: *

    - *
  1. {@code new NativeParquetWriter(filePath, schemaAddress)} — creates the native writer
    2. + *
  2. {@code new NativeParquetWriter(filePath, indexName, schemaAddress, sortColumns, reverseSorts, nullsFirst)} — creates the native writer
    3. *
  3. {@link #write(long, long)} — sends one or more Arrow batches (repeatable)
    4. *
  4. {@link #flush()} — finalizes the Parquet file and returns metadata
    5. *
  5. {@link #sync()} — fsyncs the file to durable storage (calls flush if needed)
    6. @@ -37,12 +38,23 @@ public class NativeParquetWriter { * Creates a new NativeParquetWriter. * * @param filePath the path to the Parquet file to write + * @param indexName the index name for settings lookup * @param schemaAddress the native memory address of the Arrow schema + * @param sortColumns the columns to sort by, or empty list for no sorting + * @param reverseSorts whether each sort column is descending, or empty list + * @param nullsFirst whether nulls sort first for each column, or empty list * @throws IOException if the native writer creation fails */ - public NativeParquetWriter(String filePath, long schemaAddress) throws IOException { + public NativeParquetWriter( + String filePath, + String indexName, + long schemaAddress, + List sortColumns, + List reverseSorts, + List nullsFirst + ) throws IOException { this.filePath = filePath; - RustBridge.createWriter(filePath, schemaAddress); + RustBridge.createWriter(filePath, indexName, schemaAddress, sortColumns, reverseSorts, nullsFirst); } /** diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeSettings.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeSettings.java new file mode 100644 index 0000000000000..a8ae0a7b677da --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeSettings.java @@ -0,0 +1,80 @@ +/* + * 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; + +/** + * Immutable settings passed to the native Rust writer via JNI. + * The Rust side reads values through the getter methods. + * All fields are nullable; the native side falls back to defaults when null. + */ +public class NativeSettings { + + private final String indexName; + private final String compressionType; + private final Integer compressionLevel; + private final Long pageSizeBytes; + private final Integer pageRowLimit; + private final Long dictSizeBytes; + private final Long rowGroupSizeBytes; + private final Boolean bloomFilterEnabled; + private final Double bloomFilterFpp; + private final Long bloomFilterNdv; + + private NativeSettings(Builder builder) { + this.indexName = builder.indexName; + this.compressionType = builder.compressionType; + this.compressionLevel = builder.compressionLevel; + this.pageSizeBytes = builder.pageSizeBytes; + this.pageRowLimit = builder.pageRowLimit; + this.dictSizeBytes = builder.dictSizeBytes; + this.rowGroupSizeBytes = builder.rowGroupSizeBytes; + this.bloomFilterEnabled = builder.bloomFilterEnabled; + this.bloomFilterFpp = builder.bloomFilterFpp; + this.bloomFilterNdv = builder.bloomFilterNdv; + } + + public String getIndexName() { return indexName; } + public String getCompressionType() { return compressionType; } + public Integer getCompressionLevel() { return compressionLevel; } + public Long getPageSizeBytes() { return pageSizeBytes; } + public Integer getPageRowLimit() { return pageRowLimit; } + public Long getDictSizeBytes() { return dictSizeBytes; } + public Long getRowGroupSizeBytes() { return rowGroupSizeBytes; } + public Boolean getBloomFilterEnabled() { return bloomFilterEnabled; } + public Double getBloomFilterFpp() { return bloomFilterFpp; } + public Long getBloomFilterNdv() { return bloomFilterNdv; } + + public static Builder builder() { return new Builder(); } + + public static class Builder { + private String indexName; + private String compressionType; + private Integer compressionLevel; + private Long pageSizeBytes; + private Integer pageRowLimit; + private Long dictSizeBytes; + private Long rowGroupSizeBytes; + private Boolean bloomFilterEnabled; + private Double bloomFilterFpp; + private Long bloomFilterNdv; + + public Builder indexName(String v) { this.indexName = v; return this; } + public Builder compressionType(String v) { this.compressionType = v; return this; } + public Builder compressionLevel(Integer v) { this.compressionLevel = v; return this; } + public Builder pageSizeBytes(Long v) { this.pageSizeBytes = v; return this; } + public Builder pageRowLimit(Integer v) { this.pageRowLimit = v; return this; } + public Builder dictSizeBytes(Long v) { this.dictSizeBytes = v; return this; } + public Builder rowGroupSizeBytes(Long v) { this.rowGroupSizeBytes = v; return this; } + public Builder bloomFilterEnabled(Boolean v) { this.bloomFilterEnabled = v; return this; } + public Builder bloomFilterFpp(Double v) { this.bloomFilterFpp = v; return this; } + public Builder bloomFilterNdv(Long v) { this.bloomFilterNdv = v; return this; } + + public NativeSettings build() { return new NativeSettings(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 c9086cfe4e8e6..27e75fdd51632 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 @@ -18,6 +18,8 @@ import java.lang.foreign.ValueLayout; import java.lang.invoke.MethodHandle; import java.nio.charset.StandardCharsets; +import java.nio.file.Path; +import java.util.List; public class RustBridge { @@ -27,35 +29,40 @@ public class RustBridge { private static final MethodHandle SYNC_TO_DISK; private static final MethodHandle GET_FILE_METADATA; private static final MethodHandle GET_FILTERED_BYTES; + private static final MethodHandle ON_SETTINGS_UPDATE; + private static final MethodHandle REMOVE_SETTINGS; + private static final MethodHandle MERGE_FILES; static { SymbolLookup lib = NativeLibraryLoader.symbolLookup(); Linker linker = Linker.nativeLinker(); CREATE_WRITER = linker.downcallHandle( lib.find("parquet_create_writer").orElseThrow(), - FunctionDescriptor.of(ValueLayout.JAVA_LONG, ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, ValueLayout.JAVA_LONG) + FunctionDescriptor.of( + ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, // file + ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, // index_name + ValueLayout.JAVA_LONG, // schema_address + ValueLayout.ADDRESS, ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, // sort_columns (ptrs, lens, count) + ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, // reverse_sorts (vals, count) + ValueLayout.ADDRESS, ValueLayout.JAVA_LONG // nulls_first (vals, count) + ) ); WRITE = linker.downcallHandle( lib.find("parquet_write").orElseThrow(), FunctionDescriptor.of( ValueLayout.JAVA_LONG, - ValueLayout.ADDRESS, - ValueLayout.JAVA_LONG, - ValueLayout.JAVA_LONG, - ValueLayout.JAVA_LONG + ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, + ValueLayout.JAVA_LONG, ValueLayout.JAVA_LONG ) ); FINALIZE_WRITER = linker.downcallHandle( lib.find("parquet_finalize_writer").orElseThrow(), FunctionDescriptor.of( ValueLayout.JAVA_LONG, - ValueLayout.ADDRESS, - ValueLayout.JAVA_LONG, - ValueLayout.ADDRESS, - ValueLayout.ADDRESS, - ValueLayout.ADDRESS, - ValueLayout.JAVA_LONG, - ValueLayout.ADDRESS, + ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS, ValueLayout.ADDRESS, + ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, ValueLayout.ADDRESS, ValueLayout.ADDRESS ) ); @@ -67,27 +74,71 @@ public class RustBridge { lib.find("parquet_get_file_metadata").orElseThrow(), FunctionDescriptor.of( ValueLayout.JAVA_LONG, - ValueLayout.ADDRESS, - ValueLayout.JAVA_LONG, - ValueLayout.ADDRESS, - ValueLayout.ADDRESS, - ValueLayout.ADDRESS, - ValueLayout.JAVA_LONG, - ValueLayout.ADDRESS + ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS, ValueLayout.ADDRESS, + ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, ValueLayout.ADDRESS ) ); GET_FILTERED_BYTES = linker.downcallHandle( lib.find("parquet_get_filtered_native_bytes_used").orElseThrow(), FunctionDescriptor.of(ValueLayout.JAVA_LONG, ValueLayout.ADDRESS, ValueLayout.JAVA_LONG) ); + ON_SETTINGS_UPDATE = linker.downcallHandle( + lib.find("parquet_on_settings_update").orElseThrow(), + FunctionDescriptor.of( + ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, // index_name + ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, // compression_type + ValueLayout.JAVA_LONG, // compression_level + ValueLayout.JAVA_LONG, // page_size_bytes + ValueLayout.JAVA_LONG, // page_row_limit + ValueLayout.JAVA_LONG, // dict_size_bytes + ValueLayout.JAVA_LONG, // row_group_size_bytes + ValueLayout.JAVA_LONG, // bloom_filter_enabled + ValueLayout.JAVA_DOUBLE, // bloom_filter_fpp + ValueLayout.JAVA_LONG // bloom_filter_ndv + ) + ); + REMOVE_SETTINGS = linker.downcallHandle( + lib.find("parquet_remove_settings").orElseThrow(), + FunctionDescriptor.of(ValueLayout.JAVA_LONG, ValueLayout.ADDRESS, ValueLayout.JAVA_LONG) + ); + MERGE_FILES = linker.downcallHandle( + lib.find("parquet_merge_files").orElseThrow(), + FunctionDescriptor.of( + ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS, ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, // input files (ptrs, lens, count) + ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, // output file + ValueLayout.ADDRESS, ValueLayout.JAVA_LONG // index_name + ) + ); } public static void initLogger() {} - static void createWriter(String file, long schemaAddress) throws IOException { + static void createWriter( + String file, + String indexName, + long schemaAddress, + List sortColumns, + List reverseSorts, + List nullsFirst + ) throws IOException { try (var call = new NativeCall()) { var f = call.str(file); - call.invokeIO(CREATE_WRITER, f.segment(), f.len(), schemaAddress); + var idx = call.str(indexName); + var sorts = call.strArray(sortColumns.toArray(new String[0])); + var reverseArray = marshalBoolList(call, reverseSorts); + var nullsFirstArray = marshalBoolList(call, nullsFirst); + call.invokeIO( + CREATE_WRITER, + f.segment(), f.len(), + idx.segment(), idx.len(), + schemaAddress, + sorts.ptrs(), sorts.lens(), sorts.count(), + reverseArray, (long) reverseSorts.size(), + nullsFirstArray, (long) nullsFirst.size() + ); } } @@ -107,13 +158,9 @@ static ParquetFileMetadata finalizeWriter(String file) throws IOException { var out = call.outBuffer(1024); long rc = call.invokeIO( FINALIZE_WRITER, - f.segment(), - f.len(), - versionOut, - numRowsOut, - out.data(), - (long) out.capacity(), - out.lenOut(), + f.segment(), f.len(), + versionOut, numRowsOut, + out.data(), (long) out.capacity(), out.lenOut(), crc32Out ); if (rc == 1) return null; @@ -142,7 +189,12 @@ public static ParquetFileMetadata getFileMetadata(String file) throws IOExceptio var versionOut = call.intOut(); var numRowsOut = call.longOut(); var out = call.outBuffer(1024); - call.invokeIO(GET_FILE_METADATA, f.segment(), f.len(), versionOut, numRowsOut, out.data(), (long) out.capacity(), out.lenOut()); + call.invokeIO( + GET_FILE_METADATA, + f.segment(), f.len(), + versionOut, numRowsOut, + out.data(), (long) out.capacity(), out.lenOut() + ); int createdByLen = out.actualLength(); return new ParquetFileMetadata( versionOut.get(ValueLayout.JAVA_INT, 0), @@ -162,5 +214,58 @@ public static long getFilteredNativeBytesUsed(String pathPrefix) { } } + public static void onSettingsUpdate(NativeSettings nativeSettings) throws IOException { + try (var call = new NativeCall()) { + var idx = call.str(nativeSettings.getIndexName()); + var ct = nativeSettings.getCompressionType() != null ? call.str(nativeSettings.getCompressionType()) : null; + call.invokeIO( + ON_SETTINGS_UPDATE, + idx.segment(), idx.len(), + ct != null ? ct.segment() : java.lang.foreign.MemorySegment.NULL, ct != null ? ct.len() : -1L, + nativeSettings.getCompressionLevel() != null ? (long) nativeSettings.getCompressionLevel() : -1L, + nativeSettings.getPageSizeBytes() != null ? nativeSettings.getPageSizeBytes() : -1L, + nativeSettings.getPageRowLimit() != null ? (long) nativeSettings.getPageRowLimit() : -1L, + nativeSettings.getDictSizeBytes() != null ? nativeSettings.getDictSizeBytes() : -1L, + nativeSettings.getRowGroupSizeBytes() != null ? nativeSettings.getRowGroupSizeBytes() : -1L, + nativeSettings.getBloomFilterEnabled() != null ? (nativeSettings.getBloomFilterEnabled() ? 1L : 0L) : -1L, + nativeSettings.getBloomFilterFpp() != null ? nativeSettings.getBloomFilterFpp() : -1.0, + nativeSettings.getBloomFilterNdv() != null ? nativeSettings.getBloomFilterNdv() : -1L + ); + } + } + + public static void removeSettings(String indexName) { + try (var call = new NativeCall()) { + var idx = call.str(indexName); + call.invoke(REMOVE_SETTINGS, idx.segment(), idx.len()); + } + } + + public static void mergeParquetFilesInRust(List inputFiles, String outputFile, String indexName) { + String[] paths = inputFiles.stream().map(Path::toString).toArray(String[]::new); + try (var call = new NativeCall()) { + var inputs = call.strArray(paths); + var out = call.str(outputFile); + var idx = call.str(indexName); + call.invoke( + MERGE_FILES, + inputs.ptrs(), inputs.lens(), inputs.count(), + out.segment(), out.len(), + idx.segment(), idx.len() + ); + } + } + + private static java.lang.foreign.MemorySegment marshalBoolList(NativeCall call, List bools) { + if (bools == null || bools.isEmpty()) { + return java.lang.foreign.MemorySegment.NULL; + } + var seg = call.buf(bools.size() * 8); + for (int i = 0; i < bools.size(); i++) { + seg.setAtIndex(ValueLayout.JAVA_LONG, i, bools.get(i) ? 1L : 0L); + } + return seg; + } + private RustBridge() {} } 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 137419c1a839b..f5f789d5b45c4 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 @@ -23,13 +23,18 @@ import org.opensearch.index.shard.ShardPath; import org.opensearch.index.store.FormatChecksumStrategy; import org.opensearch.index.store.PrecomputedChecksumStrategy; +import org.opensearch.parquet.ParquetSettings; +import org.opensearch.parquet.bridge.NativeSettings; import org.opensearch.parquet.bridge.RustBridge; import org.opensearch.parquet.memory.ArrowBufferPool; +import org.opensearch.parquet.merge.ParquetMergeExecutor; +import org.opensearch.parquet.merge.StreamingParquetMergeStrategy; import org.opensearch.parquet.writer.ParquetDocumentInput; import org.opensearch.parquet.writer.ParquetWriter; import org.opensearch.threadpool.ThreadPool; import java.io.IOException; +import java.io.UncheckedIOException; import java.nio.file.FileAlreadyExistsException; import java.nio.file.Files; import java.nio.file.Path; @@ -67,9 +72,10 @@ public class ParquetIndexingEngine implements IndexingExecutionEngine schemaSupplier; private final ArrowBufferPool bufferPool; - private final Settings settings; + private final IndexSettings indexSettings; private final ThreadPool threadPool; private final FormatChecksumStrategy checksumStrategy; + private final Merger parquetMerger; /** * Creates a new ParquetIndexingEngine. @@ -120,7 +126,7 @@ public ParquetIndexingEngine( this.shardPath = shardPath; this.schemaSupplier = schemaSupplier; this.bufferPool = new ArrowBufferPool(settings); - this.settings = settings; + this.indexSettings = indexSettings; this.threadPool = threadPool; this.checksumStrategy = checksumStrategy; try { @@ -130,6 +136,8 @@ public ParquetIndexingEngine( } catch (IOException e) { throw new RuntimeException(e); } + this.parquetMerger = new ParquetMergeExecutor(new StreamingParquetMergeStrategy()); + pushSettingsToRust(); } /** @@ -141,6 +149,23 @@ public FormatChecksumStrategy getChecksumStrategy() { return checksumStrategy; } + private void pushSettingsToRust() { + NativeSettings config = NativeSettings.builder() + .indexName(indexSettings.getIndex().getName()) + .compressionType(indexSettings.getValue(ParquetSettings.COMPRESSION_TYPE)) + .compressionLevel(indexSettings.getValue(ParquetSettings.COMPRESSION_LEVEL)) + .pageSizeBytes(indexSettings.getValue(ParquetSettings.PAGE_SIZE_BYTES).getBytes()) + .pageRowLimit(indexSettings.getValue(ParquetSettings.PAGE_ROW_LIMIT)) + .dictSizeBytes(indexSettings.getValue(ParquetSettings.DICT_SIZE_BYTES).getBytes()) + .rowGroupSizeBytes(indexSettings.getValue(ParquetSettings.ROW_GROUP_SIZE_BYTES).getBytes()) + .build(); + try { + RustBridge.onSettingsUpdate(config); + } catch (IOException e) { + throw new UncheckedIOException("Failed to push Parquet settings to Rust store", e); + } + } + @Override public Writer createWriter(long writerGeneration) { Path filePath = Path.of( @@ -167,7 +192,7 @@ public long getNativeBytesUsed() { @Override public Merger getMerger() { - return null; + return parquetMerger; } @Override @@ -219,6 +244,11 @@ public IndexStoreProvider getProvider() { @Override public void close() throws IOException { + try { + RustBridge.removeSettings(indexSettings.getIndex().getName()); + } catch (Exception e) { + logger.warn("Failed to remove Parquet settings from Rust store for index [{}]", indexSettings.getIndex().getName(), e); + } bufferPool.close(); } } diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/ParquetMergeExecutor.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/ParquetMergeExecutor.java new file mode 100644 index 0000000000000..52d2c2c462d07 --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/ParquetMergeExecutor.java @@ -0,0 +1,31 @@ +/* + * 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.merge; + + +import org.opensearch.index.engine.dataformat.MergeInput; +import org.opensearch.index.engine.dataformat.MergeResult; +import org.opensearch.index.engine.dataformat.Merger; + +/** + * Executes Parquet merge operations using a pluggable {@link ParquetMergeStrategy}. + */ +public class ParquetMergeExecutor implements Merger { + + private final ParquetMergeStrategy strategy; + + public ParquetMergeExecutor(ParquetMergeStrategy strategy) { + this.strategy = strategy; + } + + @Override + public MergeResult merge(MergeInput mergeInput) { + return strategy.mergeParquetFiles(mergeInput); + } +} diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/ParquetMergeStrategy.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/ParquetMergeStrategy.java new file mode 100644 index 0000000000000..e9bb508152ec7 --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/ParquetMergeStrategy.java @@ -0,0 +1,25 @@ +/* + * 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.merge; + + +import org.opensearch.index.engine.dataformat.MergeInput; +import org.opensearch.index.engine.dataformat.MergeResult; + +/** + * Interface defining a Parquet merge strategy. + */ +public interface ParquetMergeStrategy { + + /** + * Performs the actual Parquet merge. + */ + MergeResult mergeParquetFiles(MergeInput mergeInput); + +} diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/StreamingParquetMergeStrategy.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/StreamingParquetMergeStrategy.java new file mode 100644 index 0000000000000..539840d48e5ff --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/StreamingParquetMergeStrategy.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.merge; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.apache.logging.log4j.message.ParameterizedMessage; +import org.opensearch.index.engine.dataformat.DataFormat; +import org.opensearch.index.engine.dataformat.MergeInput; +import org.opensearch.index.engine.dataformat.MergeResult; +import org.opensearch.index.engine.exec.WriterFileSet; +import org.opensearch.parquet.bridge.RustBridge; +import org.opensearch.parquet.engine.ParquetDataFormat; +import org.opensearch.parquet.engine.ParquetIndexingEngine; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.*; + +/** + * Implements merging of Parquet files. + */ +public class StreamingParquetMergeStrategy implements ParquetMergeStrategy { + + private static final Logger logger = + LogManager.getLogger(StreamingParquetMergeStrategy.class); + + @Override + public MergeResult mergeParquetFiles(MergeInput mergeInput) { + + List files = mergeInput.writerFiles(); + long writerGeneration = mergeInput.newWriterGeneration(); + if (files.isEmpty()) { + throw new IllegalArgumentException("No files to merge"); + } + + List filePaths = new ArrayList<>(); + files.forEach(writerFileSet -> writerFileSet.files().forEach( + file -> filePaths.add(Path.of(writerFileSet.directory(), file)))); + + String outputDirectory = files.getFirst().directory(); + String mergedFilePath = getMergedFilePath(writerGeneration, outputDirectory); + String mergedFileName = getMergedFileName(writerGeneration); + + try { + // Merge files in Rust + RustBridge.mergeParquetFilesInRust(filePaths, mergedFilePath, mergeInput.indexName()); + + WriterFileSet mergedWriterFileSet = + WriterFileSet.builder().directory(Path.of(outputDirectory)).addFile(mergedFileName).writerGeneration(writerGeneration).build(); + + Map mergedWriterFileSetMap = Collections.singletonMap( + new ParquetDataFormat(), + mergedWriterFileSet + ); + + return new MergeResult(mergedWriterFileSetMap); + + } catch (Exception exception) { + logger.error( + () -> new ParameterizedMessage( + "Merge failed while creating merged file [{}]", + mergedFilePath + ), + exception + ); + try { + Files.deleteIfExists(Path.of(mergedFilePath)); + logger.info("Stale Merged File Deleted at : [{}]", mergedFilePath); + } catch (Exception innerException) { + logger.error( + () -> new ParameterizedMessage( + "Failed to delete stale merged file [{}]", + mergedFilePath + ), + innerException + ); + + } + throw exception; + } + + } + + private String getMergedFileName(long generation) { + // TODO: For debugging we have added extra "merged" in file name, later we can remove and keep same as writer + return ParquetIndexingEngine.FILE_NAME_PREFIX + "_merged_" + generation + ParquetIndexingEngine.FILE_NAME_EXT; + } + + private String getMergedFilePath(long generation, String outputDirectory) { + return Path.of(outputDirectory, getMergedFileName(generation)).toString(); + } +} 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 5038bf8feb36c..5557f1a34d8dd 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 @@ -13,6 +13,8 @@ import org.apache.arrow.vector.types.pojo.Schema; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.opensearch.index.IndexSettings; +import org.opensearch.index.IndexSortConfig; import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.nativebridge.spi.ArrowExport; import org.opensearch.parquet.ParquetDataFormatPlugin; @@ -23,9 +25,11 @@ import org.opensearch.parquet.memory.ArrowBufferPool; import org.opensearch.parquet.writer.FieldValuePair; import org.opensearch.parquet.writer.ParquetDocumentInput; +import org.opensearch.search.sort.SortOrder; import org.opensearch.threadpool.ThreadPool; import java.io.IOException; +import java.util.List; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; @@ -57,6 +61,7 @@ public class VSRManager implements AutoCloseable { private final AtomicReference managedVSR = new AtomicReference<>(); private final String fileName; + private final IndexSettings indexSettings; private final VSRPool vsrPool; private final ThreadPool threadPool; private final String vsrRotationThread; @@ -66,21 +71,17 @@ public class VSRManager implements AutoCloseable { /** * Creates a new VSRManager with asynchronous background writes (production default). - * - * @param fileName output Parquet file path - * @param schema Arrow schema for vector creation - * @param bufferPool shared Arrow buffer pool - * @param maxRowsPerVSR row threshold triggering VSR rotation - * @param threadPool the thread pool for background native writes */ - public VSRManager(String fileName, Schema schema, ArrowBufferPool bufferPool, int maxRowsPerVSR, ThreadPool threadPool) { - this(fileName, schema, bufferPool, maxRowsPerVSR, threadPool, true); + public VSRManager(String fileName, IndexSettings indexSettings, Schema schema, ArrowBufferPool bufferPool, + int maxRowsPerVSR, ThreadPool threadPool) { + this(fileName, indexSettings, schema, bufferPool, maxRowsPerVSR, threadPool, true); } /** * Creates a new VSRManager. * * @param fileName output Parquet file path + * @param indexSettings the index settings (sort config is read from here) * @param schema Arrow schema for vector creation * @param bufferPool shared Arrow buffer pool * @param maxRowsPerVSR row threshold triggering VSR rotation @@ -90,6 +91,7 @@ public VSRManager(String fileName, Schema schema, ArrowBufferPool bufferPool, in */ public VSRManager( String fileName, + IndexSettings indexSettings, Schema schema, ArrowBufferPool bufferPool, int maxRowsPerVSR, @@ -97,6 +99,7 @@ public VSRManager( boolean runAsync ) { this.fileName = fileName; + this.indexSettings = indexSettings; this.vsrPool = new VSRPool("pool-" + fileName, schema, bufferPool, maxRowsPerVSR); this.threadPool = threadPool; this.vsrRotationThread = runAsync ? ParquetDataFormatPlugin.PARQUET_THREAD_POOL_NAME : ThreadPool.Names.SAME; @@ -123,7 +126,7 @@ public void addDocument(ParquetDocumentInput doc) throws IOException { parquetField.createField(fieldType, activeVSR, pair.getValue()); } int rowIndex = activeVSR.getRowCount(); - BigIntVector rowIdVector = (BigIntVector) activeVSR.getVector("_row_id"); + BigIntVector rowIdVector = (BigIntVector) activeVSR.getVector("___row_id"); if (rowIdVector != null) { rowIdVector.setSafe(rowIndex, doc.getRowId()); } @@ -210,9 +213,19 @@ public void close() { } private void initializeWriter() { + // Read sort config from index settings + List sortColumns = IndexSortConfig.INDEX_SORT_FIELD_SETTING.get(indexSettings.getSettings()); + List sortOrders = IndexSortConfig.INDEX_SORT_ORDER_SETTING.get(indexSettings.getSettings()); + List reverseSorts = sortOrders.stream().map(o -> o == SortOrder.DESC).toList(); + + List missingValues = IndexSortConfig.INDEX_SORT_MISSING_SETTING.get(indexSettings.getSettings()); + List nullsFirst = missingValues.stream().map("_first"::equals).collect(java.util.stream.Collectors.toList()); + + String indexName = indexSettings.getIndex().getName(); + ArrowSchema arrowSchema = managedVSR.get().exportSchema(); try { - writer = new NativeParquetWriter(fileName, arrowSchema.memoryAddress()); + writer = new NativeParquetWriter(fileName, indexName, arrowSchema.memoryAddress(), sortColumns, reverseSorts, nullsFirst); } catch (Exception e) { throw new RuntimeException("Failed to initialize Parquet writer: " + e.getMessage(), e); } finally { 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 f02c4893a702b..22165a6f13905 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 @@ -9,7 +9,7 @@ package org.opensearch.parquet.writer; import org.apache.arrow.vector.types.pojo.Schema; -import org.opensearch.common.settings.Settings; +import org.opensearch.index.IndexSettings; import org.opensearch.index.engine.dataformat.FileInfos; import org.opensearch.index.engine.dataformat.WriteResult; import org.opensearch.index.engine.dataformat.Writer; @@ -33,7 +33,7 @@ * by the {@link VSRManager}, and flushed to a Parquet file via the native Rust writer. * *

    Writer-level settings (e.g., {@code parquet.max_rows_per_vsr}) are extracted from - * the {@link Settings} passed at construction time and propagated to the VSR layer. + * the {@link IndexSettings} passed at construction time and propagated to the VSR layer. * *

    The returned {@link FileInfos} from {@link #flush()} contains the file path, writer * generation, and row count for downstream commit tracking. @@ -54,7 +54,7 @@ public class ParquetWriter implements Writer { * @param dataFormat the Parquet data format instance * @param schema Arrow schema for vector creation * @param bufferPool shared Arrow buffer pool - * @param settings node settings for writer configuration + * @param indexSettings index settings for writer configuration * @param threadPool the thread pool for background native writes * @param checksumStrategy strategy to register pre-computed checksums on */ @@ -64,15 +64,19 @@ public ParquetWriter( ParquetDataFormat dataFormat, Schema schema, ArrowBufferPool bufferPool, - Settings settings, ThreadPool threadPool, FormatChecksumStrategy checksumStrategy + IndexSettings indexSettings, + ThreadPool threadPool ) { this.file = file; this.writerGeneration = writerGeneration; this.dataFormat = dataFormat; - this.vsrManager = new VSRManager(file, schema, bufferPool, ParquetSettings.MAX_ROWS_PER_VSR.get(settings), threadPool); this.checksumStrategy = checksumStrategy; + this.vsrManager = new VSRManager( + file, indexSettings, schema, bufferPool, + ParquetSettings.MAX_ROWS_PER_VSR.get(indexSettings.getSettings()), threadPool + ); } @Override 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 22466d27a3d60..f93074fc71d5c 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/Cargo.toml +++ b/sandbox/plugins/parquet-data-format/src/main/rust/Cargo.toml @@ -14,17 +14,13 @@ crate-type = ["rlib"] [dependencies] arrow = { workspace = true } -arrow-array = { workspace = true } -arrow-schema = { workspace = true } -arrow-buffer = { workspace = true } -log = { workspace = true } parquet = { workspace = true } lazy_static = { workspace = true } dashmap = { workspace = true } -chrono = { workspace = true } -mimalloc = { workspace = true } tempfile = { workspace = true } native-bridge-common = { workspace = true } +rayon = { workspace = true } +tokio = { workspace = true } crc32fast = { workspace = true } [dev-dependencies] 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 f015a49110ec3..3e6a430c456fc 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 @@ -16,7 +16,9 @@ use std::str; use native_bridge_common::ffm_safe; -use crate::writer::NativeParquetWriter; +use crate::native_settings::NativeSettings; +use crate::merge; +use crate::writer::{NativeParquetWriter, SETTINGS_STORE}; unsafe fn str_from_raw<'a>(ptr: *const u8, len: i64) -> Result<&'a str, String> { if ptr.is_null() { @@ -29,15 +31,70 @@ unsafe fn str_from_raw<'a>(ptr: *const u8, len: i64) -> Result<&'a str, String> str::from_utf8(bytes).map_err(|e| format!("invalid UTF-8: {}", e)) } +/// Decode a parallel (pointers, lengths, count) triple into `Vec`. +unsafe fn str_array_from_raw( + ptrs: *const *const u8, + lens: *const i64, + count: i64, +) -> Result, String> { + if count == 0 { + return Ok(vec![]); + } + if ptrs.is_null() || lens.is_null() { + return Err("null string array pointer".to_string()); + } + let n = count as usize; + let mut out = Vec::with_capacity(n); + for i in 0..n { + let p = *ptrs.add(i); + let l = *lens.add(i); + out.push(str_from_raw(p, l)?.to_string()); + } + Ok(out) +} + +/// Decode a parallel (pointers, count) array of i64 values interpreted as booleans (0 = false). +unsafe fn bool_array_from_raw( + vals: *const i64, + count: i64, +) -> Vec { + if count == 0 || vals.is_null() { + return vec![]; + } + let n = count as usize; + (0..n).map(|i| *vals.add(i) != 0).collect() +} + +// --------------------------------------------------------------------------- +// Writer lifecycle +// --------------------------------------------------------------------------- + #[ffm_safe] #[no_mangle] pub unsafe extern "C" fn parquet_create_writer( file_ptr: *const u8, file_len: i64, + index_name_ptr: *const u8, + index_name_len: i64, schema_address: i64, + sort_ptrs: *const *const u8, + sort_lens: *const i64, + sort_count: i64, + reverse_vals: *const i64, + reverse_count: i64, + nulls_first_vals: *const i64, + nulls_first_count: i64, ) -> i64 { - let filename = str_from_raw(file_ptr, file_len).map_err(|e| format!("parquet_create_writer: {}", e))?.to_string(); - NativeParquetWriter::create_writer(filename, schema_address) + let filename = str_from_raw(file_ptr, file_len) + .map_err(|e| format!("parquet_create_writer file: {}", e))?.to_string(); + let index_name = str_from_raw(index_name_ptr, index_name_len) + .map_err(|e| format!("parquet_create_writer index_name: {}", e))?.to_string(); + let sort_columns = str_array_from_raw(sort_ptrs, sort_lens, sort_count) + .map_err(|e| format!("parquet_create_writer sort_columns: {}", e))?; + let reverse_sorts = bool_array_from_raw(reverse_vals, reverse_count); + let nulls_first = bool_array_from_raw(nulls_first_vals, nulls_first_count); + + NativeParquetWriter::create_writer(filename, index_name, schema_address, sort_columns, reverse_sorts, nulls_first) .map(|_| 0) .map_err(|e| e.to_string()) } @@ -141,3 +198,127 @@ pub unsafe extern "C" fn parquet_get_filtered_native_bytes_used( 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 } + +// --------------------------------------------------------------------------- +// Settings management +// --------------------------------------------------------------------------- + +/// Update native settings for an index. Nullable fields use sentinel -1 for "not set". +#[ffm_safe] +#[no_mangle] +pub unsafe extern "C" fn parquet_on_settings_update( + index_name_ptr: *const u8, + index_name_len: i64, + compression_type_ptr: *const u8, + compression_type_len: i64, + compression_level: i64, + page_size_bytes: i64, + page_row_limit: i64, + dict_size_bytes: i64, + row_group_size_bytes: i64, + bloom_filter_enabled: i64, + bloom_filter_fpp: f64, + bloom_filter_ndv: i64, +) -> i64 { + let index_name = str_from_raw(index_name_ptr, index_name_len) + .map_err(|e| format!("parquet_on_settings_update index_name: {}", e))?.to_string(); + + let compression_type = if compression_type_ptr.is_null() || compression_type_len < 0 { + None + } else { + Some(str_from_raw(compression_type_ptr, compression_type_len) + .map_err(|e| format!("parquet_on_settings_update compression_type: {}", e))?.to_string()) + }; + + fn opt_i32(v: i64) -> Option { if v < 0 { None } else { Some(v as i32) } } + fn opt_usize(v: i64) -> Option { if v < 0 { None } else { Some(v as usize) } } + fn opt_bool(v: i64) -> Option { if v < 0 { None } else { Some(v != 0) } } + fn opt_f64(v: f64) -> Option { if v < 0.0 { None } else { Some(v) } } + fn opt_u64(v: i64) -> Option { if v < 0 { None } else { Some(v as u64) } } + + let config = NativeSettings { + index_name: Some(index_name.clone()), + compression_type, + compression_level: opt_i32(compression_level), + page_size_bytes: opt_usize(page_size_bytes), + page_row_limit: opt_usize(page_row_limit), + dict_size_bytes: opt_usize(dict_size_bytes), + row_group_size_bytes: opt_usize(row_group_size_bytes), + bloom_filter_enabled: opt_bool(bloom_filter_enabled), + bloom_filter_fpp: opt_f64(bloom_filter_fpp), + bloom_filter_ndv: opt_u64(bloom_filter_ndv), + ..Default::default() + }; + + SETTINGS_STORE.insert(index_name, config); + Ok(0) +} + +#[ffm_safe] +#[no_mangle] +pub unsafe extern "C" fn parquet_remove_settings( + index_name_ptr: *const u8, + index_name_len: i64, +) -> i64 { + let index_name = str_from_raw(index_name_ptr, index_name_len) + .map_err(|e| format!("parquet_remove_settings: {}", e))?.to_string(); + SETTINGS_STORE.remove(&index_name); + Ok(0) +} + +// --------------------------------------------------------------------------- +// Merge +// --------------------------------------------------------------------------- + +#[ffm_safe] +#[no_mangle] +pub unsafe extern "C" fn parquet_merge_files( + input_ptrs: *const *const u8, + input_lens: *const i64, + input_count: i64, + output_ptr: *const u8, + output_len: i64, + index_name_ptr: *const u8, + index_name_len: i64, +) -> i64 { + let input_files = str_array_from_raw(input_ptrs, input_lens, input_count) + .map_err(|e| format!("parquet_merge_files inputs: {}", e))?; + let output_path = str_from_raw(output_ptr, output_len) + .map_err(|e| format!("parquet_merge_files output: {}", e))?; + let index_name = str_from_raw(index_name_ptr, index_name_len) + .map_err(|e| format!("parquet_merge_files index_name: {}", e))?; + + let (sort_cols, reverse_flags, nulls_first_flags) = match SETTINGS_STORE.get(index_name) { + Some(s) => { + let sc = s.sort_columns.clone(); + let rf = s.reverse_sorts.clone(); + let nf = s.nulls_first.clone(); + if !sc.is_empty() && rf.is_empty() { + crate::log_info!("parquet_merge_files: sort columns present but reverse_sorts is empty for index '{}', defaulting to ascending", index_name); + } + if !sc.is_empty() && nf.is_empty() { + crate::log_info!("parquet_merge_files: sort columns present but nulls_first is empty for index '{}', defaulting to nulls last", index_name); + } + (sc, rf, nf) + } + None => { + crate::log_info!("parquet_merge_files: no settings found for index '{}', proceeding with unsorted merge", index_name); + (vec![], vec![], vec![]) + } + }; + + if sort_cols.is_empty() { + merge::merge_unsorted(&input_files, output_path, index_name) + } else { + merge::merge_sorted( + &input_files, + output_path, + index_name, + &sort_cols, + &reverse_flags, + &nulls_first_flags, + ) + } + .map(|_| 0) + .map_err(|e| format!("{}", e)) +} diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/field_config.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/field_config.rs new file mode 100644 index 0000000000000..a13b904e3f8d3 --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/field_config.rs @@ -0,0 +1,45 @@ +/* + * 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. + */ + +#[derive(Debug, Clone, Default)] +pub struct FieldConfig { + pub compression_type: Option, + pub compression_level: Option, +} + +impl FieldConfig { + pub fn new() -> Self { + Self::default() + } + + pub fn is_empty(&self) -> bool { + self.compression_type.is_none() && self.compression_level.is_none() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_field_config_default() { + let config = FieldConfig::default(); + assert!(config.is_empty()); + } + + #[test] + fn test_field_config_construction() { + let config = FieldConfig { + compression_type: Some("SNAPPY".to_string()), + compression_level: Some(1), + }; + assert_eq!(config.compression_type, Some("SNAPPY".to_string())); + assert_eq!(config.compression_level, Some(1)); + assert!(!config.is_empty()); + } +} 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 c13fd3e8b5f10..ff62af3543296 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 @@ -14,5 +14,14 @@ mod tests; pub mod writer; pub mod ffm; +pub mod native_settings; +pub mod field_config; +pub mod writer_properties_builder; +pub mod rate_limited_writer; +pub mod merge; +pub use native_settings::NativeSettings; +pub use field_config::FieldConfig; +pub use writer_properties_builder::WriterPropertiesBuilder; +pub use writer::SETTINGS_STORE; pub use native_bridge_common::{log_info, log_error, log_debug}; 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 new file mode 100644 index 0000000000000..f68ea432830f0 --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/context.rs @@ -0,0 +1,208 @@ +use std::fs::File; +use std::path::Path; +use std::sync::Arc; + +use arrow::array::RecordBatch; +use arrow::compute::concat_batches; +use arrow::datatypes::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; +use parquet::arrow::arrow_writer::{ArrowRowGroupWriterFactory, compute_leaves}; +use parquet::file::writer::SerializedFileWriter; +use parquet::schema::types::SchemaDescriptor; +use rayon::prelude::*; +use tokio::sync::{mpsc as tokio_mpsc, oneshot}; + +use crate::rate_limited_writer::RateLimitedWriter; +use crate::writer_properties_builder::WriterPropertiesBuilder; +use crate::{log_debug, SETTINGS_STORE}; + +use super::error::{MergeError, MergeResult}; +use super::io_task::{ + get_merge_pool, spawn_io_task, IoCommand, 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, +/// IO channel, buffered batches, and counters. Used by both sorted and +/// unsorted merge paths. +pub struct MergeContext { + data_schema: Arc, + output_schema: Arc, + rg_writer_factory: ArrowRowGroupWriterFactory, + io_tx: tokio_mpsc::Sender, + output_chunks: Vec, + output_row_count: usize, + output_flush_rows: usize, + row_group_index: usize, + next_row_id: i64, + total_rows_written: usize, +} + +impl MergeContext { + /// Creates a new merge context: builds union schemas, opens the output + /// writer, and spawns the background IO task. + pub fn new( + arrow_schemas: Vec, + parquet_descriptors: &[SchemaDescriptor], + output_path: &str, + index_name: &str, + output_flush_rows: usize, + ) -> 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() + ))); + } + } + + let union_data_schema = ArrowSchema::try_merge(arrow_schemas).map_err(|e| { + MergeError::Logic(format!( + "Failed to compute union schema across input files: {}", + e + )) + })?; + let data_schema = Arc::new(union_data_schema); + + let mut output_fields: Vec = data_schema + .fields() + .iter() + .map(|f| f.as_ref().clone()) + .collect(); + output_fields.push(ArrowField::new( + ROW_ID_COLUMN_NAME, + ArrowDataType::Int64, + false, + )); + let output_schema = Arc::new(ArrowSchema::new(output_fields)); + + let parquet_root = build_parquet_root_schema(parquet_descriptors)?; + + let output_file = File::create(output_path)?; + let throttled_writer = + RateLimitedWriter::new(output_file, RATE_LIMIT_MB_PER_SEC).map_err(MergeError::Io)?; + + let config = SETTINGS_STORE + .get(index_name) + .map(|r| r.clone()) + .unwrap_or_default(); + let writer_props = Arc::new(WriterPropertiesBuilder::build(&config)); + + let writer = SerializedFileWriter::new(throttled_writer, parquet_root, writer_props)?; + let rg_writer_factory = ArrowRowGroupWriterFactory::new(&writer, output_schema.clone()); + let io_tx = spawn_io_task(writer); + + Ok(Self { + data_schema, + output_schema, + rg_writer_factory, + io_tx, + output_chunks: Vec::new(), + output_row_count: 0, + output_flush_rows, + row_group_index: 0, + next_row_id: 0, + total_rows_written: 0, + }) + } + + pub fn data_schema(&self) -> &Arc { + &self.data_schema + } + + /// Buffers a batch (already padded to data_schema) and auto-flushes when + /// the row count threshold is reached. + pub fn push_batch(&mut self, batch: RecordBatch) -> MergeResult<()> { + self.output_row_count += batch.num_rows(); + self.output_chunks.push(batch); + if self.output_row_count >= self.output_flush_rows { + self.flush()?; + } + Ok(()) + } + + /// Concat buffered batches, append row IDs, encode columns in parallel, + /// and send the encoded row group to the IO task. + pub fn flush(&mut self) -> MergeResult<()> { + if self.output_chunks.is_empty() { + return Ok(()); + } + + let merged = concat_batches(&self.data_schema, self.output_chunks.as_slice())?; + self.output_chunks.clear(); + let n = merged.num_rows(); + + let with_id = append_row_id(&merged, self.next_row_id, &self.output_schema)?; + drop(merged); + + let col_writers = self + .rg_writer_factory + .create_column_writers(self.row_group_index)?; + + let mut leaves_and_writers = Vec::new(); + { + let mut writer_iter = col_writers.into_iter(); + for (arr, field) in with_id.columns().iter().zip(self.output_schema.fields()) { + for leaf in compute_leaves(field, arr)? { + let col_writer = writer_iter.next().ok_or_else(|| { + MergeError::Logic("Fewer column writers than leaf columns".into()) + })?; + leaves_and_writers.push((leaf, col_writer)); + } + } + } + + let chunk_results: Vec< + Result, + > = get_merge_pool().install(|| { + leaves_and_writers + .into_par_iter() + .map(|(leaf, mut col_writer)| { + col_writer.write(&leaf)?; + col_writer.close() + }) + .collect() + }); + + let mut encoded_chunks = Vec::with_capacity(chunk_results.len()); + for r in chunk_results { + encoded_chunks.push(r?); + } + + self.io_tx + .blocking_send(IoCommand::WriteRowGroup(encoded_chunks)) + .map_err(|_| MergeError::Logic("IO task terminated unexpectedly".into()))?; + + self.row_group_index += 1; + self.next_row_id += n as i64; + self.total_rows_written += n; + self.output_row_count = 0; + + log_debug!( + "[RUST] Flushed row group {}: {} rows (total: {})", + self.row_group_index - 1, + n, + self.total_rows_written + ); + + Ok(()) + } + + /// Final flush + close the IO task. Returns Parquet metadata. + pub fn finish(mut self) -> MergeResult { + self.flush()?; + + let (reply_tx, reply_rx) = + oneshot::channel::>(); + + self.io_tx + .blocking_send(IoCommand::Close(reply_tx)) + .map_err(|_| MergeError::Logic("IO task terminated before close".into()))?; + + drop(self.io_tx); + + reply_rx + .blocking_recv() + .map_err(|_| MergeError::Logic("IO task terminated during close".into()))? + } +} 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 new file mode 100644 index 0000000000000..8c48ed1d4b89b --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/cursor.rs @@ -0,0 +1,216 @@ +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 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::schema::projection_indices_excluding_row_id; +/// A cursor over a single sorted Parquet input file. +/// +/// Each cursor reads batches sequentially and prefetches the next batch on the +/// shared Rayon pool to overlap IO with merge computation. +pub struct FileCursor { + reader: Arc>, + prefetch_rx: std::sync::mpsc::Receiver>>, + prefetch_tx: std::sync::mpsc::SyncSender>>, + prefetch_pending: bool, + pub current_batch: Option, + pub row_idx: usize, + pub file_id: usize, + pub sort_col_indices: Vec, + pub sort_col_types: Vec, + pub nulls_first: Vec, +} + +impl FileCursor { + /// Opens a Parquet file and creates a cursor positioned at the first row. + /// + /// Returns `(cursor, projected_arrow_schema, parquet_schema_descriptor)` + /// so the caller can build union schemas without re-opening the file. + pub fn new( + path: &str, + file_id: usize, + sort_columns: &[String], + nulls_first: &[bool], + batch_size: usize, + ) -> MergeResult<(Self, Arc, SchemaDescriptor)> { + let file = File::open(path)?; + let builder = ParquetRecordBatchReaderBuilder::try_new(file)?; + let schema = builder.schema().clone(); + + let mut sort_col_types = Vec::with_capacity(sort_columns.len()); + for col_name in sort_columns { + let dt = schema + .fields() + .iter() + .find(|f| f.name() == col_name.as_str()) + .map(|f| f.data_type().clone()) + .ok_or_else(|| { + MergeError::Logic(format!( + "Sort column '{}' not found in file '{}' (cursor {})", + col_name, path, file_id + )) + })?; + sort_col_types.push(dt); + } + + let parquet_schema_descr = builder.parquet_schema().clone(); + let projection_indices = projection_indices_excluding_row_id(&schema); + + let projection = + parquet::arrow::ProjectionMask::roots(&parquet_schema_descr, projection_indices); + + let mut reader = builder + .with_batch_size(batch_size) + .with_projection(projection) + .build()?; + + let first_batch = match reader.next() { + Some(Ok(b)) if b.num_rows() > 0 => b, + Some(Err(e)) => return Err(e.into()), + _ => { + return Err(MergeError::Logic(format!( + "File '{}' (cursor {}) yielded no rows despite passing validation", + path, file_id + ))); + } + }; + + let projected_schema = first_batch.schema(); + + let mut sort_col_indices = Vec::with_capacity(sort_columns.len()); + for col_name in sort_columns { + let idx = projected_schema + .fields() + .iter() + .position(|f| f.name() == col_name.as_str()) + .ok_or_else(|| { + MergeError::Logic(format!( + "Sort column '{}' not found after projection in file '{}'", + col_name, path + )) + })?; + sort_col_indices.push(idx); + } + + let (prefetch_tx, prefetch_rx) = + std::sync::mpsc::sync_channel::>>(1); + + let reader = Arc::new(Mutex::new(reader)); + + let mut cursor = Self { + reader, + prefetch_rx, + prefetch_tx, + prefetch_pending: false, + current_batch: Some(first_batch), + row_idx: 0, + file_id, + sort_col_indices, + sort_col_types, + nulls_first: nulls_first.to_vec(), + }; + + cursor.start_prefetch(); + + Ok((cursor, projected_schema, parquet_schema_descr)) + } + + fn start_prefetch(&mut self) { + if self.prefetch_pending { + return; + } + self.prefetch_pending = true; + + let reader = Arc::clone(&self.reader); + let tx = self.prefetch_tx.clone(); + + get_merge_pool().spawn(move || { + let mut reader = reader.lock().unwrap(); + let result = match reader.next() { + Some(Ok(batch)) if batch.num_rows() > 0 => Some(Ok(batch)), + Some(Err(e)) => Some(Err(MergeError::Arrow(e))), + _ => None, + }; + let _ = tx.send(result); + }); + } + + pub fn load_next_batch(&mut self) -> MergeResult { + self.current_batch = None; + + match self.prefetch_rx.recv() { + Ok(Some(Ok(batch))) => { + self.current_batch = Some(batch); + self.row_idx = 0; + self.prefetch_pending = false; + self.start_prefetch(); + Ok(true) + } + Ok(Some(Err(e))) => { + self.prefetch_pending = false; + Err(e) + } + Ok(None) | Err(_) => { + self.prefetch_pending = false; + Ok(false) + } + } + } + + #[inline] + pub fn current_sort_values(&self) -> MergeResult> { + let batch = self + .current_batch + .as_ref() + .ok_or_else(|| MergeError::Logic("Cursor exhausted".into()))?; + get_sort_values(batch, self.row_idx, &self.sort_col_indices, &self.sort_col_types, &self.nulls_first) + } + + #[inline] + pub fn last_sort_values(&self) -> MergeResult> { + let batch = self + .current_batch + .as_ref() + .ok_or_else(|| MergeError::Logic("Cursor exhausted".into()))?; + get_sort_values( + batch, + batch.num_rows() - 1, + &self.sort_col_indices, + &self.sort_col_types, + &self.nulls_first, + ) + } + + #[inline] + pub fn batch_height(&self) -> usize { + self.current_batch.as_ref().map_or(0, |b| b.num_rows()) + } + + #[inline] + pub fn take_slice(&self, start: usize, len: usize) -> RecordBatch { + self.current_batch.as_ref().unwrap().slice(start, len) + } + + pub fn advance(&mut self) -> MergeResult { + if self.current_batch.is_none() { + return Ok(false); + } + self.row_idx += 1; + if self.row_idx >= self.current_batch.as_ref().unwrap().num_rows() { + self.current_batch = None; + return self.load_next_batch(); + } + Ok(true) + } + + pub fn advance_past_batch(&mut self) -> MergeResult { + self.current_batch = None; + self.load_next_batch() + } +} diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/error.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/error.rs new file mode 100644 index 0000000000000..1c8faef6cda32 --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/error.rs @@ -0,0 +1,48 @@ +use std::error::Error; + +/// Result type alias for merge operations. +pub type MergeResult = Result; + +/// Unified error type for all merge failures. +#[derive(Debug)] +pub enum MergeError { + /// Error from the Arrow compute or array layer. + Arrow(arrow::error::ArrowError), + /// Error from the Parquet reader or writer. + Parquet(parquet::errors::ParquetError), + /// Filesystem or network IO error. + Io(std::io::Error), + /// Logic or invariant violation within the merge algorithm. + Logic(String), +} + +impl std::fmt::Display for MergeError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + MergeError::Arrow(e) => write!(f, "Arrow error: {e}"), + MergeError::Parquet(e) => write!(f, "Parquet error: {e}"), + MergeError::Io(e) => write!(f, "IO error: {e}"), + MergeError::Logic(s) => write!(f, "{s}"), + } + } +} + +impl Error for MergeError {} + +impl From for MergeError { + fn from(e: arrow::error::ArrowError) -> Self { + MergeError::Arrow(e) + } +} + +impl From for MergeError { + fn from(e: parquet::errors::ParquetError) -> Self { + MergeError::Parquet(e) + } +} + +impl From for MergeError { + fn from(e: std::io::Error) -> Self { + MergeError::Io(e) + } +} diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/heap.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/heap.rs new file mode 100644 index 0000000000000..8fefd64af4521 --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/heap.rs @@ -0,0 +1,179 @@ +use std::cmp::Ordering; +use std::sync::Arc; + +use arrow::array::{AsArray, RecordBatch}; +use arrow::datatypes::{ + DataType as ArrowDataType, Date32Type, Date64Type, DurationMicrosecondType, + DurationMillisecondType, DurationNanosecondType, DurationSecondType, Float32Type, Float64Type, + Int16Type, Int32Type, Int64Type, Int8Type, TimestampMicrosecondType, + TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, UInt32Type, +}; + +use super::error::{MergeError, MergeResult}; + +// ============================================================================= +// SortKey — typed sort value with null ordering baked in +// ============================================================================= + +#[derive(Debug, Clone)] +pub enum SortKey { + NullFirst, + NullLast, + Int(i64), + Float(f64), + Bytes(Vec), +} + +impl Eq for SortKey {} + +impl PartialEq for SortKey { + fn eq(&self, other: &Self) -> bool { + self.cmp(other) == Ordering::Equal + } +} + +impl Ord for SortKey { + fn cmp(&self, other: &Self) -> Ordering { + match (self, other) { + (SortKey::NullFirst, SortKey::NullFirst) => Ordering::Equal, + (SortKey::NullFirst, _) => Ordering::Less, + (_, SortKey::NullFirst) => Ordering::Greater, + (SortKey::NullLast, SortKey::NullLast) => Ordering::Equal, + (SortKey::NullLast, _) => Ordering::Greater, + (_, SortKey::NullLast) => Ordering::Less, + (SortKey::Int(a), SortKey::Int(b)) => a.cmp(b), + (SortKey::Float(a), SortKey::Float(b)) => a.total_cmp(b), + (SortKey::Bytes(a), SortKey::Bytes(b)) => a.cmp(b), + // Same column always produces the same variant; cross-variant is unreachable. + _ => Ordering::Equal, + } + } +} + +impl PartialOrd for SortKey { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +// ============================================================================= +// Sort-direction helpers +// ============================================================================= + +/// Lexicographic comparison of two sort-key tuples, respecting per-column +/// sort direction. Returns `Ordering::Equal` when all values match. +#[inline(always)] +pub fn cmp_sort_values(a: &[SortKey], b: &[SortKey], reverse_sorts: &[bool]) -> Ordering { + for (i, (av, bv)) in a.iter().zip(b.iter()).enumerate() { + let ord = av.cmp(bv); + if ord != Ordering::Equal { + let reverse = reverse_sorts.get(i).copied().unwrap_or(false); + return if reverse { ord.reverse() } else { ord }; + } + } + Ordering::Equal +} + +// ============================================================================= +// HeapItem for k-way merge +// ============================================================================= + +#[derive(Debug)] +pub struct HeapItem { + pub sort_values: Vec, + pub file_id: usize, + pub reverse_sorts: Arc>, +} + +impl Eq for HeapItem {} + +impl PartialEq for HeapItem { + fn eq(&self, other: &Self) -> bool { + self.sort_values == other.sort_values + } +} + +impl Ord for HeapItem { + fn cmp(&self, other: &Self) -> Ordering { + // Swap other/self so max-heap behaves as min-heap. + cmp_sort_values(&other.sort_values, &self.sort_values, &self.reverse_sorts) + } +} + +impl PartialOrd for HeapItem { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +// ============================================================================= +// Sort value extraction +// ============================================================================= + +#[inline] +pub fn get_sort_value( + batch: &RecordBatch, + row: usize, + col_idx: usize, + dtype: &ArrowDataType, + null_first: bool, +) -> MergeResult { + let col = batch.column(col_idx); + if col.is_null(row) { + return Ok(if null_first { SortKey::NullFirst } else { SortKey::NullLast }); + } + let key = match dtype { + // Integer types → SortKey::Int + ArrowDataType::Int64 => SortKey::Int(col.as_primitive::().value(row)), + ArrowDataType::Int32 => SortKey::Int(col.as_primitive::().value(row) as i64), + ArrowDataType::Int16 => SortKey::Int(col.as_primitive::().value(row) as i64), + ArrowDataType::Int8 => SortKey::Int(col.as_primitive::().value(row) as i64), + ArrowDataType::UInt32 => SortKey::Int(col.as_primitive::().value(row) as i64), + ArrowDataType::Date32 => SortKey::Int(col.as_primitive::().value(row) as i64), + ArrowDataType::Date64 => SortKey::Int(col.as_primitive::().value(row)), + ArrowDataType::Timestamp(unit, _) => SortKey::Int(match unit { + arrow::datatypes::TimeUnit::Second => col.as_primitive::().value(row), + arrow::datatypes::TimeUnit::Millisecond => col.as_primitive::().value(row), + arrow::datatypes::TimeUnit::Microsecond => col.as_primitive::().value(row), + arrow::datatypes::TimeUnit::Nanosecond => col.as_primitive::().value(row), + }), + ArrowDataType::Duration(unit) => SortKey::Int(match unit { + arrow::datatypes::TimeUnit::Second => col.as_primitive::().value(row), + arrow::datatypes::TimeUnit::Millisecond => col.as_primitive::().value(row), + arrow::datatypes::TimeUnit::Microsecond => col.as_primitive::().value(row), + arrow::datatypes::TimeUnit::Nanosecond => col.as_primitive::().value(row), + }), + + // Float types → SortKey::Float + ArrowDataType::Float64 => SortKey::Float(col.as_primitive::().value(row)), + ArrowDataType::Float32 => SortKey::Float(col.as_primitive::().value(row) as f64), + + // String types → SortKey::Bytes + ArrowDataType::Utf8 => SortKey::Bytes(col.as_string::().value(row).as_bytes().to_vec()), + ArrowDataType::LargeUtf8 => SortKey::Bytes(col.as_string::().value(row).as_bytes().to_vec()), + + other => { + return Err(MergeError::Logic(format!( + "Unsupported sort column type: {:?}", + other + ))); + } + }; + Ok(key) +} + +#[inline] +pub fn get_sort_values( + batch: &RecordBatch, + row: usize, + col_indices: &[usize], + dtypes: &[ArrowDataType], + nulls_first: &[bool], +) -> MergeResult> { + let mut values = Vec::with_capacity(col_indices.len()); + for (i, (col_idx, dtype)) in col_indices.iter().zip(dtypes.iter()).enumerate() { + let nf = nulls_first.get(i).copied().unwrap_or(false); + values.push(get_sort_value(batch, row, *col_idx, dtype, nf)?); + } + Ok(values) +} 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 new file mode 100644 index 0000000000000..0cf810c0b2aac --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/io_task.rs @@ -0,0 +1,175 @@ +use std::fs::File; +use std::sync::OnceLock; + +use parquet::file::metadata::ParquetMetaData; +use parquet::file::writer::SerializedFileWriter; + +use rayon::ThreadPool; + +use tokio::runtime::Runtime; +use tokio::sync::{mpsc as tokio_mpsc, oneshot}; +use tokio::task::JoinHandle; + +use crate::rate_limited_writer::RateLimitedWriter; +use crate::log_error; + +use super::error::{MergeError, MergeResult}; +// ============================================================================= +// Constants +// ============================================================================= + +/// Number of rows to request per Parquet read batch. +pub const BATCH_SIZE: usize = 100_000; + +/// Approximate number of rows to buffer before flushing a row group. +pub const OUTPUT_FLUSH_ROWS: usize = 1_000_000; + +/// Disk write rate limit in MB/s. +pub const RATE_LIMIT_MB_PER_SEC: f64 = 20.0; + +/// Number of threads in the shared Rayon pool for parallel column encoding. +const RAYON_NUM_THREADS: usize = 4; + +/// Bounded channel capacity between the merge loop and the IO task. +const IO_CHANNEL_BUFFER: usize = 2; + +// ============================================================================= +// Process-wide shared Rayon thread pool +// ============================================================================= + +static MERGE_POOL: OnceLock = OnceLock::new(); + +pub fn get_merge_pool() -> &'static ThreadPool { + MERGE_POOL.get_or_init(|| { + rayon::ThreadPoolBuilder::new() + .num_threads(RAYON_NUM_THREADS) + .thread_name(|idx| format!("parquet-merge-{}", idx)) + .build() + .expect("Failed to build parquet-merge Rayon thread pool") + }) +} + +// ============================================================================= +// Process-wide shared Tokio runtime for async IO +// ============================================================================= + +static IO_RUNTIME: OnceLock = OnceLock::new(); + +fn get_io_runtime() -> &'static Runtime { + IO_RUNTIME.get_or_init(|| { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(4) + .thread_name("parquet-io") + .enable_all() + .build() + .expect("Failed to build tokio IO runtime") + }) +} + +// ============================================================================= +// IO task protocol +// ============================================================================= + +/// Commands sent from the merge loop to the background IO task. +pub enum IoCommand { + WriteRowGroup(Vec), + Close(oneshot::Sender>), +} + +async fn drain_on_error(rx: &mut tokio_mpsc::Receiver, msg: &str) { + while let Some(cmd) = rx.recv().await { + if let IoCommand::Close(reply) = cmd { + let _ = reply.send(Err(MergeError::Logic( + format!("Prior IO write failed: {msg}"), + ))); + } + } +} + +/// Spawns the background IO task on the shared Tokio runtime. +/// +/// The IO task owns the `SerializedFileWriter` and receives encoded row groups +/// over a bounded channel. Each disk write is dispatched to `spawn_blocking` +/// but is **not** awaited immediately — this allows the merge loop to prepare +/// the next row group while the current one is still being flushed to disk. +pub fn spawn_io_task( + writer: SerializedFileWriter>, +) -> tokio_mpsc::Sender { + let (tx, mut rx) = tokio_mpsc::channel::(IO_CHANNEL_BUFFER); + + get_io_runtime().spawn(async move { + let mut writer: Option>> = Some(writer); + let mut in_flight: Option< + JoinHandle>>>, + > = None; + + while let Some(cmd) = rx.recv().await { + match cmd { + IoCommand::WriteRowGroup(chunks) => { + if let Some(handle) = in_flight.take() { + match handle.await { + Ok(Ok(w)) => writer = Some(w), + Ok(Err(e)) => { + let msg = format!("{e}"); + log_error!("[RUST] IO write error during merge: {}", e); + drain_on_error(&mut rx, &msg).await; + return; + } + Err(e) => { + let msg = format!("{e}"); + log_error!("[RUST] IO spawn_blocking panicked during merge: {}", e); + drain_on_error(&mut rx, &msg).await; + return; + } + } + } + + let w = writer.take().unwrap(); + in_flight = Some(tokio::task::spawn_blocking(move || { + let mut w = w; + let mut rg_writer = w.next_row_group()?; + for chunk in chunks { + chunk.append_to_row_group(&mut rg_writer)?; + } + rg_writer.close()?; + Ok(w) + })); + } + + IoCommand::Close(reply) => { + if let Some(handle) = in_flight.take() { + match handle.await { + Ok(Ok(w)) => writer = Some(w), + Ok(Err(e)) => { + let _ = reply.send(Err(e)); + return; + } + Err(e) => { + let _ = reply.send(Err(MergeError::Logic( + format!("IO panic during final write: {e}"), + ))); + return; + } + } + } + + let w = writer.take().unwrap(); + let result = tokio::task::spawn_blocking(move || { + w.close().map_err(MergeError::from) + }) + .await; + + let _ = match result { + Ok(r) => reply.send(r), + Err(e) => reply.send(Err(MergeError::Logic( + format!("Close panicked: {e}"), + ))), + }; + return; + } + } + } + }); + + tx +} 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 new file mode 100644 index 0000000000000..71e026e4b79c7 --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/mod.rs @@ -0,0 +1,12 @@ +mod context; +mod cursor; +pub mod error; +pub mod heap; +pub mod io_task; +pub mod schema; +mod sorted; +mod unsorted; + +pub use error::{MergeError, MergeResult}; +pub use sorted::merge_sorted; +pub use unsorted::merge_unsorted; diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/schema.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/schema.rs new file mode 100644 index 0000000000000..3376d9545fdbd --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/schema.rs @@ -0,0 +1,101 @@ +use std::collections::HashSet; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int64Array, RecordBatch}; +use arrow::datatypes::Schema as ArrowSchema; +use parquet::basic::Repetition; +use parquet::schema::types::Type; + +use super::error::MergeResult; + +/// Reserved column name for the synthetic row identifier added during merge. +pub const ROW_ID_COLUMN_NAME: &str = "___row_id"; + +/// Builds the output Parquet schema as the union of pre-read schema descriptors. +/// +/// The output schema contains every column seen across all inputs, except: +/// - Any existing `___row_id` column is removed. +/// - A fresh `___row_id` INT64 REQUIRED column is appended at the end. +pub fn build_parquet_root_schema( + schema_descriptors: &[parquet::schema::types::SchemaDescriptor], +) -> MergeResult> { + let mut seen_names: HashSet = HashSet::new(); + let mut parquet_fields: Vec> = Vec::new(); + + for descr in schema_descriptors { + let root = descr.root_schema(); + for field in root.get_fields() { + if field.name() != ROW_ID_COLUMN_NAME + && seen_names.insert(field.name().to_string()) + { + parquet_fields.push(Arc::new(field.as_ref().clone())); + } + } + } + + let row_id_type = + Type::primitive_type_builder(ROW_ID_COLUMN_NAME, parquet::basic::Type::INT64) + .with_repetition(Repetition::REQUIRED) + .build()?; + parquet_fields.push(Arc::new(row_id_type)); + + let parquet_root = Type::group_type_builder("schema") + .with_fields(parquet_fields) + .build()?; + + Ok(Arc::new(parquet_root)) +} + +/// Returns column indices that exclude `___row_id`, for use as a projection mask. +pub fn projection_indices_excluding_row_id(schema: &ArrowSchema) -> Vec { + schema + .fields() + .iter() + .enumerate() + .filter(|(_, f)| f.name() != ROW_ID_COLUMN_NAME) + .map(|(i, _)| i) + .collect() +} + +/// Pads a batch to conform to the target schema by adding null-filled columns +/// for any fields present in `target_schema` but missing from the batch. +/// +/// Returns the batch unchanged (no copy) when schemas already match. +pub fn pad_batch_to_schema( + batch: &RecordBatch, + target_schema: &Arc, +) -> MergeResult { + let batch_schema = batch.schema(); + if batch_schema.fields() == target_schema.fields() { + return Ok(batch.clone()); + } + + let num_rows = batch.num_rows(); + let mut columns: Vec = Vec::with_capacity(target_schema.fields().len()); + + for field in target_schema.fields() { + match batch_schema.index_of(field.name()) { + Ok(col_idx) => columns.push(batch.column(col_idx).clone()), + Err(_) => { + columns.push(arrow::array::new_null_array(field.data_type(), num_rows)); + } + } + } + + Ok(RecordBatch::try_new(target_schema.clone(), columns)?) +} + +/// Appends a `___row_id` column with sequential values `[start_id, start_id + N)` +/// to the given batch, producing a new batch with the output schema. +pub fn append_row_id( + batch: &RecordBatch, + start_id: i64, + output_schema: &Arc, +) -> MergeResult { + let n = batch.num_rows() as i64; + let row_ids = Int64Array::from_iter_values(start_id..start_id + n); + let mut columns: Vec = batch.columns().to_vec(); + columns.push(Arc::new(row_ids)); + let result = RecordBatch::try_new(output_schema.clone(), columns)?; + Ok(result) +} 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 new file mode 100644 index 0000000000000..e39d09c2da87c --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/sorted.rs @@ -0,0 +1,204 @@ +use std::cmp::Ordering; +use std::collections::BinaryHeap; +use std::sync::Arc; + +use arrow::datatypes::Schema as ArrowSchema; +use parquet::schema::types::SchemaDescriptor; + +use crate::{log_debug, log_info}; + +use super::context::MergeContext; +use super::cursor::FileCursor; +use super::heap::{cmp_sort_values, get_sort_values, HeapItem}; +use super::io_task::{get_merge_pool, BATCH_SIZE, OUTPUT_FLUSH_ROWS}; +use super::schema::pad_batch_to_schema; + +/// Performs a streaming k-way merge with an explicit sort direction per column. +pub fn merge_sorted( + input_files: &[String], + output_path: &str, + index_name: &str, + sort_columns: &[String], + reverse_sorts: &[bool], + nulls_first: &[bool], +) -> super::MergeResult<()> { + let batch_size = BATCH_SIZE; + let output_flush_rows = OUTPUT_FLUSH_ROWS; + if input_files.is_empty() { + return Ok(()); + } + + if sort_columns.is_empty() { + return Err(super::MergeError::Logic( + "merge_sorted called with empty sort_columns; use merge_unsorted instead".into(), + )); + } + + let pool = get_merge_pool(); + let direction_label = if reverse_sorts.iter().all(|&r| !r) { + "ascending" + } else if reverse_sorts.iter().all(|&r| r) { + "descending" + } else { + "mixed" + }; + + log_info!( + "[RUST] Starting streaming merge ({}): {} input files, sort_columns={:?}, \ + batch_size={}, flush_rows={}, merge_threads={}, output='{}'", + direction_label, + input_files.len(), + sort_columns, + batch_size, + output_flush_rows, + pool.current_num_threads(), + output_path + ); + + // ── Phase 1: Initialize cursors and collect schemas ───────────────── + let mut cursors: Vec = Vec::with_capacity(input_files.len()); + let mut arrow_schemas: Vec = Vec::with_capacity(input_files.len()); + let mut parquet_descriptors: Vec = Vec::with_capacity(input_files.len()); + + for (file_id, path) in input_files.iter().enumerate() { + log_debug!("[RUST] Opening cursor {} for file: {}", file_id, path); + let (cursor, projected_schema, parquet_descr) = + FileCursor::new(path, file_id, sort_columns, nulls_first, batch_size)?; + cursors.push(cursor); + arrow_schemas.push(projected_schema.as_ref().clone()); + parquet_descriptors.push(parquet_descr); + } + + let num_cursors = cursors.len(); + + // ── Phase 2: Create MergeContext (union schemas, writer, IO task) ─── + let mut ctx = MergeContext::new( + arrow_schemas, + &parquet_descriptors, + output_path, + index_name, + output_flush_rows, + )?; + + log_info!( + "[RUST] Merge initialized ({}): {} cursors", + direction_label, + num_cursors + ); + + // ── Phase 3: Seed the heap ────────────────────────────────────────── + let reverse_sorts_arc = Arc::new(reverse_sorts.to_vec()); + let mut heap: BinaryHeap = BinaryHeap::with_capacity(num_cursors); + for cursor in &cursors { + let sv = cursor.current_sort_values()?; + heap.push(HeapItem { + sort_values: sv, + file_id: cursor.file_id, + reverse_sorts: Arc::clone(&reverse_sorts_arc), + }); + } + + // ── Phase 4: K-way merge loop — three-tier cascade ────────────────── + while let Some(item) = heap.pop() { + let file_id = item.file_id; + + // TIER 1: Single cursor remaining — drain it + if heap.is_empty() { + let cursor = &mut cursors[file_id]; + loop { + let remaining = cursor.batch_height() - cursor.row_idx; + if remaining > 0 { + let slice = cursor.take_slice(cursor.row_idx, remaining); + let padded = pad_batch_to_schema(&slice, ctx.data_schema())?; + ctx.push_batch(padded)?; + } + if !cursor.advance_past_batch()? { + break; + } + } + break; + } + + // TIER 2 & 3: Multiple cursors active + let cursor = &mut cursors[file_id]; + + loop { + let heap_top = &heap.peek().unwrap().sort_values; + + // TIER 2: Entire remaining batch fits before heap top + let last_val = cursor.last_sort_values()?; + if cmp_sort_values(&last_val, heap_top, reverse_sorts) != Ordering::Greater { + let remaining = cursor.batch_height() - cursor.row_idx; + let slice = cursor.take_slice(cursor.row_idx, remaining); + let padded = pad_batch_to_schema(&slice, ctx.data_schema())?; + ctx.push_batch(padded)?; + + if !cursor.advance_past_batch()? { + break; + } + continue; + } + + // TIER 3: Binary search for the exact boundary + let run_start = cursor.row_idx; + let batch_h = cursor.batch_height(); + let batch = cursor.current_batch.as_ref().unwrap(); + + let mut lo = run_start; + let mut hi = batch_h - 1; + + while lo + 1 < hi { + let mid = lo + (hi - lo) / 2; + let mid_val = get_sort_values( + batch, + mid, + &cursor.sort_col_indices, + &cursor.sort_col_types, + &cursor.nulls_first, + )?; + + if cmp_sort_values(&mid_val, heap_top, reverse_sorts) != Ordering::Greater { + lo = mid; + } else { + hi = mid; + } + } + let run_end = lo; + + let run_len = run_end - run_start + 1; + if run_len > 0 { + let slice = cursor.take_slice(run_start, run_len); + let padded = pad_batch_to_schema(&slice, ctx.data_schema())?; + ctx.push_batch(padded)?; + } + + cursor.row_idx = run_end; + if !cursor.advance()? { + break; + } + + let next_val = cursor.current_sort_values()?; + if cmp_sort_values(&next_val, heap_top, reverse_sorts) == Ordering::Greater { + heap.push(HeapItem { + sort_values: next_val, + file_id, + reverse_sorts: Arc::clone(&reverse_sorts_arc), + }); + break; + } + } + } + + // ── Phase 5: Close ────────────────────────────────────────────────── + let _metadata = ctx.finish()?; + + log_info!( + "[RUST] Merge complete ({}): {} total rows written to '{}' in {} row groups", + direction_label, + _metadata.file_metadata().num_rows(), + output_path, + _metadata.num_row_groups() + ); + + Ok(()) +} 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 new file mode 100644 index 0000000000000..34e2a98ead0e6 --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/unsorted.rs @@ -0,0 +1,82 @@ +use std::fs::File; + +use arrow::array::RecordBatchReader; +use arrow::datatypes::Schema as ArrowSchema; +use parquet::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder}; +use parquet::schema::types::SchemaDescriptor; + +use crate::{log_debug, log_info}; + +use super::context::MergeContext; +use super::error::MergeResult; +use super::io_task::{BATCH_SIZE, OUTPUT_FLUSH_ROWS}; +use super::schema::{pad_batch_to_schema, projection_indices_excluding_row_id}; + +/// Unsorted merge: reads each input file sequentially, pads to union schema, +/// rewrites `___row_id` with globally sequential values. No sorting performed. +pub fn merge_unsorted( + input_files: &[String], + output_path: &str, + index_name: &str, +) -> MergeResult<()> { + log_info!( + "[RUST] Starting unsorted merge: {} input files, output='{}'", + input_files.len(), + output_path + ); + + // Single pass: collect schemas and build readers. + 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()); + + for path in input_files { + let file = File::open(path)?; + let builder = ParquetRecordBatchReaderBuilder::try_new(file)?; + let schema = builder.schema().clone(); + let parquet_descr = builder.parquet_schema().clone(); + + 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()); + parquet_descriptors.push(parquet_descr); + readers.push(reader); + } + + let mut ctx = MergeContext::new( + arrow_schemas, + &parquet_descriptors, + output_path, + index_name, + OUTPUT_FLUSH_ROWS, + )?; + + // Iterate readers for data. + for (file_idx, reader) in readers.into_iter().enumerate() { + log_debug!( + "[RUST] Unsorted merge: processing file {} of {}", + file_idx + 1, + input_files.len() + ); + + for batch_result in reader { + let batch = batch_result?; + let padded = pad_batch_to_schema(&batch, ctx.data_schema())?; + ctx.push_batch(padded)?; + } + } + + let _metadata = ctx.finish()?; + + log_info!( + "[RUST] Unsorted merge complete: {} total rows written to '{}' in {} row groups", + _metadata.file_metadata().num_rows(), + output_path, + _metadata.num_row_groups() + ); + + Ok(()) +} diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/native_settings.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/native_settings.rs new file mode 100644 index 0000000000000..293b5cda4ebcc --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/native_settings.rs @@ -0,0 +1,125 @@ +/* + * 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. + */ + +use std::collections::HashMap; + +use crate::field_config::FieldConfig; + +#[derive(Debug, Clone, Default)] +pub struct NativeSettings { + pub index_name: Option, + pub compression_level: Option, + pub compression_type: Option, + pub page_size_bytes: Option, + pub page_row_limit: Option, + pub dict_size_bytes: Option, + pub row_group_size_bytes: Option, + pub field_configs: Option>, + pub custom_settings: Option>, + pub bloom_filter_enabled: Option, + pub bloom_filter_fpp: Option, + pub bloom_filter_ndv: Option, + pub sort_columns: Vec, + pub reverse_sorts: Vec, + pub nulls_first: Vec, +} + +impl NativeSettings { + pub fn new() -> Self { + Self::default() + } + + pub fn get_compression_type(&self) -> &str { + self.compression_type.as_deref().unwrap_or("LZ4_RAW") + } + + pub fn get_compression_level(&self) -> i32 { + self.compression_level.unwrap_or(2) + } + + pub fn get_page_size_bytes(&self) -> usize { + self.page_size_bytes.unwrap_or(1024 * 1024) + } + + pub fn get_page_row_limit(&self) -> usize { + self.page_row_limit.unwrap_or(20000) + } + + pub fn get_dict_size_bytes(&self) -> usize { + self.dict_size_bytes.unwrap_or(2 * 1024 * 1024) + } + + pub fn get_row_group_size_bytes(&self) -> usize { + self.row_group_size_bytes.unwrap_or(128 * 1024 * 1024) + } + + pub fn get_bloom_filter_enabled(&self) -> bool { + self.bloom_filter_enabled.unwrap_or(true) + } + + pub fn get_bloom_filter_fpp(&self) -> f64 { + self.bloom_filter_fpp.unwrap_or(0.1) + } + + pub fn get_bloom_filter_ndv(&self) -> u64 { + self.bloom_filter_ndv.unwrap_or(100_000) + } + + pub fn get_field_config(&self, field_name: &str) -> Option<&FieldConfig> { + self.field_configs.as_ref()?.get(field_name) + } + + pub fn has_field_configs(&self) -> bool { + self.field_configs.as_ref().map_or(false, |configs| !configs.is_empty()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_defaults() { + let config = NativeSettings::default(); + assert_eq!(config.get_compression_type(), "LZ4_RAW"); + assert_eq!(config.get_compression_level(), 2); + assert_eq!(config.get_page_row_limit(), 20000); + assert_eq!(config.get_dict_size_bytes(), 2 * 1024 * 1024); + } + + #[test] + fn test_struct_construction() { + let config = NativeSettings { + compression_type: Some("SNAPPY".to_string()), + compression_level: Some(1), + ..Default::default() + }; + assert_eq!(config.get_compression_type(), "SNAPPY"); + assert_eq!(config.get_compression_level(), 1); + } + + #[test] + fn test_field_configs() { + use crate::field_config::FieldConfig; + use std::collections::HashMap; + + let mut field_configs = HashMap::new(); + field_configs.insert("timestamp".to_string(), FieldConfig { + compression_type: Some("SNAPPY".to_string()), + compression_level: None, + }); + let config = NativeSettings { + compression_type: Some("ZSTD".to_string()), + field_configs: Some(field_configs), + ..Default::default() + }; + assert!(config.has_field_configs()); + let fc = config.get_field_config("timestamp").unwrap(); + assert_eq!(fc.compression_type, Some("SNAPPY".to_string())); + } +} diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/rate_limited_writer.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/rate_limited_writer.rs new file mode 100644 index 0000000000000..32826276b0fd7 --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/rate_limited_writer.rs @@ -0,0 +1,213 @@ +/* + * 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. + */ + +use std::io::{Result, Write}; +use std::sync::{Arc, RwLock}; +use std::thread::sleep; +use std::time::{Duration, Instant}; + +// TODO: Make this value dynamic based on resource availability (e.g., adjust ±x% based on IOPS pressure) +const MIN_PAUSE_CHECK_MSEC: f64 = 20.0; +const BYTES_PER_MB: f64 = 1024.0 * 1024.0; +const MAX_MIN_PAUSE_CHECK_BYTES: usize = 1024 * 1024; // 1 MB +const MSEC_TO_SEC: f64 = 1000.0; + +/// Configuration for rate limiting behavior. +struct RateLimiterConfig { + /// Maximum throughput in megabytes per second + mb_per_sec: f64, + /// Minimum bytes to write before checking if pause is needed + min_pause_check_bytes: usize, +} + +/// A writer that rate-limits write operations to a specified throughput. +/// +/// This writer wraps another writer and ensures that data is written at a maximum +/// rate specified in megabytes per second. It uses periodic pauses to maintain +/// the target rate, checking after a minimum number of bytes have been written. +/// +/// # Rate Limiting Strategy +/// +/// The rate limiter works by: +/// 1. Tracking bytes written since the last pause +/// 2. Periodically checking if enough time has elapsed for the bytes written +/// 3. Sleeping if the write rate exceeds the configured limit +/// +/// The minimum pause check interval is calculated to avoid excessive overhead +/// from frequent time checks, defaulting to 25ms worth of data or 1MB, whichever +/// is smaller. +/// +/// # Thread Safety +/// +/// The rate limit can be updated dynamically via `set_mb_per_sec()`. The configuration +/// is protected by a `RwLock`, allowing concurrent reads while ensuring safe updates. +/// If the lock becomes poisoned (due to a panic in another thread), the writer will +/// gracefully degrade by skipping rate limiting rather than propagating the panic. +/// +/// +/// # Special Cases +/// +/// - Setting `mb_per_sec` to `0.0` disables rate limiting entirely +/// - Negative values are rejected with an error +/// - Lock poisoning is handled gracefully by skipping rate limiting +pub struct RateLimitedWriter { + inner: W, + rate_limiter_config: Arc>, + bytes_since_last_pause: usize, + last_pause_time: Instant, +} + +impl RateLimitedWriter { + /// Creates a new rate-limited writer with the specified throughput limit. + /// + /// # Arguments + /// + /// * `inner` - The underlying writer to wrap + /// * `mb_per_sec` - Maximum write rate in megabytes per second (must be non-negative) + /// + /// # Returns + /// + /// Returns `Ok(RateLimitedWriter)` on success, or an error if `mb_per_sec` is negative. + /// + /// + /// # Errors + /// + /// Returns `Err` with `ErrorKind::InvalidInput` if `mb_per_sec` is negative. + pub fn new(inner: W, mb_per_sec: f64) -> Result { + if mb_per_sec < 0.0 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("mbPerSec must be non-negative: got: {}", mb_per_sec), + )); + } + + let min_pause_check_bytes = Self::calculate_min_pause_check_bytes(mb_per_sec); + Ok(Self { + inner, + rate_limiter_config: Arc::new(RwLock::new(RateLimiterConfig { + mb_per_sec, + min_pause_check_bytes, + })), + bytes_since_last_pause: 0, + last_pause_time: Instant::now(), + }) + } + + /// Updates the rate limit dynamically. + /// + /// This method allows changing the throughput limit while the writer is in use. + /// The new rate takes effect immediately for subsequent write operations. + /// + /// # Arguments + /// + /// * `mb_per_sec` - New maximum write rate in megabytes per second (must be non-negative) + /// + /// # Returns + /// + /// Returns `Ok(())` on success, or an error if the rate is invalid or the lock is poisoned. + /// + /// + /// # Errors + /// + /// Returns `Err` with: + /// - `ErrorKind::InvalidInput` if `mb_per_sec` is negative + /// - `ErrorKind::Other` if the internal lock is poisoned + pub fn set_mb_per_sec(&mut self, mb_per_sec: f64) -> Result<()> { + if mb_per_sec < 0.0 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("mbPerSec must be non-negative: got: {}", mb_per_sec), + )); + } + + let min_pause_check_bytes = Self::calculate_min_pause_check_bytes(mb_per_sec); + + let mut config = self.rate_limiter_config.write().map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::Other, + format!("Failed to acquire write lock: {}", e), + ) + })?; + + config.mb_per_sec = mb_per_sec; + config.min_pause_check_bytes = min_pause_check_bytes; + + Ok(()) + } + + /// Calculates the minimum number of bytes to write before checking if a pause is needed. + /// + /// This is based on the configured rate and a minimum pause check interval to avoid + /// excessive overhead from frequent time checks. The result is capped at 1MB. + fn calculate_min_pause_check_bytes(mb_per_sec: f64) -> usize { + let bytes = (MIN_PAUSE_CHECK_MSEC / MSEC_TO_SEC) * mb_per_sec * BYTES_PER_MB; + std::cmp::min(MAX_MIN_PAUSE_CHECK_BYTES, bytes as usize) + } + + /// Pauses execution if the write rate exceeds the configured limit. + /// + /// Calculates the target time for writing the given number of bytes based on + /// the configured rate, and sleeps if insufficient time has elapsed since the + /// last pause. If the lock is poisoned, rate limiting is skipped. + /// + /// # Arguments + /// + /// * `bytes` - Number of bytes written since the last pause + fn pause(&mut self, bytes: usize) { + let config = match self.rate_limiter_config.read() { + Ok(config) => config, + Err(_) => { + // Lock is poisoned, skip rate limiting this time + return; + } + }; + + if config.mb_per_sec == 0.0 { + return; + } + + let elapsed = self.last_pause_time.elapsed().as_secs_f64(); + let target_time = bytes as f64 / (config.mb_per_sec * BYTES_PER_MB); + + if target_time > elapsed { + let sleep_time = Duration::from_secs_f64(target_time - elapsed); + sleep(sleep_time); + } + + self.last_pause_time = Instant::now(); + } +} + +impl Write for RateLimitedWriter { + fn write(&mut self, buf: &[u8]) -> Result { + let n = self.inner.write(buf)?; + self.bytes_since_last_pause += n; + + let current_min_pause_check_bytes = { + match self.rate_limiter_config.read() { + Ok(config) => config.min_pause_check_bytes, + Err(_) => { + // Lock is poisoned, use a safe default + MAX_MIN_PAUSE_CHECK_BYTES + } + } + }; + + if self.bytes_since_last_pause > current_min_pause_check_bytes { + self.pause(self.bytes_since_last_pause); + self.bytes_since_last_pause = 0; + } + Ok(n) + } + + fn flush(&mut self) -> Result<()> { + self.inner.flush() + } +} + + 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 2a80157518ec8..7b46092b09ddc 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 @@ -7,10 +7,13 @@ */ use arrow::array::{Int32Array, StringArray, StructArray}; +use arrow::compute::concat_batches; use arrow::datatypes::{DataType, Field, Schema}; use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; use arrow::record_batch::RecordBatch; -use arrow_array::Array; +use arrow::array::Array; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use std::fs::File; use std::sync::Arc; use tempfile::tempdir; @@ -33,12 +36,16 @@ pub fn cleanup_ffi_schema(schema_ptr: i64) { } pub fn create_test_ffi_data() -> Result<(i64, i64), Box> { + create_test_ffi_data_with_ids(vec![1, 2, 3], vec![Some("Alice"), Some("Bob"), None]) +} + +pub fn create_test_ffi_data_with_ids(ids: Vec, names: Vec>) -> Result<(i64, i64), Box> { let schema = Arc::new(Schema::new(vec![ Field::new("id", DataType::Int32, false), Field::new("name", DataType::Utf8, true), ])); - let id_array = Arc::new(Int32Array::from(vec![1, 2, 3])); - let name_array = Arc::new(StringArray::from(vec![Some("Alice"), Some("Bob"), None])); + let id_array = Arc::new(Int32Array::from(ids)); + let name_array = Arc::new(StringArray::from(names)); let record_batch = RecordBatch::try_new(schema.clone(), vec![id_array, name_array])?; let struct_array = StructArray::from(record_batch); let array_data = struct_array.into_data(); @@ -65,7 +72,16 @@ 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(filename.to_string(), schema_ptr); + let result = NativeParquetWriter::create_writer(filename.to_string(), "test-index".to_string(), schema_ptr, vec![], vec![], vec![]); + assert!(result.is_ok()); + (schema, schema_ptr) +} + +pub fn create_sorted_writer_and_assert_success(filename: &str, sort_column: &str, reverse: bool) -> (Arc, i64) { + let (schema, schema_ptr) = create_test_ffi_schema(); + let result = NativeParquetWriter::create_writer( + filename.to_string(), "test-index".to_string(), schema_ptr, vec![sort_column.to_string()], vec![reverse], vec![false] + ); assert!(result.is_ok()); (schema, schema_ptr) } @@ -105,3 +121,18 @@ pub fn close_writer_and_get_metadata(filename: &str, schema_ptr: i64) -> crate:: cleanup_ffi_schema(schema_ptr); result.unwrap().unwrap() } + +pub fn read_parquet_file(filename: &str) -> Vec { + let file = File::open(filename).unwrap(); + let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + let reader = builder.build().unwrap(); + reader.collect::, _>>().unwrap() +} + +pub fn read_parquet_file_sorted_ids(filename: &str) -> Vec { + let batches = read_parquet_file(filename); + let combined = concat_batches(&batches[0].schema(), &batches).unwrap(); + let id_col = combined.column(0) + .as_any().downcast_ref::().unwrap(); + (0..id_col.len()).map(|i| id_col.value(i)).collect() +} 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 9efcc961be225..25c4e375a7264 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 @@ -6,10 +6,15 @@ * compatible open source license. */ +use std::path::Path; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use std::thread; +use tempfile::tempdir; + use crate::test_utils::*; -use crate::writer::{NativeParquetWriter, WRITER_MANAGER, FILE_MANAGER}; +use crate::writer::NativeParquetWriter; -use parquet::file::reader::FileReader; use std::fs::File; use std::io::Read; @@ -17,8 +22,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!(WRITER_MANAGER.contains_key(&filename)); - assert!(FILE_MANAGER.contains_key(&filename)); + assert!(NativeParquetWriter::has_writer(&filename)); close_writer_and_cleanup_schema(&filename, schema_ptr); } @@ -26,16 +30,15 @@ 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(invalid_path.to_string(), schema_ptr); + let result = NativeParquetWriter::create_writer(invalid_path.to_string(), "test-index".to_string(), schema_ptr, vec![], vec![], vec![]); assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("No such file or directory")); cleanup_ffi_schema(schema_ptr); } #[test] fn test_create_writer_invalid_schema_pointer() { let (_temp_dir, filename) = get_temp_file_path("invalid_schema.parquet"); - let result = NativeParquetWriter::create_writer(filename, 0); + let result = NativeParquetWriter::create_writer(filename, "test-index".to_string(), 0, vec![], vec![], vec![]); assert!(result.is_err()); assert!(result.unwrap_err().to_string().contains("Invalid schema address")); } @@ -44,9 +47,11 @@ fn test_create_writer_invalid_schema_pointer() { 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 result2 = NativeParquetWriter::create_writer(filename.clone(), schema_ptr); + let (_, schema_ptr2) = create_test_ffi_schema(); + let result2 = NativeParquetWriter::create_writer(filename.clone(), "test-index".to_string(), schema_ptr2, vec![], vec![], vec![]); assert!(result2.is_err()); assert!(result2.unwrap_err().to_string().contains("Writer already exists")); + cleanup_ffi_schema(schema_ptr2); close_writer_and_cleanup_schema(&filename, schema_ptr); } @@ -70,6 +75,17 @@ fn test_write_data_no_writer() { cleanup_ffi_data(array_ptr, schema_ptr); } +#[test] +fn test_write_data_multiple_batches() { + let (_temp_dir, filename) = get_temp_file_path("multi_write_ffi.parquet"); + let (_schema, schema_ptr) = create_writer_and_assert_success(&filename); + for _ in 0..3 { + let (array_ptr, data_schema_ptr) = write_ffi_data_to_writer(&filename); + cleanup_ffi_data(array_ptr, data_schema_ptr); + } + close_writer_and_cleanup_schema(&filename, schema_ptr); +} + #[test] fn test_write_data_invalid_pointers() { let (_temp_dir, filename) = get_temp_file_path("invalid_ffi.parquet"); @@ -98,16 +114,11 @@ fn test_write_data_incompatible_schema() { fn test_finalize_writer_success() { let (_temp_dir, filename) = get_temp_file_path("test_close.parquet"); 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()); assert!(result.is_ok()); - let finalize_result = result.unwrap(); - assert!(finalize_result.is_some()); - let finalize_result = finalize_result.unwrap(); - assert_eq!(finalize_result.metadata.file_metadata().num_rows(), 0); - assert!(finalize_result.metadata.file_metadata().version() > 0); - assert!(!WRITER_MANAGER.contains_key(&filename)); - assert!(FILE_MANAGER.contains_key(&filename)); - FILE_MANAGER.remove(&filename); + assert!(Path::new(&filename).exists()); cleanup_ffi_schema(schema_ptr); } @@ -125,9 +136,7 @@ fn test_finalize_writer_with_data_returns_correct_metadata() { let metadata = result.unwrap().unwrap(); assert_eq!(metadata.metadata.file_metadata().num_rows(), 6); assert!(metadata.metadata.file_metadata().version() > 0); - assert_eq!(metadata.metadata.file_metadata().schema_descr().num_columns(), 3); // root + 2 fields (id, name) assert_ne!(metadata.crc32, 0, "CRC32 should be non-zero for a file with data"); - FILE_MANAGER.remove(&filename); cleanup_ffi_schema(schema_ptr); } @@ -142,17 +151,13 @@ fn test_close_nonexistent_writer() { fn test_close_multiple_times_same_file() { let (_temp_dir, filename) = get_temp_file_path("test.parquet"); 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()); assert!(result1.is_ok()); - let metadata = result1.unwrap(); - assert!(metadata.is_some()); - assert_eq!(metadata.unwrap().metadata.num_rows, 0); - assert!(!WRITER_MANAGER.contains_key(&filename)); - assert!(FILE_MANAGER.contains_key(&filename)); - let result2 = NativeParquetWriter::finalize_writer(filename.clone()); + let result2 = NativeParquetWriter::finalize_writer(filename); assert!(result2.is_err()); assert!(result2.unwrap_err().to_string().contains("Writer not found")); - FILE_MANAGER.remove(&filename); cleanup_ffi_schema(schema_ptr); } @@ -160,18 +165,115 @@ fn test_close_multiple_times_same_file() { fn test_sync_to_disk_success() { let (_temp_dir, filename) = get_temp_file_path("test_flush.parquet"); let (_schema, schema_ptr) = create_writer_and_assert_success(&filename); - assert!(FILE_MANAGER.contains_key(&filename)); - let result = NativeParquetWriter::sync_to_disk(filename.clone()); + let (array_ptr, data_schema_ptr) = write_ffi_data_to_writer(&filename); + cleanup_ffi_data(array_ptr, data_schema_ptr); + let _ = NativeParquetWriter::finalize_writer(filename.clone()); + let result = NativeParquetWriter::sync_to_disk(filename); assert!(result.is_ok()); - assert!(!FILE_MANAGER.contains_key(&filename)); - close_writer_and_cleanup_schema(&filename, schema_ptr); + cleanup_ffi_schema(schema_ptr); } #[test] fn test_flush_nonexistent_file() { let result = NativeParquetWriter::sync_to_disk("nonexistent.parquet".to_string()); assert!(result.is_err()); - assert_eq!(result.unwrap_err().to_string(), "File not found"); +} + +#[test] +fn test_complete_writer_lifecycle() { + let (_temp_dir, filename) = get_temp_file_path("complete_workflow.parquet"); + let file_path = Path::new(&filename); + let (_schema, schema_ptr) = create_writer_and_assert_success(&filename); + + for _ in 0..3 { + let (array_ptr, data_schema_ptr) = write_ffi_data_to_writer(&filename); + cleanup_ffi_data(array_ptr, data_schema_ptr); + } + + let close_result = NativeParquetWriter::finalize_writer(filename.clone()); + assert!(close_result.is_ok()); + assert!(close_result.unwrap().is_some()); + + assert!(NativeParquetWriter::sync_to_disk(filename.clone()).is_ok()); + assert!(file_path.exists()); + assert!(file_path.metadata().unwrap().len() > 0); + + cleanup_ffi_schema(schema_ptr); +} + +#[test] +fn test_sorted_writer_ascending() { + let (_temp_dir, filename) = get_temp_file_path("sorted_asc.parquet"); + let (_schema, schema_ptr) = create_sorted_writer_and_assert_success(&filename, "id", false); + + 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(); + 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(); + cleanup_ffi_data(ap2, sp2); + + NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + + let ids = read_parquet_file_sorted_ids(&filename); + assert_eq!(ids, vec![10, 20, 30, 40, 50], "Data should be sorted ascending by id"); + + cleanup_ffi_schema(schema_ptr); +} + +#[test] +fn test_sorted_writer_descending() { + let (_temp_dir, filename) = get_temp_file_path("sorted_desc.parquet"); + let (_schema, schema_ptr) = create_sorted_writer_and_assert_success(&filename, "id", true); + + 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(); + 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(); + cleanup_ffi_data(ap2, sp2); + + NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + + let ids = read_parquet_file_sorted_ids(&filename); + assert_eq!(ids, vec![50, 40, 30, 20, 10], "Data should be sorted descending by id"); + + cleanup_ffi_schema(schema_ptr); +} + +#[test] +fn test_unsorted_writer_preserves_insertion_order() { + let (_temp_dir, filename) = get_temp_file_path("unsorted.parquet"); + let (_schema, schema_ptr) = create_writer_and_assert_success(&filename); + + 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(); + 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(); + cleanup_ffi_data(ap2, sp2); + + NativeParquetWriter::finalize_writer(filename.clone()).unwrap(); + + let ids = read_parquet_file_sorted_ids(&filename); + assert_eq!(ids, vec![30, 10, 50, 20, 40], "Data should preserve insertion order"); + + cleanup_ffi_schema(schema_ptr); } #[test] @@ -183,68 +285,30 @@ fn test_get_filtered_writer_memory_usage_with_writers() { let (_schema2, schema_ptr2) = create_writer_and_assert_success(&filename2); let result = NativeParquetWriter::get_filtered_writer_memory_usage(prefix); assert!(result.is_ok()); - let _memory_usage = result.unwrap(); - assert!(_memory_usage >= 0); + assert!(result.unwrap() >= 0); close_writer_and_cleanup_schema(&filename1, schema_ptr1); close_writer_and_cleanup_schema(&filename2, schema_ptr2); } +// CRC32 tests -/// Computes CRC32 of a file by reading it from disk in chunks. -/// This is the "re-read" baseline that the streaming checksum must match. fn compute_file_crc32(path: &str) -> u32 { let mut file = File::open(path).unwrap(); let mut hasher = crc32fast::Hasher::new(); let mut buf = [0u8; 64 * 1024]; loop { let n = file.read(&mut buf).unwrap(); - if n == 0 { - break; - } + if n == 0 { break; } hasher.update(&buf[..n]); } hasher.finalize() } -/// Verifies that the streaming CRC32 computed during write (via Crc32Writer) -/// exactly matches a CRC32 computed by re-reading the finalized file from disk. -/// -/// This proves the streaming approach is correct and eliminates the need for -/// a second I/O pass over the file. -#[test] -fn test_streaming_crc32_matches_reread_crc32_empty_file() { - let (_temp_dir, filename) = get_temp_file_path("crc32_empty.parquet"); - let (_schema, schema_ptr) = create_writer_and_assert_success(&filename); - - // Finalize with zero rows — still writes the Parquet magic bytes + footer - let result = NativeParquetWriter::finalize_writer(filename.clone()); - assert!(result.is_ok()); - let finalize_result = result.unwrap().unwrap(); - let streaming_crc32 = finalize_result.crc32; - - // Re-read the file and compute CRC32 independently - let reread_crc32 = compute_file_crc32(&filename); - - assert_eq!( - streaming_crc32, reread_crc32, - "Streaming CRC32 ({:#010x}) must match re-read CRC32 ({:#010x}) for empty Parquet file", - streaming_crc32, reread_crc32 - ); - assert_ne!(streaming_crc32, 0, "CRC32 should be non-zero even for an empty Parquet file (magic bytes + footer)"); - - FILE_MANAGER.remove(&filename); - cleanup_ffi_schema(schema_ptr); -} - -/// Verifies streaming CRC32 matches re-read CRC32 for a file with actual data. -/// Writes multiple batches to exercise the full write path (row groups, column -/// chunks, compression, bloom filters, footer). #[test] -fn test_streaming_crc32_matches_reread_crc32_with_data() { +fn test_crc32_matches_reread_with_data() { let (_temp_dir, filename) = get_temp_file_path("crc32_with_data.parquet"); let (_schema, schema_ptr) = create_writer_and_assert_success(&filename); - // Write 3 batches (9 rows total) to exercise multiple write() calls 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(); @@ -256,41 +320,23 @@ fn test_streaming_crc32_matches_reread_crc32_with_data() { let finalize_result = result.unwrap().unwrap(); let streaming_crc32 = finalize_result.crc32; - // Verify metadata is correct assert_eq!(finalize_result.metadata.file_metadata().num_rows(), 9); - // Re-read the file and compute CRC32 independently let reread_crc32 = compute_file_crc32(&filename); + assert_eq!(streaming_crc32, reread_crc32); + assert_ne!(streaming_crc32, 0); - assert_eq!( - streaming_crc32, reread_crc32, - "Streaming CRC32 ({:#010x}) must match re-read CRC32 ({:#010x}) for Parquet file with {} rows", - streaming_crc32, reread_crc32, finalize_result.metadata.file_metadata().num_rows() - ); - assert_ne!(streaming_crc32, 0, "CRC32 should be non-zero for a file with data"); - - // Verify the file is a valid Parquet file by reading it back - let file = File::open(&filename).unwrap(); - let reader = parquet::file::reader::SerializedFileReader::new(file).unwrap(); - assert_eq!(reader.metadata().file_metadata().num_rows(), 9); - - FILE_MANAGER.remove(&filename); cleanup_ffi_schema(schema_ptr); } -/// Verifies that two different files produce different CRC32 values, -/// confirming the checksum is content-dependent and not a constant. #[test] -fn test_streaming_crc32_differs_for_different_content() { - // File 1: empty +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 crc32_empty = result1.unwrap().unwrap().crc32; - FILE_MANAGER.remove(&filename1); cleanup_ffi_schema(schema_ptr1); - // File 2: with data 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(); @@ -298,12 +344,148 @@ fn test_streaming_crc32_differs_for_different_content() { cleanup_ffi_data(array_ptr, data_schema_ptr); let result2 = NativeParquetWriter::finalize_writer(filename2.clone()); let crc32_with_data = result2.unwrap().unwrap().crc32; - FILE_MANAGER.remove(&filename2); cleanup_ffi_schema(schema_ptr2); - assert_ne!( - crc32_empty, crc32_with_data, - "Empty file CRC32 ({:#010x}) should differ from file-with-data CRC32 ({:#010x})", - crc32_empty, crc32_with_data - ); + assert_ne!(crc32_empty, crc32_with_data); +} + +// Concurrency tests + +#[test] +fn test_concurrent_writer_creation() { + let temp_dir = tempdir().unwrap(); + let success_count = Arc::new(AtomicUsize::new(0)); + let mut handles = vec![]; + + for i in 0..10 { + let temp_dir_path = temp_dir.path().to_path_buf(); + let success_count = Arc::clone(&success_count); + + let handle = thread::spawn(move || { + 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(filename.clone(), "test-index".to_string(), schema_ptr, vec![], vec![], vec![]).is_ok() { + success_count.fetch_add(1, Ordering::SeqCst); + let (ap, sp) = create_test_ffi_data().unwrap(); + let _ = NativeParquetWriter::write_data(filename.clone(), ap, sp); + cleanup_ffi_data(ap, sp); + let _ = NativeParquetWriter::finalize_writer(filename); + } + cleanup_ffi_schema(schema_ptr); + }); + handles.push(handle); + } + + for handle in handles { + handle.join().unwrap(); + } + + assert_eq!(success_count.load(Ordering::SeqCst), 10); +} + +#[test] +fn test_concurrent_close_operations_same_file() { + let (_temp_dir, filename) = get_temp_file_path("close_race.parquet"); + 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 success_count = Arc::new(AtomicUsize::new(0)); + let mut handles = vec![]; + + for _ in 0..3 { + let filename = filename.clone(); + let success_count = Arc::clone(&success_count); + + let handle = thread::spawn(move || { + if NativeParquetWriter::finalize_writer(filename).is_ok() { + success_count.fetch_add(1, Ordering::SeqCst); + } + }); + handles.push(handle); + } + + for handle in handles { + handle.join().unwrap(); + } + + assert_eq!(success_count.load(Ordering::SeqCst), 1); + cleanup_ffi_schema(schema_ptr); +} + +#[test] +fn test_concurrent_writes_same_file() { + let (_temp_dir, filename) = get_temp_file_path("concurrent_write_ffi.parquet"); + let (_schema, schema_ptr) = create_writer_and_assert_success(&filename); + + let success_count = Arc::new(AtomicUsize::new(0)); + let mut handles = vec![]; + + for _ in 0..5 { + let filename = filename.clone(); + 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() { + success_count.fetch_add(1, Ordering::SeqCst); + } + cleanup_ffi_data(array_ptr, data_schema_ptr); + }); + handles.push(handle); + } + + for handle in handles { + handle.join().unwrap(); + } + + assert_eq!(success_count.load(Ordering::SeqCst), 5); + close_writer_and_cleanup_schema(&filename, schema_ptr); +} + +#[test] +fn test_concurrent_writes_different_files() { + let temp_dir = tempdir().unwrap(); + let file_count = 8; + let success_count = Arc::new(AtomicUsize::new(0)); + let mut handles = vec![]; + let mut filenames = vec![]; + let mut schema_ptrs = vec![]; + + for i in 0..file_count { + let file_path = temp_dir.path().join(format!("concurrent_write_{}.parquet", i)); + let filename = file_path.to_string_lossy().to_string(); + let (_schema, schema_ptr) = create_writer_and_assert_success(&filename); + filenames.push(filename); + schema_ptrs.push(schema_ptr); + } + + for i in 0..file_count { + let filename = filenames[i].clone(); + let success_count = Arc::clone(&success_count); + + 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() { + success_count.fetch_add(1, Ordering::SeqCst); + } + cleanup_ffi_data(array_ptr, data_schema_ptr); + } + }); + handles.push(handle); + } + + for handle in handles { + handle.join().unwrap(); + } + + assert_eq!(success_count.load(Ordering::SeqCst), file_count * 2); + + for (i, filename) in filenames.iter().enumerate() { + close_writer_and_cleanup_schema(filename, schema_ptrs[i]); + } } 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 36bb2fe795d7d..214316fe16db5 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 @@ -8,74 +8,85 @@ use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; use arrow::record_batch::RecordBatch; +use arrow::compute::{concat_batches, lexsort_to_indices, take, SortColumn}; use dashmap::DashMap; use lazy_static::lazy_static; -use parquet::arrow::ArrowWriter; -use parquet::basic::Compression; -use parquet::file::properties::WriterProperties; +use parquet::arrow::{arrow_reader::ParquetRecordBatchReaderBuilder, ArrowWriter}; use parquet::file::reader::{FileReader, SerializedFileReader}; use std::fs::File; -use std::io::Write; +use std::io::Read; +use std::path::Path; use std::sync::{Arc, Mutex}; -use crate::{log_error, log_debug}; - -/// A write wrapper that computes CRC32 as bytes flow through. -/// Wraps a File and tracks the running checksum without buffering. -pub struct Crc32Writer { - inner: File, - hasher: crc32fast::Hasher, -} - -impl Crc32Writer { - fn new(file: File) -> Self { - Self { - inner: file, - hasher: crc32fast::Hasher::new(), - } - } - - /// Finalizes and returns the CRC32 checksum of all bytes written. - fn checksum(&self) -> u32 { - self.hasher.clone().finalize() - } -} - -impl Write for Crc32Writer { - fn write(&mut self, buf: &[u8]) -> std::io::Result { - let n = self.inner.write(buf)?; - self.hasher.update(&buf[..n]); - Ok(n) - } - - fn flush(&mut self) -> std::io::Result<()> { - self.inner.flush() - } -} +use crate::{log_info, log_error, log_debug}; +use crate::merge::schema::ROW_ID_COLUMN_NAME; +use crate::native_settings::NativeSettings; +use crate::writer_properties_builder::WriterPropertiesBuilder; /// Result from finalizing a writer: Parquet metadata + whole-file CRC32. +#[derive(Debug)] pub struct FinalizeResult { pub metadata: parquet::file::metadata::ParquetMetaData, pub crc32: u32, } +/// 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 { + writer: Arc>>, + file_handle: File, + index_name: String, + sort_columns: Vec, + reverse_sorts: Vec, + nulls_first: Vec, +} + lazy_static! { - pub static ref WRITER_MANAGER: DashMap>>> = DashMap::new(); - pub static ref FILE_MANAGER: DashMap = DashMap::new(); + /// Unified per-writer registry. Keyed by temp filename. + static ref WRITERS: DashMap = DashMap::new(); + pub static ref SETTINGS_STORE: DashMap = DashMap::new(); } pub struct NativeParquetWriter; impl NativeParquetWriter { - pub fn create_writer(filename: String, schema_address: i64) -> Result<(), Box> { - log_debug!("create_writer called for file: {}, schema_address: {}", filename, schema_address); + /// 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); + path.parent() + .unwrap_or_else(|| Path::new("")) + .join(format!("temp-{}", path.file_name().unwrap().to_str().unwrap())) + .to_string_lossy() + .to_string() + } + + pub fn create_writer( + filename: String, + index_name: String, + schema_address: i64, + sort_columns: Vec, + reverse_sorts: Vec, + nulls_first: Vec, + ) -> Result<(), Box> { + log_info!( + "create_writer called for file: {}, index: {}, schema_address: {}, sort_columns: {:?}, reverse_sorts: {:?}, nulls_first: {:?}", + filename, index_name, schema_address, sort_columns, reverse_sorts, nulls_first + ); if (schema_address as *mut u8).is_null() { log_error!("ERROR: Invalid schema address (null pointer) for file: {}", filename); return Err("Invalid schema address".into()); } - if WRITER_MANAGER.contains_key(&filename) { - log_error!("ERROR: Writer already exists for file: {}", filename); + + 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()); } @@ -83,27 +94,47 @@ impl NativeParquetWriter { let schema = Arc::new(arrow::datatypes::Schema::try_from(&arrow_schema)?); log_debug!("Schema created with {} fields", schema.fields().len()); - let file = File::create(&filename)?; + let file = File::create(&temp_filename)?; let file_clone = file.try_clone()?; - FILE_MANAGER.insert(filename.clone(), file_clone); - - let props = WriterProperties::builder() - .set_compression(Compression::LZ4_RAW) - .set_bloom_filter_enabled(true) - .set_bloom_filter_fpp(0.1) - .set_bloom_filter_ndv(100000) - .build(); - let crc_writer = Crc32Writer::new(file); - let writer = ArrowWriter::try_new(crc_writer, schema, Some(props))?; - WRITER_MANAGER.insert(filename, Arc::new(Mutex::new(writer))); + + let config: NativeSettings = SETTINGS_STORE + .get(&index_name) + .map(|r| r.clone()) + .unwrap_or_default(); + let props = WriterPropertiesBuilder::build(&config); + + SETTINGS_STORE.entry(index_name.clone()).and_modify(|s| { + s.sort_columns = sort_columns.clone(); + s.reverse_sorts = reverse_sorts.clone(); + s.nulls_first = nulls_first.clone(); + }).or_insert_with(|| { + let mut s = NativeSettings::default(); + s.sort_columns = sort_columns.clone(); + s.reverse_sorts = reverse_sorts.clone(); + s.nulls_first = nulls_first.clone(); + s + }); + + let writer = ArrowWriter::try_new(file, schema, Some(props))?; + + WRITERS.insert(temp_filename, WriterState { + writer: Arc::new(Mutex::new(writer)), + file_handle: file_clone, + index_name, + sort_columns, + reverse_sorts, + nulls_first, + }); + Ok(()) } pub fn write_data(filename: String, array_address: i64, schema_address: i64) -> Result<(), Box> { - log_debug!("write_data called for file: {}", filename); + let temp_filename = Self::temp_filename(&filename); + log_debug!("write_data called for file: {} (temp: {})", filename, 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: {}", filename); + log_error!("ERROR: Invalid FFI addresses for file: {}", temp_filename); return Err("Invalid FFI addresses (null pointers)".into()); } @@ -118,12 +149,12 @@ impl NativeParquetWriter { let record_batch = RecordBatch::try_new(schema, struct_array.columns().to_vec())?; log_debug!("Created RecordBatch with {} rows and {} columns", record_batch.num_rows(), record_batch.num_columns()); - if let Some(writer_arc) = WRITER_MANAGER.get(&filename) { - let mut writer = writer_arc.lock().unwrap(); + if let Some(state) = WRITERS.get(&temp_filename) { + let mut writer = state.writer.lock().unwrap(); writer.write(&record_batch)?; Ok(()) } else { - log_error!("ERROR: No writer found for file: {}", filename); + log_error!("ERROR: No writer found for temp file: {}", temp_filename); Err("Writer not found".into()) } } else { @@ -134,50 +165,291 @@ impl NativeParquetWriter { } pub fn finalize_writer(filename: String) -> Result, Box> { - log_debug!("finalize_writer called for file: {}", filename); + let temp_filename = Self::temp_filename(&filename); + log_info!("finalize_writer called for file: {} (temp: {})", filename, temp_filename); - if let Some((_, writer_arc)) = WRITER_MANAGER.remove(&filename) { + if let Some((_, state)) = WRITERS.remove(&temp_filename) { + let WriterState { writer: writer_arc, file_handle: _file, index_name, sort_columns, reverse_sorts, nulls_first } = state; match Arc::try_unwrap(writer_arc) { Ok(mutex) => { - let mut writer = mutex.into_inner().unwrap(); - let parquet_metadata = writer.finish()?; - let file_metadata = parquet_metadata.file_metadata(); - log_debug!("Successfully finalized writer for file: {}, num_rows={}", filename, file_metadata.num_rows()); - let crc32 = writer.inner().checksum(); - log_debug!("CRC32 for file {}: {:#010x}", filename, crc32); - Ok(Some(FinalizeResult { metadata: parquet_metadata, crc32 })) + let writer = mutex.into_inner().unwrap(); + match writer.close() { + Ok(_) => { + log_info!("Successfully closed temp writer for: {}", temp_filename); + // _file is dropped here, closing the file handle + + Self::sort_and_rewrite_parquet(&temp_filename, &filename, &index_name, &sort_columns, &reverse_sorts, &nulls_first)?; + + let _ = std::fs::remove_file(&temp_filename); + + // Compute CRC32 by reading the final sorted file + let crc32 = Self::compute_file_crc32(&filename)?; + log_debug!("CRC32 for file {}: {:#010x}", filename, crc32); + + // Read full ParquetMetaData from the final file + let file = File::open(&filename)?; + let reader = SerializedFileReader::new(file)?; + let parquet_metadata = reader.metadata().clone(); + + Ok(Some(FinalizeResult { metadata: parquet_metadata, crc32 })) + } + 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 file: {}", filename); + 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 file: {}", filename); + log_error!("ERROR: Writer not found for temp file: {}", temp_filename); Err("Writer not found".into()) } } + fn compute_file_crc32(path: &str) -> Result> { + let mut file = File::open(path)?; + let mut hasher = crc32fast::Hasher::new(); + let mut buf = [0u8; 64 * 1024]; + loop { + let n = file.read(&mut buf)?; + if n == 0 { break; } + hasher.update(&buf[..n]); + } + Ok(hasher.finalize()) + } + + fn sort_and_rewrite_parquet( + temp_filename: &str, + output_filename: &str, + index_name: &str, + sort_columns: &[String], + reverse_sorts: &[bool], + nulls_first: &[bool], + ) -> Result<(), Box> { + log_info!( + "sort_and_rewrite_parquet: temp={}, output={}, sort_columns={:?}, reverse_sorts={:?}, nulls_first={:?}", + temp_filename, output_filename, sort_columns, reverse_sorts, nulls_first + ); + + if sort_columns.is_empty() { + log_info!("No sort columns specified, renaming temp file to final"); + std::fs::rename(temp_filename, output_filename)?; + return Ok(()); + } + + let file_size = std::fs::metadata(temp_filename)?.len(); + const MAX_MEMORY_SIZE: u64 = 32 * 1024 * 1024; + + if file_size <= MAX_MEMORY_SIZE { + Self::sort_small_file(temp_filename, output_filename, index_name, sort_columns, reverse_sorts, nulls_first) + } else { + Self::sort_large_file(temp_filename, output_filename, index_name, sort_columns, reverse_sorts, nulls_first) + } + } + + /// In-memory sort for small files: read all batches, concat, sort, rewrite row IDs, write. + fn sort_small_file( + temp_filename: &str, + output_filename: &str, + index_name: &str, + sort_columns: &[String], + reverse_sorts: &[bool], + nulls_first: &[bool], + ) -> Result<(), Box> { + log_info!("Using in-memory sort for small file: {}", temp_filename); + + let file = File::open(temp_filename)?; + let builder = ParquetRecordBatchReaderBuilder::try_new(file)?; + let arrow_reader = builder.with_batch_size(2048).build()?; + + let mut batches = Vec::new(); + for batch_result in arrow_reader { + batches.push(batch_result?); + } + + if batches.is_empty() { + log_info!("No data to sort in file: {}", temp_filename); + std::fs::rename(temp_filename, output_filename)?; + return Ok(()); + } + + let schema = batches[0].schema(); + let combined_batch = concat_batches(&schema, &batches)?; + let sorted_batch = Self::sort_batch(&combined_batch, sort_columns, reverse_sorts, nulls_first)?; + let final_batch = Self::rewrite_row_ids(&sorted_batch, &schema)?; + + Self::write_final_file(output_filename, index_name, &final_batch, schema)?; + Ok(()) + } + + /// For large files: read in batches, sort each batch individually, write each + /// as a separate sorted chunk file, then use the streaming k-way merge to + /// produce the final globally-sorted output. + fn sort_large_file( + temp_filename: &str, + output_filename: &str, + index_name: &str, + sort_columns: &[String], + reverse_sorts: &[bool], + nulls_first: &[bool], + ) -> Result<(), Box> { + log_info!("Using streaming merge sort for large file: {}", temp_filename); + + let file = File::open(temp_filename)?; + let builder = ParquetRecordBatchReaderBuilder::try_new(file)?; + let arrow_reader = builder.with_batch_size(8192).build()?; + + let mut chunk_paths: Vec = Vec::new(); + let mut batch_count = 0; + let temp_dir = std::env::temp_dir(); + + for batch_result in arrow_reader { + let batch = batch_result?; + let schema = batch.schema(); + let sorted_batch = Self::sort_batch(&batch, sort_columns, reverse_sorts, nulls_first)?; + + let chunk_filename = temp_dir + .join(format!("sort_chunk_{}_{}.parquet", batch_count, std::process::id())) + .to_string_lossy() + .to_string(); + Self::write_final_file(&chunk_filename, index_name, &sorted_batch, schema)?; + + chunk_paths.push(chunk_filename); + batch_count += 1; + } + + if chunk_paths.is_empty() { + log_info!("No data to sort in file: {}", temp_filename); + std::fs::rename(temp_filename, output_filename)?; + return Ok(()); + } + + log_info!("Created {} sorted chunks, merging via streaming k-way merge", batch_count); + + // Use the streaming merge to produce the final sorted file + crate::merge::merge_sorted( + &chunk_paths, + output_filename, + index_name, + sort_columns, + reverse_sorts, + nulls_first, + ).map_err(|e| -> Box { format!("Streaming merge failed: {}", e).into() })?; + + // Clean up temp chunk files + for path in &chunk_paths { + let _ = std::fs::remove_file(path); + } + + Ok(()) + } + + fn sort_batch( + batch: &RecordBatch, + sort_columns: &[String], + reverse_sorts: &[bool], + nulls_first: &[bool], + ) -> Result> { + let columns: Vec = sort_columns + .iter() + .enumerate() + .map(|(i, col_name)| { + let reverse = reverse_sorts.get(i).copied().unwrap_or(false); + let nf = nulls_first.get(i).copied().unwrap_or(false); + let options = arrow::compute::SortOptions { + descending: reverse, + nulls_first: nf, + }; + let col_index = batch.schema().index_of(col_name) + .map_err(|_| format!("Sort column '{}' not found in schema", col_name))?; + Ok(SortColumn { + values: batch.column(col_index).clone(), + options: Some(options), + }) + }) + .collect::, Box>>()?; + + let indices = lexsort_to_indices(&columns, None)?; + let sorted_columns: Result, _> = batch + .columns() + .iter() + .map(|col| take(col.as_ref(), &indices, None)) + .collect(); + + Ok(RecordBatch::try_new(batch.schema(), sorted_columns?)?) + } + + /// If a ___row_id column exists, rewrite it with sequential values 0..N. + fn rewrite_row_ids( + batch: &RecordBatch, + schema: &Arc, + ) -> Result> { + use arrow::array::Int64Array; + + if let Some(row_id_idx) = schema.fields().iter().position(|f| f.name() == ROW_ID_COLUMN_NAME) { + log_info!("Rewriting ___row_id column with sequential values 0..{}", batch.num_rows()); + let sequential_ids = Int64Array::from_iter_values( + (0..batch.num_rows() as u64).map(|x| x as i64) + ); + let mut new_columns = batch.columns().to_vec(); + new_columns[row_id_idx] = Arc::new(sequential_ids); + Ok(RecordBatch::try_new(schema.clone(), new_columns)?) + } else { + Ok(batch.clone()) + } + } + + fn write_final_file( + output_filename: &str, + index_name: &str, + batch: &RecordBatch, + schema: Arc, + ) -> Result<(), Box> { + let config = SETTINGS_STORE + .get(index_name) + .map(|r| r.clone()) + .unwrap_or_default(); + let props = WriterPropertiesBuilder::build(&config); + let file = File::create(output_filename)?; + let mut writer = ArrowWriter::try_new(file, schema, Some(props))?; + writer.write(batch)?; + writer.close()?; + log_info!("Successfully wrote final file: {}", output_filename); + Ok(()) + } + pub fn sync_to_disk(filename: String) -> Result<(), Box> { log_debug!("sync_to_disk called for file: {}", filename); - if let Some(file) = FILE_MANAGER.get_mut(&filename) { - file.sync_all()?; - log_debug!("Successfully fsynced file: {}", filename); - drop(file); - FILE_MANAGER.remove(&filename); - Ok(()) - } else { - log_error!("ERROR: File not found for fsync: {}", filename); - Err("File not found".into()) + let file = match File::open(&filename) { + Ok(f) => f, + Err(e) => { + log_error!("ERROR: Failed to open file for fsync: {}", filename); + return Err(e.into()); + } + }; + + match file.sync_all() { + Ok(_) => { + log_debug!("Successfully fsynced file: {}", filename); + Ok(()) + } + Err(e) => { + log_error!("ERROR: Failed to fsync file: {}", filename); + Err(e.into()) + } } } pub fn get_filtered_writer_memory_usage(path_prefix: String) -> Result> { let mut total_memory = 0; - for entry in WRITER_MANAGER.iter() { + for entry in WRITERS.iter() { if entry.key().starts_with(&path_prefix) { - if let Ok(writer) = entry.value().lock() { + if let Ok(writer) = entry.value().writer.lock() { total_memory += writer.memory_size(); } } diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/writer_properties_builder.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/writer_properties_builder.rs new file mode 100644 index 0000000000000..8af048cb2290a --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/writer_properties_builder.rs @@ -0,0 +1,198 @@ +/* + * 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. + */ + +use parquet::basic::{Compression, ZstdLevel, GzipLevel, BrotliLevel}; +use parquet::file::properties::WriterProperties; + +use crate::native_settings::NativeSettings; + +/// Builder for converting NativeSettings into Parquet WriterProperties. +/// +/// This struct follows the Single Responsibility Principle by focusing +/// solely on the conversion logic from configuration to Parquet properties. +/// +/// # Design Principles +/// +/// - **Single Responsibility**: Only handles WriterProperties construction +/// - **Open/Closed**: Can be extended with new compression types without modification +/// - **Dependency Inversion**: Depends on NativeSettings abstraction +pub struct WriterPropertiesBuilder; + +impl WriterPropertiesBuilder { + /// Builds WriterProperties from a NativeSettings. + /// + /// This method applies both index-level and field-level configurations + /// to create a complete WriterProperties instance for Parquet writing. + /// + /// # Arguments + /// + /// * `config` - The native settings to convert + /// + /// # Returns + /// + /// A fully configured WriterProperties instance + pub fn build(config: &NativeSettings) -> WriterProperties { + let mut builder = WriterProperties::builder(); + + // Apply compression settings + builder = Self::apply_compression_settings(builder, config); + + // Apply page settings + builder = Self::apply_page_settings(builder, config); + + // Apply row group settings + builder = Self::apply_row_group_settings(builder, config); + + // Apply dictionary settings + builder = Self::apply_dictionary_settings(builder, config); + + // Apply bloom filter settings + builder = Self::apply_bloom_filter_settings(builder, config); + + // Apply field-level configurations + builder = Self::apply_field_configs(builder, config); + + builder.build() + } + + /// Applies compression settings to the builder. + fn apply_compression_settings( + mut builder: parquet::file::properties::WriterPropertiesBuilder, + config: &NativeSettings + ) -> parquet::file::properties::WriterPropertiesBuilder { + let compression = Self::parse_compression_type( + config.get_compression_type(), + config.get_compression_level() + ); + builder = builder.set_compression(compression); + builder + } + + /// Applies page size and row limit settings. + fn apply_page_settings( + mut builder: parquet::file::properties::WriterPropertiesBuilder, + config: &NativeSettings + ) -> parquet::file::properties::WriterPropertiesBuilder { + builder = builder.set_data_page_size_limit(config.get_page_size_bytes()); + builder = builder.set_data_page_row_count_limit(config.get_page_row_limit()); + builder + } + + /// Applies row group size and row count settings. + fn apply_row_group_settings( + builder: parquet::file::properties::WriterPropertiesBuilder, + config: &NativeSettings + ) -> parquet::file::properties::WriterPropertiesBuilder { + builder + .set_max_row_group_size(config.get_row_group_size_bytes()) + } + + /// Applies dictionary encoding settings. + fn apply_dictionary_settings( + mut builder: parquet::file::properties::WriterPropertiesBuilder, + config: &NativeSettings + ) -> parquet::file::properties::WriterPropertiesBuilder { + builder = builder.set_dictionary_page_size_limit(config.get_dict_size_bytes()); + builder + } + + /// Applies bloom filter settings. + fn apply_bloom_filter_settings( + mut builder: parquet::file::properties::WriterPropertiesBuilder, + config: &NativeSettings + ) -> parquet::file::properties::WriterPropertiesBuilder { + builder = builder.set_bloom_filter_enabled(config.get_bloom_filter_enabled()); + builder = builder.set_bloom_filter_fpp(config.get_bloom_filter_fpp()); + builder = builder.set_bloom_filter_ndv(config.get_bloom_filter_ndv()); + builder + } + + /// Applies field-level configurations. + fn apply_field_configs( + mut builder: parquet::file::properties::WriterPropertiesBuilder, + config: &NativeSettings + ) -> parquet::file::properties::WriterPropertiesBuilder { + if let Some(field_configs) = &config.field_configs { + for (field_name, field_config) in field_configs { + if let Some(compression_type) = &field_config.compression_type { + let compression = Self::parse_compression_type( + compression_type, + field_config.compression_level.unwrap_or(3) + ); + builder = builder.set_column_compression(field_name.clone().into(), compression); + } + } + } + builder + } + + /// Parses compression type string to Parquet Compression enum. + /// + /// # Arguments + /// + /// * `compression_type` - String representation of compression type + /// * `level` - Compression level (algorithm-dependent) + /// + /// # Returns + /// + /// Appropriate Compression enum variant + fn parse_compression_type(compression_type: &str, level: i32) -> Compression { + match compression_type.to_uppercase().as_str() { + "ZSTD" => Compression::ZSTD( + ZstdLevel::try_new(level).unwrap_or(ZstdLevel::default()) + ), + "SNAPPY" => Compression::SNAPPY, + "GZIP" => Compression::GZIP( + GzipLevel::try_new(level as u32).unwrap_or_default() + ), + "LZ4" => Compression::LZ4, + "BROTLI" => Compression::BROTLI( + BrotliLevel::try_new(level as u32).unwrap_or_default() + ), + "LZ4_RAW" => Compression::LZ4_RAW, + "UNCOMPRESSED" => Compression::UNCOMPRESSED, + _ => Compression::ZSTD(ZstdLevel::try_new(level).unwrap_or(ZstdLevel::default())) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::native_settings::NativeSettings; + + #[test] + fn test_build_with_compression() { + let config = NativeSettings { + compression_type: Some("ZSTD".to_string()), + compression_level: Some(5), + ..Default::default() + }; + + let props = WriterPropertiesBuilder::build(&config); + assert_ne!(props.compression(&parquet::schema::types::ColumnPath::from("test")), Compression::UNCOMPRESSED); + } + + #[test] + fn test_parse_compression_types() { + assert!(matches!( + WriterPropertiesBuilder::parse_compression_type("ZSTD", 3), + Compression::ZSTD(_) + )); + + assert!(matches!( + WriterPropertiesBuilder::parse_compression_type("SNAPPY", 0), + Compression::SNAPPY + )); + + assert!(matches!( + WriterPropertiesBuilder::parse_compression_type("GZIP", 6), + Compression::GZIP(_) + )); + } +} 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 new file mode 100644 index 0000000000000..c056071588835 --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/rust/tests/merge_integration_tests.rs @@ -0,0 +1,175 @@ +/* + * 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. + */ + +use arrow::array::{Array, PrimitiveArray}; +use arrow::array::types::TimestampMillisecondType; +use opensearch_parquet_format::merge::{merge_sorted, merge_unsorted}; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use parquet::file::reader::{FileReader, SerializedFileReader}; +use std::fs::File; +use std::path::Path; +use tempfile::tempdir; + +/// Helper: collect all parquet files in a directory (sorted by name). +fn list_parquet_files(dir: &str) -> Vec { + let mut files: Vec = std::fs::read_dir(dir) + .expect("cannot read directory") + .filter_map(|e| { + let p = e.ok()?.path(); + if p.extension().and_then(|s| s.to_str()) == Some("parquet") + && !p.file_name()?.to_str()?.starts_with("merged") + { + Some(p.to_string_lossy().to_string()) + } else { + None + } + }) + .collect(); + files.sort(); + files +} + +/// Helper: count total rows across input files. +fn count_rows_in_files(files: &[String]) -> i64 { + files + .iter() + .map(|f| { + let reader = SerializedFileReader::new(File::open(f).unwrap()).unwrap(); + reader.metadata().file_metadata().num_rows() + }) + .sum() +} + +/// Helper: count rows in a single parquet file. +fn count_rows(path: &str) -> i64 { + let reader = SerializedFileReader::new(File::open(path).unwrap()).unwrap(); + reader.metadata().file_metadata().num_rows() +} + +const INPUT_DIR: &str = "/Users/shaikumm/Downloads/files"; + +#[test] +fn test_unsorted_merge_real_files() { + if !Path::new(INPUT_DIR).exists() { + eprintln!("Skipping: {} not found", INPUT_DIR); + return; + } + + let files = list_parquet_files(INPUT_DIR); + assert!(!files.is_empty(), "No parquet files found in {}", INPUT_DIR); + println!("Found {} input files", files.len()); + + let expected_rows = count_rows_in_files(&files); + println!("Total input rows: {}", expected_rows); + + let tmp = tempdir().unwrap(); + let output = tmp.path().join("merged_unsorted.parquet"); + let output_str = output.to_string_lossy().to_string(); + + // Empty sort columns → unsorted merge + merge_unsorted(&files, &output_str, "test-index").unwrap(); + + assert!(output.exists(), "Output file was not created"); + let actual_rows = count_rows(&output_str); + println!("Output rows: {}", actual_rows); + assert_eq!(actual_rows, expected_rows, "Row count mismatch"); +} + +/// Verify that ___row_id in the output is monotonically increasing (0, 1, 2, ...). +fn verify_row_id_order(path: &str) { + let file = File::open(path).unwrap(); + let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + let schema = builder.schema().clone(); + let col_idx = schema.index_of("___row_id").expect("___row_id not in output"); + let reader = builder.build().unwrap(); + + let mut expected: i64 = 0; + for batch in reader { + let batch = batch.unwrap(); + let col = batch.column(col_idx).as_any() + .downcast_ref::() + .expect("___row_id should be Int64"); + for i in 0..col.len() { + assert!(!col.is_null(i), "___row_id should never be null"); + assert_eq!(col.value(i), expected, "___row_id gap at row {}", expected); + expected += 1; + } + } + println!("Verified ___row_id is sequential 0..{}", expected); +} + + +#[test] +fn test_sorted_merge_real_files() { + if !Path::new(INPUT_DIR).exists() { + eprintln!("Skipping: {} not found", INPUT_DIR); + return; + } + + let files = list_parquet_files(INPUT_DIR); + assert!(!files.is_empty(), "No parquet files found in {}", INPUT_DIR); + + let expected_rows = count_rows_in_files(&files); + println!("Total input rows: {}", expected_rows); + + let tmp = tempdir().unwrap(); + let output = tmp.path().join("merged_sorted.parquet"); + let output_str = output.to_string_lossy().to_string(); + + // Sort by EventDate ascending (each input file is pre-sorted by EventDate) + let sort_cols = vec!["EventDate".to_string()]; + let reverse = vec![false]; + let nulls_first = vec![false]; + + merge_sorted(&files, &output_str, "test-index", &sort_cols, &reverse, &nulls_first) + .unwrap(); + + assert!(output.exists(), "Output file was not created"); + let actual_rows = count_rows(&output_str); + println!("Output rows: {}", actual_rows); + assert_eq!(actual_rows, expected_rows, "Row count mismatch"); + + // Verify ___row_id is sequential 0..N + verify_row_id_order(&output_str); + + // Verify EventDate is non-decreasing in the merged output + let file = File::open(&output_str).unwrap(); + let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + let out_schema = builder.schema().clone(); + let col_idx = out_schema.index_of("EventDate").unwrap(); + let reader = builder.build().unwrap(); + + let mut prev: Option = None; + let mut rows_checked: i64 = 0; + let mut out_of_order: i64 = 0; + + for batch in reader { + let batch = batch.unwrap(); + let col = batch.column(col_idx).as_any() + .downcast_ref::>() + .unwrap(); + for i in 0..col.len() { + if col.is_null(i) { continue; } + let val = col.value(i); + if let Some(p) = prev { + if val < p { + out_of_order += 1; + if out_of_order <= 5 { + eprintln!("Out of order at row {}: prev={}, cur={}", rows_checked, p, val); + } + } + } + prev = Some(val); + rows_checked += 1; + } + } + + println!("Verified EventDate sort order across {} non-null rows", rows_checked); + assert_eq!(out_of_order, 0, "Found {} out-of-order rows in EventDate", out_of_order); +} + 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 new file mode 100644 index 0000000000000..d1ec9f3527821 --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/rust/tests/sort_types_tests.rs @@ -0,0 +1,497 @@ +/* + * 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. + */ + +//! Tests for merge_sorted across all supported sort column types: +//! Int64, Int32, Float64, Float32, Utf8, and multi-column combinations. + +use std::fs::File; +use std::sync::Arc; + +use arrow::array::*; +use arrow::datatypes::{DataType, Field, Schema}; +use opensearch_parquet_format::merge::merge_sorted; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use tempfile::tempdir; + +/// Write a single RecordBatch to a new Parquet file. +fn write_parquet(path: &str, batch: &RecordBatch) { + let file = File::create(path).unwrap(); + let mut writer = ArrowWriter::try_new(file, batch.schema(), None).unwrap(); + writer.write(batch).unwrap(); + writer.close().unwrap(); +} + +/// Read all values of a typed primitive column from a Parquet file. +fn read_primitive_col( + path: &str, + col_name: &str, +) -> Vec> { + let file = File::open(path).unwrap(); + let reader = ParquetRecordBatchReaderBuilder::try_new(file) + .unwrap() + .build() + .unwrap(); + let mut vals = Vec::new(); + for batch in reader { + let batch = batch.unwrap(); + let idx = batch.schema().index_of(col_name).unwrap(); + let col = batch.column(idx).as_primitive::(); + for i in 0..col.len() { + if col.is_null(i) { + vals.push(None); + } else { + vals.push(Some(col.value(i))); + } + } + } + vals +} + +/// Read all string values from a Utf8 column. +fn read_string_col(path: &str, col_name: &str) -> Vec> { + let file = File::open(path).unwrap(); + let reader = ParquetRecordBatchReaderBuilder::try_new(file) + .unwrap() + .build() + .unwrap(); + let mut vals = Vec::new(); + for batch in reader { + let batch = batch.unwrap(); + let idx = batch.schema().index_of(col_name).unwrap(); + let col = batch.column(idx).as_string::(); + for i in 0..col.len() { + if col.is_null(i) { + vals.push(None); + } else { + vals.push(Some(col.value(i).to_string())); + } + } + } + vals +} + +/// Count rows in a Parquet file. +fn count_rows(path: &str) -> usize { + let file = File::open(path).unwrap(); + let reader = ParquetRecordBatchReaderBuilder::try_new(file) + .unwrap() + .build() + .unwrap(); + reader.map(|b| b.unwrap().num_rows()).sum() +} + +// ─── Int64 ────────────────────────────────────────────────────────────────── + +#[test] +fn test_merge_sort_by_int64() { + let schema = Arc::new(Schema::new(vec![ + Field::new("val", DataType::Int64, false), + ])); + + // File A: [1, 3, 5] File B: [2, 4, 6] File C: [0, 7, 8] + let batches = vec![ + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(vec![1, 3, 5]))]).unwrap(), + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(vec![2, 4, 6]))]).unwrap(), + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(vec![0, 7, 8]))]).unwrap(), + ]; + + let tmp = tempdir().unwrap(); + let files: Vec = batches.iter().enumerate().map(|(i, b)| { + let p = tmp.path().join(format!("input_{}.parquet", i)); + let s = p.to_string_lossy().to_string(); + write_parquet(&s, b); + s + }).collect(); + + let output = tmp.path().join("merged.parquet").to_string_lossy().to_string(); + merge_sorted(&files, &output, "test", &["val".into()], &[false], &[false]).unwrap(); + + let vals = read_primitive_col::(&output, "val"); + let vals: Vec = vals.into_iter().map(|v| v.unwrap()).collect(); + assert_eq!(vals, vec![0, 1, 2, 3, 4, 5, 6, 7, 8]); + assert_eq!(count_rows(&output), 9); +} + +// ─── Int64 with nulls ─────────────────────────────────────────────────────── + +#[test] +fn test_merge_sort_by_int64_with_nulls() { + let schema = Arc::new(Schema::new(vec![ + Field::new("val", DataType::Int64, true), + ])); + + // Each file pre-sorted: nulls last, then ascending + let batches = vec![ + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(vec![Some(1), Some(5), None]))]).unwrap(), + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(vec![Some(2), Some(4), None]))]).unwrap(), + ]; + + let tmp = tempdir().unwrap(); + let files: Vec = batches.iter().enumerate().map(|(i, b)| { + let p = tmp.path().join(format!("input_{}.parquet", i)); + let s = p.to_string_lossy().to_string(); + write_parquet(&s, b); + s + }).collect(); + + let output = tmp.path().join("merged.parquet").to_string_lossy().to_string(); + merge_sorted(&files, &output, "test", &["val".into()], &[false], &[false]).unwrap(); + + let vals = read_primitive_col::(&output, "val"); + assert_eq!(vals, vec![Some(1), Some(2), Some(4), Some(5), None, None]); +} + +// ─── Int32 ────────────────────────────────────────────────────────────────── + +#[test] +fn test_merge_sort_by_int32() { + let schema = Arc::new(Schema::new(vec![ + Field::new("val", DataType::Int32, false), + ])); + + let batches = vec![ + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![10, 30]))]).unwrap(), + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![20, 40]))]).unwrap(), + ]; + + let tmp = tempdir().unwrap(); + let files: Vec = batches.iter().enumerate().map(|(i, b)| { + let p = tmp.path().join(format!("input_{}.parquet", i)); + let s = p.to_string_lossy().to_string(); + write_parquet(&s, b); + s + }).collect(); + + let output = tmp.path().join("merged.parquet").to_string_lossy().to_string(); + merge_sorted(&files, &output, "test", &["val".into()], &[false], &[false]).unwrap(); + + let vals = read_primitive_col::(&output, "val"); + let vals: Vec = vals.into_iter().map(|v| v.unwrap()).collect(); + assert_eq!(vals, vec![10, 20, 30, 40]); +} + +// ─── Float64 ──────────────────────────────────────────────────────────────── + +#[test] +fn test_merge_sort_by_float64() { + let schema = Arc::new(Schema::new(vec![ + Field::new("val", DataType::Float64, false), + ])); + + let batches = vec![ + RecordBatch::try_new(schema.clone(), vec![Arc::new(Float64Array::from(vec![1.1, 3.3, 5.5]))]).unwrap(), + RecordBatch::try_new(schema.clone(), vec![Arc::new(Float64Array::from(vec![2.2, 4.4, 6.6]))]).unwrap(), + ]; + + let tmp = tempdir().unwrap(); + let files: Vec = batches.iter().enumerate().map(|(i, b)| { + let p = tmp.path().join(format!("input_{}.parquet", i)); + let s = p.to_string_lossy().to_string(); + write_parquet(&s, b); + s + }).collect(); + + let output = tmp.path().join("merged.parquet").to_string_lossy().to_string(); + merge_sorted(&files, &output, "test", &["val".into()], &[false], &[false]).unwrap(); + + let vals = read_primitive_col::(&output, "val"); + let vals: Vec = vals.into_iter().map(|v| v.unwrap()).collect(); + assert_eq!(vals, vec![1.1, 2.2, 3.3, 4.4, 5.5, 6.6]); +} + +// ─── Float64 with nulls ───────────────────────────────────────────────────── + +#[test] +fn test_merge_sort_by_float64_with_nulls() { + let schema = Arc::new(Schema::new(vec![ + Field::new("val", DataType::Float64, true), + ])); + + let batches = vec![ + RecordBatch::try_new(schema.clone(), vec![Arc::new(Float64Array::from(vec![None, Some(1.5), Some(4.0)]))]).unwrap(), + RecordBatch::try_new(schema.clone(), vec![Arc::new(Float64Array::from(vec![None, Some(2.5), Some(3.0)]))]).unwrap(), + ]; + + let tmp = tempdir().unwrap(); + let files: Vec = batches.iter().enumerate().map(|(i, b)| { + let p = tmp.path().join(format!("input_{}.parquet", i)); + let s = p.to_string_lossy().to_string(); + write_parquet(&s, b); + s + }).collect(); + + let output = tmp.path().join("merged.parquet").to_string_lossy().to_string(); + merge_sorted(&files, &output, "test", &["val".into()], &[false], &[true]).unwrap(); + + let vals = read_primitive_col::(&output, "val"); + assert_eq!(vals, vec![None, None, Some(1.5), Some(2.5), Some(3.0), Some(4.0)]); +} + +// ─── Float32 ──────────────────────────────────────────────────────────────── + +#[test] +fn test_merge_sort_by_float32() { + let schema = Arc::new(Schema::new(vec![ + Field::new("val", DataType::Float32, false), + ])); + + let batches = vec![ + RecordBatch::try_new(schema.clone(), vec![Arc::new(Float32Array::from(vec![1.0f32, 3.0]))]).unwrap(), + RecordBatch::try_new(schema.clone(), vec![Arc::new(Float32Array::from(vec![2.0f32, 4.0]))]).unwrap(), + ]; + + let tmp = tempdir().unwrap(); + let files: Vec = batches.iter().enumerate().map(|(i, b)| { + let p = tmp.path().join(format!("input_{}.parquet", i)); + let s = p.to_string_lossy().to_string(); + write_parquet(&s, b); + s + }).collect(); + + let output = tmp.path().join("merged.parquet").to_string_lossy().to_string(); + merge_sorted(&files, &output, "test", &["val".into()], &[false], &[false]).unwrap(); + + let vals = read_primitive_col::(&output, "val"); + let vals: Vec = vals.into_iter().map(|v| v.unwrap()).collect(); + assert_eq!(vals, vec![1.0, 2.0, 3.0, 4.0]); +} + +// ─── Float32 with nulls ───────────────────────────────────────────────────── + +#[test] +fn test_merge_sort_by_float32_with_nulls() { + let schema = Arc::new(Schema::new(vec![ + Field::new("val", DataType::Float32, true), + ])); + + let batches = vec![ + RecordBatch::try_new(schema.clone(), vec![Arc::new(Float32Array::from(vec![Some(1.0f32), Some(3.0), None]))]).unwrap(), + RecordBatch::try_new(schema.clone(), vec![Arc::new(Float32Array::from(vec![Some(2.0f32), None, None]))]).unwrap(), + ]; + + let tmp = tempdir().unwrap(); + let files: Vec = batches.iter().enumerate().map(|(i, b)| { + let p = tmp.path().join(format!("input_{}.parquet", i)); + let s = p.to_string_lossy().to_string(); + write_parquet(&s, b); + s + }).collect(); + + let output = tmp.path().join("merged.parquet").to_string_lossy().to_string(); + merge_sorted(&files, &output, "test", &["val".into()], &[false], &[false]).unwrap(); + + let vals = read_primitive_col::(&output, "val"); + assert_eq!(vals, vec![Some(1.0), Some(2.0), Some(3.0), None, None, None]); +} + +// ─── Utf8 (String / keyword) ─────────────────────────────────────────────── + +#[test] +fn test_merge_sort_by_string() { + let schema = Arc::new(Schema::new(vec![ + Field::new("val", DataType::Utf8, false), + ])); + + let batches = vec![ + RecordBatch::try_new(schema.clone(), vec![Arc::new(StringArray::from(vec!["apple", "cherry", "fig"]))]).unwrap(), + RecordBatch::try_new(schema.clone(), vec![Arc::new(StringArray::from(vec!["banana", "date", "grape"]))]).unwrap(), + ]; + + let tmp = tempdir().unwrap(); + let files: Vec = batches.iter().enumerate().map(|(i, b)| { + let p = tmp.path().join(format!("input_{}.parquet", i)); + let s = p.to_string_lossy().to_string(); + write_parquet(&s, b); + s + }).collect(); + + let output = tmp.path().join("merged.parquet").to_string_lossy().to_string(); + merge_sorted(&files, &output, "test", &["val".into()], &[false], &[false]).unwrap(); + + let vals = read_string_col(&output, "val"); + let vals: Vec = vals.into_iter().map(|v| v.unwrap()).collect(); + assert_eq!(vals, vec!["apple", "banana", "cherry", "date", "fig", "grape"]); +} + +// ─── Utf8 with nulls ──────────────────────────────────────────────────────── + +#[test] +fn test_merge_sort_by_string_with_nulls() { + let schema = Arc::new(Schema::new(vec![ + Field::new("val", DataType::Utf8, true), + ])); + + let batches = vec![ + RecordBatch::try_new(schema.clone(), vec![ + Arc::new(StringArray::from(vec![None, Some("banana"), Some("fig")])), + ]).unwrap(), + RecordBatch::try_new(schema.clone(), vec![ + Arc::new(StringArray::from(vec![None, Some("apple"), Some("cherry")])), + ]).unwrap(), + ]; + + let tmp = tempdir().unwrap(); + let files: Vec = batches.iter().enumerate().map(|(i, b)| { + let p = tmp.path().join(format!("input_{}.parquet", i)); + let s = p.to_string_lossy().to_string(); + write_parquet(&s, b); + s + }).collect(); + + let output = tmp.path().join("merged.parquet").to_string_lossy().to_string(); + merge_sorted(&files, &output, "test", &["val".into()], &[false], &[true]).unwrap(); + + let vals = read_string_col(&output, "val"); + assert_eq!(vals, vec![None, None, Some("apple".into()), Some("banana".into()), Some("cherry".into()), Some("fig".into())]); +} + +// ─── Descending sort ──────────────────────────────────────────────────────── + +#[test] +fn test_merge_sort_descending() { + let schema = Arc::new(Schema::new(vec![ + Field::new("val", DataType::Int64, false), + ])); + + // Each file sorted descending + let batches = vec![ + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(vec![8, 5, 2]))]).unwrap(), + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(vec![7, 4, 1]))]).unwrap(), + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(vec![9, 6, 3]))]).unwrap(), + ]; + + let tmp = tempdir().unwrap(); + let files: Vec = batches.iter().enumerate().map(|(i, b)| { + let p = tmp.path().join(format!("input_{}.parquet", i)); + let s = p.to_string_lossy().to_string(); + write_parquet(&s, b); + s + }).collect(); + + let output = tmp.path().join("merged.parquet").to_string_lossy().to_string(); + merge_sorted(&files, &output, "test", &["val".into()], &[true], &[false]).unwrap(); + + let vals = read_primitive_col::(&output, "val"); + let vals: Vec = vals.into_iter().map(|v| v.unwrap()).collect(); + assert_eq!(vals, vec![9, 8, 7, 6, 5, 4, 3, 2, 1]); +} + +// ─── Multi-column: String + Int64 ────────────────────────────────────────── + +#[test] +fn test_merge_sort_multi_column_string_and_int() { + let schema = Arc::new(Schema::new(vec![ + Field::new("category", DataType::Utf8, false), + Field::new("priority", DataType::Int64, false), + ])); + + // File A: (alpha,1), (alpha,3), (beta,1) + // File B: (alpha,2), (beta,2), (beta,3) + // Sorted by (category ASC, priority ASC) + let batches = vec![ + RecordBatch::try_new(schema.clone(), vec![ + Arc::new(StringArray::from(vec!["alpha", "alpha", "beta"])), + Arc::new(Int64Array::from(vec![1, 3, 1])), + ]).unwrap(), + RecordBatch::try_new(schema.clone(), vec![ + Arc::new(StringArray::from(vec!["alpha", "beta", "beta"])), + Arc::new(Int64Array::from(vec![2, 2, 3])), + ]).unwrap(), + ]; + + let tmp = tempdir().unwrap(); + let files: Vec = batches.iter().enumerate().map(|(i, b)| { + let p = tmp.path().join(format!("input_{}.parquet", i)); + let s = p.to_string_lossy().to_string(); + write_parquet(&s, b); + s + }).collect(); + + let output = tmp.path().join("merged.parquet").to_string_lossy().to_string(); + merge_sorted( + &files, &output, "test", + &["category".into(), "priority".into()], + &[false, false], + &[false, false], + ).unwrap(); + + let cats = read_string_col(&output, "category"); + let cats: Vec = cats.into_iter().map(|v| v.unwrap()).collect(); + let pris = read_primitive_col::(&output, "priority"); + let pris: Vec = pris.into_iter().map(|v| v.unwrap()).collect(); + + assert_eq!(cats, vec!["alpha", "alpha", "alpha", "beta", "beta", "beta"]); + assert_eq!(pris, vec![1, 2, 3, 1, 2, 3]); +} + +// ─── Nulls ────────────────────────────────────────────────────────────────── + +#[test] +fn test_merge_sort_with_nulls_first() { + let schema = Arc::new(Schema::new(vec![ + Field::new("val", DataType::Int64, true), + ])); + + // Each file pre-sorted with nulls first, then ascending + // File A: [null, 2, 5] File B: [null, 1, 4] + let batches = vec![ + RecordBatch::try_new(schema.clone(), vec![ + Arc::new(Int64Array::from(vec![None, Some(2), Some(5)])), + ]).unwrap(), + RecordBatch::try_new(schema.clone(), vec![ + Arc::new(Int64Array::from(vec![None, Some(1), Some(4)])), + ]).unwrap(), + ]; + + let tmp = tempdir().unwrap(); + let files: Vec = batches.iter().enumerate().map(|(i, b)| { + let p = tmp.path().join(format!("input_{}.parquet", i)); + let s = p.to_string_lossy().to_string(); + write_parquet(&s, b); + s + }).collect(); + + let output = tmp.path().join("merged.parquet").to_string_lossy().to_string(); + merge_sorted(&files, &output, "test", &["val".into()], &[false], &[true]).unwrap(); + + let vals = read_primitive_col::(&output, "val"); + // nulls_first=true → nulls come first, then ascending + assert_eq!(vals, vec![None, None, Some(1), Some(2), Some(4), Some(5)]); +} + +#[test] +fn test_merge_sort_with_nulls_last() { + let schema = Arc::new(Schema::new(vec![ + Field::new("val", DataType::Int64, true), + ])); + + let batches = vec![ + RecordBatch::try_new(schema.clone(), vec![ + Arc::new(Int64Array::from(vec![Some(1), Some(3), None])), + ]).unwrap(), + RecordBatch::try_new(schema.clone(), vec![ + Arc::new(Int64Array::from(vec![Some(2), None, None])), + ]).unwrap(), + ]; + + let tmp = tempdir().unwrap(); + let files: Vec = batches.iter().enumerate().map(|(i, b)| { + let p = tmp.path().join(format!("input_{}.parquet", i)); + let s = p.to_string_lossy().to_string(); + write_parquet(&s, b); + s + }).collect(); + + let output = tmp.path().join("merged.parquet").to_string_lossy().to_string(); + merge_sorted(&files, &output, "test", &["val".into()], &[false], &[false]).unwrap(); + + let vals = read_primitive_col::(&output, "val"); + // nulls_first=false → values ascending, then nulls + assert_eq!(vals, vec![Some(1), Some(2), Some(3), None, None, None]); +} 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 8a0bc1c6c8778..2a67511ba9c8f 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 @@ -6,8 +6,8 @@ * compatible open source license. */ -use parquet_dataformat_jni::test_utils::*; -use parquet_dataformat_jni::writer::NativeParquetWriter; +use opensearch_parquet_format::test_utils::*; +use opensearch_parquet_format::writer::NativeParquetWriter; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use std::thread; @@ -27,7 +27,6 @@ fn test_complete_writer_lifecycle() { let metadata = close_writer_and_get_metadata(&filename, schema_ptr); assert_eq!(metadata.metadata.file_metadata().num_rows(), 9); // 3 batches × 3 rows assert!(metadata.metadata.file_metadata().version() > 0); - assert_eq!(metadata.metadata.file_metadata().schema_descr().num_columns(), 3); // root + 2 fields assert!(NativeParquetWriter::sync_to_disk(filename.clone()).is_ok()); assert!(file_path.exists()); @@ -51,7 +50,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(filename.clone(), schema_ptr).is_ok() { + if NativeParquetWriter::create_writer(filename.clone(), "test-index".to_string(), schema_ptr, vec![], vec![], vec![]).is_ok() { success_count.fetch_add(1, Ordering::SeqCst); let _ = NativeParquetWriter::finalize_writer(filename); } @@ -174,7 +173,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(filename.clone(), schema_ptr).is_ok() { + if NativeParquetWriter::create_writer(filename.clone(), "test-index".to_string(), schema_ptr, vec![], vec![], vec![]).is_ok() { 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(); cleanup_ffi_data(array_ptr, data_schema_ptr); 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 693f35a846a44..20da8fe6a3f02 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 @@ -28,6 +28,7 @@ import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.Path; +import java.util.Collections; import java.util.List; import java.util.function.Consumer; @@ -144,14 +145,14 @@ public void testWriteAfterFlushThrows() throws Exception { public void testCreateWriterWithNonExistentDirectory() { expectThrows(IOException.class, () -> { try (ArrowExport export = exportSchema()) { - new NativeParquetWriter("/nonexistent/dir/file.parquet", export.getSchemaAddress()); + new NativeParquetWriter("/nonexistent/dir/file.parquet", "test-index", export.getSchemaAddress(), Collections.emptyList(), Collections.emptyList(), Collections.emptyList()); } }); } public void testCreateWriterWithInvalidSchemaAddress() { String filePath = createTempDir().resolve("bad-schema.parquet").toString(); - expectThrows(Exception.class, () -> new NativeParquetWriter(filePath, 0L)); + expectThrows(Exception.class, () -> new NativeParquetWriter(filePath, "test-index", 0L, Collections.emptyList(), Collections.emptyList(), Collections.emptyList())); } public void testWriteWithSchemaMismatch() throws Exception { @@ -235,7 +236,7 @@ public void testWriteWithNullAddresses() throws Exception { private NativeParquetWriter createWriter(String filePath) throws Exception { try (ArrowExport export = exportSchema()) { - return new NativeParquetWriter(filePath, export.getSchemaAddress()); + return new NativeParquetWriter(filePath, "test-index", export.getSchemaAddress(), Collections.emptyList(), Collections.emptyList(), Collections.emptyList()); } } diff --git a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/vsr/VSRManagerTests.java b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/vsr/VSRManagerTests.java index 6ea57eadd03ed..a646625d827cf 100644 --- a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/vsr/VSRManagerTests.java +++ b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/vsr/VSRManagerTests.java @@ -13,7 +13,10 @@ 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.Version; +import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.common.settings.Settings; +import org.opensearch.index.IndexSettings; import org.opensearch.index.mapper.NumberFieldMapper; import org.opensearch.parquet.ParquetDataFormatPlugin; import org.opensearch.parquet.bridge.ParquetFileMetadata; @@ -31,6 +34,7 @@ public class VSRManagerTests extends OpenSearchTestCase { private ArrowBufferPool bufferPool; private Schema schema; private ThreadPool threadPool; + private IndexSettings indexSettings; @Override public void setUp() throws Exception { @@ -38,6 +42,11 @@ public void setUp() throws Exception { RustBridge.initLogger(); bufferPool = new ArrowBufferPool(Settings.EMPTY); schema = new Schema(List.of(new Field("val", FieldType.nullable(new ArrowType.Int(32, true)), null))); + Settings indexSettingsBuilder = Settings.builder() + .put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT) + .build(); + IndexMetadata indexMetadata = IndexMetadata.builder("test-index").settings(indexSettingsBuilder).build(); + indexSettings = new IndexSettings(indexMetadata, Settings.EMPTY); Settings settings = Settings.builder().put("node.name", "vsrmanager-test").build(); threadPool = new ThreadPool( settings, @@ -60,7 +69,7 @@ public void tearDown() throws Exception { public void testConstructionInitializesActiveVSR() throws Exception { String filePath = createTempDir().resolve("init.parquet").toString(); - VSRManager manager = new VSRManager(filePath, schema, bufferPool, 50000, threadPool); + VSRManager manager = new VSRManager(filePath, indexSettings, schema, bufferPool, 50000, threadPool); assertNotNull(manager.getActiveManagedVSR()); assertEquals(VSRState.ACTIVE, manager.getActiveManagedVSR().getState()); // flush handles freeze + close internally @@ -69,7 +78,7 @@ public void testConstructionInitializesActiveVSR() throws Exception { public void testFlushWithNoDataReturnsMetadata() throws Exception { String filePath = createTempDir().resolve("empty.parquet").toString(); - VSRManager manager = new VSRManager(filePath, schema, bufferPool, 50000, threadPool); + VSRManager manager = new VSRManager(filePath, indexSettings, schema, bufferPool, 50000, threadPool); ParquetFileMetadata metadata = manager.flush(); assertNotNull(metadata); assertEquals(0, metadata.numRows()); @@ -77,7 +86,7 @@ public void testFlushWithNoDataReturnsMetadata() throws Exception { public void testFlushWithData() throws Exception { String filePath = createTempDir().resolve("data.parquet").toString(); - VSRManager manager = new VSRManager(filePath, schema, bufferPool, 50000, threadPool); + VSRManager manager = new VSRManager(filePath, indexSettings, schema, bufferPool, 50000, threadPool); ManagedVSR active = manager.getActiveManagedVSR(); IntVector vec = (IntVector) active.getVector("val"); @@ -93,7 +102,7 @@ public void testFlushWithData() throws Exception { public void testAddDocument() throws Exception { String filePath = createTempDir().resolve("add-doc.parquet").toString(); - VSRManager manager = new VSRManager(filePath, schema, bufferPool, 50000, threadPool); + VSRManager manager = new VSRManager(filePath, indexSettings, schema, bufferPool, 50000, threadPool); NumberFieldMapper.NumberFieldType valField = new NumberFieldMapper.NumberFieldType("val", NumberFieldMapper.NumberType.INTEGER); ParquetDocumentInput doc = new ParquetDocumentInput(); @@ -109,7 +118,7 @@ public void testAddDocument() throws Exception { public void testSyncAfterFlush() throws Exception { String filePath = createTempDir().resolve("sync.parquet").toString(); - VSRManager manager = new VSRManager(filePath, schema, bufferPool, 50000, threadPool); + VSRManager manager = new VSRManager(filePath, indexSettings, schema, bufferPool, 50000, threadPool); ManagedVSR active = manager.getActiveManagedVSR(); IntVector vec = (IntVector) active.getVector("val"); @@ -123,7 +132,7 @@ public void testSyncAfterFlush() throws Exception { public void testMaybeRotateNoOpBelowThreshold() throws Exception { String filePath = createTempDir().resolve("norotate.parquet").toString(); - VSRManager manager = new VSRManager(filePath, schema, bufferPool, 50000, threadPool); + VSRManager manager = new VSRManager(filePath, indexSettings, schema, bufferPool, 50000, threadPool); ManagedVSR original = manager.getActiveManagedVSR(); original.setRowCount(100); manager.maybeRotateActiveVSR(); @@ -133,7 +142,7 @@ public void testMaybeRotateNoOpBelowThreshold() throws Exception { public void testMaybeRotateAtThreshold() throws Exception { String filePath = createTempDir().resolve("rotate.parquet").toString(); - VSRManager manager = new VSRManager(filePath, schema, bufferPool, 50000, threadPool); + VSRManager manager = new VSRManager(filePath, indexSettings, schema, bufferPool, 50000, threadPool); ManagedVSR original = manager.getActiveManagedVSR(); original.setRowCount(50000); @@ -147,7 +156,7 @@ public void testMaybeRotateAtThreshold() throws Exception { public void testFlushAfterRotation() throws Exception { String filePath = createTempDir().resolve("rotate-flush.parquet").toString(); - VSRManager manager = new VSRManager(filePath, schema, bufferPool, 50000, threadPool); + VSRManager manager = new VSRManager(filePath, indexSettings, schema, bufferPool, 50000, threadPool); // Fill first VSR to trigger rotation ManagedVSR first = manager.getActiveManagedVSR(); @@ -171,7 +180,7 @@ public void testFlushAfterRotation() throws Exception { public void testRotationAwaitsWhenFrozenSlotOccupied() throws Exception { String filePath = createTempDir().resolve("double-rotate.parquet").toString(); - VSRManager manager = new VSRManager(filePath, schema, bufferPool, 100, threadPool); + VSRManager manager = new VSRManager(filePath, indexSettings, schema, bufferPool, 100, threadPool); // Fill first VSR to trigger rotation (async write submitted) ManagedVSR first = manager.getActiveManagedVSR(); @@ -206,7 +215,7 @@ public void testRotationAwaitsWhenFrozenSlotOccupied() throws Exception { public void testRotationWritesHappenOnBackgroundThread() throws Exception { String filePath = createTempDir().resolve("bg-thread.parquet").toString(); - VSRManager manager = new VSRManager(filePath, schema, bufferPool, 100, threadPool); + VSRManager manager = new VSRManager(filePath, indexSettings, schema, bufferPool, 100, threadPool); // Fill and rotate ManagedVSR first = manager.getActiveManagedVSR(); @@ -235,7 +244,7 @@ public void testRotationWritesHappenOnBackgroundThread() throws Exception { public void testFlushAwaitsBackgroundWrite() throws Exception { String filePath = createTempDir().resolve("flush-await.parquet").toString(); - VSRManager manager = new VSRManager(filePath, schema, bufferPool, 100, threadPool); + VSRManager manager = new VSRManager(filePath, indexSettings, schema, bufferPool, 100, threadPool); // Fill and rotate to trigger background write ManagedVSR first = manager.getActiveManagedVSR(); @@ -260,7 +269,7 @@ public void testFlushAwaitsBackgroundWrite() throws Exception { public void testCloseAwaitsBackgroundWrite() throws Exception { String filePath = createTempDir().resolve("close-await.parquet").toString(); - VSRManager manager = new VSRManager(filePath, schema, bufferPool, 100, threadPool); + VSRManager manager = new VSRManager(filePath, indexSettings, schema, bufferPool, 100, threadPool); // Fill and rotate to trigger background write ManagedVSR first = manager.getActiveManagedVSR(); diff --git a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/writer/ParquetWriterTests.java b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/writer/ParquetWriterTests.java index 7fa90cf358ed5..a4355cafa41d1 100644 --- a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/writer/ParquetWriterTests.java +++ b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/writer/ParquetWriterTests.java @@ -10,7 +10,10 @@ import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; +import org.opensearch.Version; +import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.common.settings.Settings; +import org.opensearch.index.IndexSettings; import org.opensearch.index.engine.dataformat.FileInfos; import org.opensearch.index.engine.dataformat.WriteResult; import org.opensearch.index.mapper.KeywordFieldMapper; @@ -39,6 +42,7 @@ public class ParquetWriterTests extends OpenSearchTestCase { private MappedFieldType scoreField; private Schema schema; private ThreadPool threadPool; + private IndexSettings indexSettings; @Override public void setUp() throws Exception { @@ -49,6 +53,11 @@ public void setUp() throws Exception { nameField = new KeywordFieldMapper.KeywordFieldType("name"); scoreField = new NumberFieldMapper.NumberFieldType("score", NumberFieldMapper.NumberType.LONG); schema = buildSchema(List.of(idField, nameField, scoreField)); + Settings indexSettingsBuilder = Settings.builder() + .put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT) + .build(); + IndexMetadata indexMetadata = IndexMetadata.builder("test-index").settings(indexSettingsBuilder).build(); + indexSettings = new IndexSettings(indexMetadata, Settings.EMPTY); Settings settings = Settings.builder().put("node.name", "parquetwriter-test").build(); threadPool = new ThreadPool( settings, @@ -71,63 +80,7 @@ public void tearDown() throws Exception { public void testAddDocReturnsSuccess() throws Exception { String filePath = createTempDir().resolve("success.parquet").toString(); - ParquetWriter writer = new ParquetWriter( - filePath, - 1L, - new ParquetDataFormat(), - schema, - bufferPool, - Settings.EMPTY, - threadPool, - null - ); - - ParquetDocumentInput doc = new ParquetDocumentInput(); - doc.addField(idField, 1); - doc.addField(nameField, "alice"); - doc.addField(scoreField, 100L); - WriteResult result = writer.addDoc(doc); - assertTrue(result instanceof WriteResult.Success); - doc.close(); - writer.flush(); - } - - public void testSingleDocumentFlush() throws Exception { - String filePath = createTempDir().resolve("single.parquet").toString(); - ParquetWriter writer = new ParquetWriter( - filePath, - 1L, - new ParquetDataFormat(), - schema, - bufferPool, - Settings.EMPTY, - threadPool, - null - ); - - ParquetDocumentInput doc = new ParquetDocumentInput(); - doc.addField(idField, 42); - doc.addField(nameField, "bob"); - doc.addField(scoreField, 500L); - writer.addDoc(doc); - doc.close(); - - writer.flush(); - assertEquals(1, RustBridge.getFileMetadata(filePath).numRows()); - } - - public void testMultipleDocumentsFlush() throws Exception { - String filePath = createTempDir().resolve("multi.parquet").toString(); - ParquetWriter writer = new ParquetWriter( - filePath, - 1L, - new ParquetDataFormat(), - schema, - bufferPool, - Settings.EMPTY, - threadPool, - null - ); + ParquetWriter writer = new ParquetWriter(filePath, 1L, new ParquetDataFormat(), schema, bufferPool, indexSettings, threadPool); for (int i = 0; i < 10; i++) { ParquetDocumentInput doc = new ParquetDocumentInput(); @@ -146,31 +99,7 @@ public void testMultipleDocumentsFlush() throws Exception { public void testFlushWithNoDocuments() throws Exception { String filePath = createTempDir().resolve("empty.parquet").toString(); - ParquetWriter writer = new ParquetWriter( - filePath, - 1L, - new ParquetDataFormat(), - schema, - bufferPool, - Settings.EMPTY, - threadPool, - null - ); - assertEquals(FileInfos.empty(), writer.flush()); - } - - public void testSyncAfterFlush() throws Exception { - String filePath = createTempDir().resolve("sync.parquet").toString(); - ParquetWriter writer = new ParquetWriter( - filePath, - 1L, - new ParquetDataFormat(), - schema, - bufferPool, - Settings.EMPTY, - threadPool, - null - ); + ParquetWriter writer = new ParquetWriter(filePath, 1L, new ParquetDataFormat(), schema, bufferPool, indexSettings, threadPool); ParquetDocumentInput doc = new ParquetDocumentInput(); doc.addField(idField, 1); diff --git a/server/src/main/java/org/opensearch/index/engine/dataformat/MergeInput.java b/server/src/main/java/org/opensearch/index/engine/dataformat/MergeInput.java index b9b312bc39dcc..1c1c8eb323a41 100644 --- a/server/src/main/java/org/opensearch/index/engine/dataformat/MergeInput.java +++ b/server/src/main/java/org/opensearch/index/engine/dataformat/MergeInput.java @@ -21,14 +21,14 @@ * @opensearch.experimental */ @ExperimentalApi -public record MergeInput(List writerFiles, RowIdMapping rowIdMapping, long newWriterGeneration) { +public record MergeInput(List writerFiles, RowIdMapping rowIdMapping, long newWriterGeneration, String indexName) { public MergeInput { writerFiles = List.copyOf(writerFiles); } private MergeInput(Builder builder) { - this(new ArrayList<>(builder.fileMetadataList), builder.rowIdMapping, builder.newWriterGeneration); + this(new ArrayList<>(builder.fileMetadataList), builder.rowIdMapping, builder.newWriterGeneration, builder.indexName); } /** @@ -48,6 +48,7 @@ public static class Builder { private List fileMetadataList = new ArrayList<>(); private RowIdMapping rowIdMapping; private long newWriterGeneration; + private String indexName; private Builder() {} @@ -95,6 +96,17 @@ public Builder newWriterGeneration(long newWriterGeneration) { return this; } + /** + * Sets the index name for settings lookup during merge. + * + * @param indexName the index name + * @return this builder + */ + public Builder indexName(String indexName) { + this.indexName = indexName; + return this; + } + /** * Builds an immutable {@link MergeInput}. * From 7363827b16914a3e673a0b4da33bff5663553e52 Mon Sep 17 00:00:00 2001 From: Shailesh-Kumar-Singh Date: Wed, 22 Apr 2026 14:16:37 +0530 Subject: [PATCH 02/10] add tests Signed-off-by: Shailesh-Kumar-Singh --- .../parquet/engine/ParquetIndexingEngine.java | 15 ++++++++------- .../opensearch/parquet/writer/ParquetWriter.java | 3 +-- .../parquet-data-format/src/main/rust/Cargo.toml | 1 + .../src/main/rust/src/merge/context.rs | 8 ++++++++ .../src/main/rust/src/merge/cursor.rs | 9 +++++++++ .../src/main/rust/src/merge/error.rs | 8 ++++++++ .../src/main/rust/src/merge/heap.rs | 8 ++++++++ .../src/main/rust/src/merge/io_task.rs | 8 ++++++++ .../src/main/rust/src/merge/mod.rs | 8 ++++++++ .../src/main/rust/src/merge/schema.rs | 8 ++++++++ .../src/main/rust/src/merge/sorted.rs | 8 ++++++++ .../src/main/rust/src/merge/unsorted.rs | 8 ++++++++ .../engine/ParquetIndexingEngineTests.java | 16 +++++++++++++--- .../opensearch/parquet/vsr/VSRManagerTests.java | 2 ++ .../parquet/writer/ParquetWriterTests.java | 6 ++++-- 15 files changed, 102 insertions(+), 14 deletions(-) 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 f5f789d5b45c4..9c96487c3d04f 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 @@ -150,14 +150,15 @@ public FormatChecksumStrategy getChecksumStrategy() { } private void pushSettingsToRust() { + Settings settings = indexSettings.getSettings(); NativeSettings config = NativeSettings.builder() .indexName(indexSettings.getIndex().getName()) - .compressionType(indexSettings.getValue(ParquetSettings.COMPRESSION_TYPE)) - .compressionLevel(indexSettings.getValue(ParquetSettings.COMPRESSION_LEVEL)) - .pageSizeBytes(indexSettings.getValue(ParquetSettings.PAGE_SIZE_BYTES).getBytes()) - .pageRowLimit(indexSettings.getValue(ParquetSettings.PAGE_ROW_LIMIT)) - .dictSizeBytes(indexSettings.getValue(ParquetSettings.DICT_SIZE_BYTES).getBytes()) - .rowGroupSizeBytes(indexSettings.getValue(ParquetSettings.ROW_GROUP_SIZE_BYTES).getBytes()) + .compressionType(ParquetSettings.COMPRESSION_TYPE.get(settings)) + .compressionLevel(ParquetSettings.COMPRESSION_LEVEL.get(settings)) + .pageSizeBytes(ParquetSettings.PAGE_SIZE_BYTES.get(settings).getBytes()) + .pageRowLimit(ParquetSettings.PAGE_ROW_LIMIT.get(settings)) + .dictSizeBytes(ParquetSettings.DICT_SIZE_BYTES.get(settings).getBytes()) + .rowGroupSizeBytes(ParquetSettings.ROW_GROUP_SIZE_BYTES.get(settings).getBytes()) .build(); try { RustBridge.onSettingsUpdate(config); @@ -179,7 +180,7 @@ public Writer createWriter(long writerGeneration) { dataFormat, schemaSupplier.get(), bufferPool, - settings, + indexSettings, threadPool, checksumStrategy ); 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 22165a6f13905..66ff27b22c2ef 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 @@ -64,10 +64,9 @@ public ParquetWriter( ParquetDataFormat dataFormat, Schema schema, ArrowBufferPool bufferPool, + IndexSettings indexSettings, ThreadPool threadPool, FormatChecksumStrategy checksumStrategy - IndexSettings indexSettings, - ThreadPool threadPool ) { this.file = file; this.writerGeneration = writerGeneration; 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 f93074fc71d5c..9b40ecb1c21b5 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/Cargo.toml +++ b/sandbox/plugins/parquet-data-format/src/main/rust/Cargo.toml @@ -25,3 +25,4 @@ crc32fast = { workspace = true } [dev-dependencies] opensearch-parquet-format = { path = ".", features = ["test-utils"] } + 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 f68ea432830f0..cc47b67371eeb 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 @@ -1,3 +1,11 @@ +/* + * 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. + */ + use std::fs::File; use std::path::Path; use std::sync::Arc; 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 8c48ed1d4b89b..04a8e32223514 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 @@ -1,3 +1,11 @@ +/* + * 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. + */ + use std::fs::File; use std::sync::{Arc, Mutex}; @@ -10,6 +18,7 @@ use super::error::{MergeError, MergeResult}; use super::heap::{get_sort_values, SortKey}; use super::io_task::get_merge_pool; use super::schema::projection_indices_excluding_row_id; + /// A cursor over a single sorted Parquet input file. /// /// Each cursor reads batches sequentially and prefetches the next batch on the diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/error.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/error.rs index 1c8faef6cda32..3913604276a41 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/error.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/error.rs @@ -1,3 +1,11 @@ +/* + * 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. + */ + use std::error::Error; /// Result type alias for merge operations. diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/heap.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/heap.rs index 8fefd64af4521..20d366d277292 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/heap.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/heap.rs @@ -1,3 +1,11 @@ +/* + * 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. + */ + use std::cmp::Ordering; use std::sync::Arc; 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 0cf810c0b2aac..9b87e3e51e297 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 @@ -1,3 +1,11 @@ +/* + * 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. + */ + use std::fs::File; use std::sync::OnceLock; 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 71e026e4b79c7..ee76348c99b7b 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 @@ -1,3 +1,11 @@ +/* + * 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. + */ + mod context; mod cursor; pub mod error; diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/schema.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/schema.rs index 3376d9545fdbd..6c2465e426fd5 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/schema.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/schema.rs @@ -1,3 +1,11 @@ +/* + * 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. + */ + use std::collections::HashSet; use std::sync::Arc; 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 e39d09c2da87c..d13230173ba8b 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 @@ -1,3 +1,11 @@ +/* + * 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. + */ + use std::cmp::Ordering; use std::collections::BinaryHeap; use std::sync::Arc; 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 34e2a98ead0e6..01b06b56c2d02 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 @@ -1,3 +1,11 @@ +/* + * 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. + */ + use std::fs::File; use arrow::array::RecordBatchReader; diff --git a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/engine/ParquetIndexingEngineTests.java b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/engine/ParquetIndexingEngineTests.java index 92504864cf60f..215e5bd5655bb 100644 --- a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/engine/ParquetIndexingEngineTests.java +++ b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/engine/ParquetIndexingEngineTests.java @@ -11,7 +11,10 @@ import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.opensearch.common.settings.Settings; +import org.opensearch.Version; +import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.core.index.shard.ShardId; +import org.opensearch.index.IndexSettings; import org.opensearch.index.engine.dataformat.FileInfos; import org.opensearch.index.engine.dataformat.RefreshInput; import org.opensearch.index.engine.dataformat.RefreshResult; @@ -127,8 +130,8 @@ public void testRefreshWithNullInput() throws Exception { assertTrue(result.refreshedSegments().isEmpty()); } - public void testGetMergerReturnsNull() { - assertNull(engine.getMerger()); + public void testGetMergerReturnsNonNull() { + assertNotNull(engine.getMerger()); } public void testGetNextWriterGenerationThrows() { @@ -164,7 +167,14 @@ private ParquetIndexingEngine createEngine() { Path dataPath = tempDir.resolve(indexUUID).resolve("0"); Files.createDirectories(dataPath.resolve("parquet")); ShardPath shardPath = new ShardPath(false, dataPath, dataPath, shardId); - return new ParquetIndexingEngine(Settings.EMPTY, new ParquetDataFormat(), shardPath, () -> schema, null, threadPool); + Settings indexSettingsBuilder = Settings.builder() + .put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT) + .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1) + .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0) + .build(); + IndexMetadata indexMetadata = IndexMetadata.builder("test_index").settings(indexSettingsBuilder).build(); + IndexSettings indexSettings = new IndexSettings(indexMetadata, Settings.EMPTY); + return new ParquetIndexingEngine(Settings.EMPTY, new ParquetDataFormat(), shardPath, () -> schema, indexSettings, threadPool); } catch (Exception e) { throw new RuntimeException(e); } diff --git a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/vsr/VSRManagerTests.java b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/vsr/VSRManagerTests.java index a646625d827cf..4d289e7c3ea1f 100644 --- a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/vsr/VSRManagerTests.java +++ b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/vsr/VSRManagerTests.java @@ -44,6 +44,8 @@ public void setUp() throws Exception { schema = new Schema(List.of(new Field("val", FieldType.nullable(new ArrowType.Int(32, true)), null))); Settings indexSettingsBuilder = Settings.builder() .put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT) + .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1) + .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0) .build(); IndexMetadata indexMetadata = IndexMetadata.builder("test-index").settings(indexSettingsBuilder).build(); indexSettings = new IndexSettings(indexMetadata, Settings.EMPTY); diff --git a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/writer/ParquetWriterTests.java b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/writer/ParquetWriterTests.java index a4355cafa41d1..6f9089889bc07 100644 --- a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/writer/ParquetWriterTests.java +++ b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/writer/ParquetWriterTests.java @@ -55,6 +55,8 @@ public void setUp() throws Exception { schema = buildSchema(List.of(idField, nameField, scoreField)); Settings indexSettingsBuilder = Settings.builder() .put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT) + .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1) + .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0) .build(); IndexMetadata indexMetadata = IndexMetadata.builder("test-index").settings(indexSettingsBuilder).build(); indexSettings = new IndexSettings(indexMetadata, Settings.EMPTY); @@ -80,7 +82,7 @@ public void tearDown() throws Exception { public void testAddDocReturnsSuccess() throws Exception { String filePath = createTempDir().resolve("success.parquet").toString(); - ParquetWriter writer = new ParquetWriter(filePath, 1L, new ParquetDataFormat(), schema, bufferPool, indexSettings, threadPool); + ParquetWriter writer = new ParquetWriter(filePath, 1L, new ParquetDataFormat(), schema, bufferPool, indexSettings, threadPool, null); for (int i = 0; i < 10; i++) { ParquetDocumentInput doc = new ParquetDocumentInput(); @@ -99,7 +101,7 @@ public void testAddDocReturnsSuccess() throws Exception { public void testFlushWithNoDocuments() throws Exception { String filePath = createTempDir().resolve("empty.parquet").toString(); - ParquetWriter writer = new ParquetWriter(filePath, 1L, new ParquetDataFormat(), schema, bufferPool, indexSettings, threadPool); + ParquetWriter writer = new ParquetWriter(filePath, 1L, new ParquetDataFormat(), schema, bufferPool, indexSettings, threadPool, null); ParquetDocumentInput doc = new ParquetDocumentInput(); doc.addField(idField, 1); From b55eb4a4135f21e68592b6984bbb97b465fad2bf Mon Sep 17 00:00:00 2001 From: Shailesh-Kumar-Singh Date: Wed, 22 Apr 2026 18:35:33 +0530 Subject: [PATCH 03/10] add ColumnMapping optimization Signed-off-by: Shailesh-Kumar-Singh --- gradle.properties | 1 + .../src/main/rust/src/merge/context.rs | 9 +- .../src/main/rust/src/merge/schema.rs | 87 +++++++++++++------ .../src/main/rust/src/merge/sorted.rs | 20 +++-- .../src/main/rust/src/merge/unsorted.rs | 13 ++- 5 files changed, 89 insertions(+), 41 deletions(-) diff --git a/gradle.properties b/gradle.properties index 47c3efdfbd2a0..0fa2b072de8f4 100644 --- a/gradle.properties +++ b/gradle.properties @@ -32,3 +32,4 @@ systemProp.jdk.tls.client.protocols=TLSv1.2,TLSv1.3 # jvm args for faster test execution by default systemProp.tests.jvm.argline=-XX:TieredStopAtLevel=1 -XX:ReservedCodeCacheSize=64m +systemProp.sandbox.enabled=true 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 cc47b67371eeb..8eb8782f2076e 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 @@ -136,8 +136,13 @@ impl MergeContext { return Ok(()); } - let merged = concat_batches(&self.data_schema, self.output_chunks.as_slice())?; - self.output_chunks.clear(); + let merged = if self.output_chunks.len() == 1 { + self.output_chunks.pop().unwrap() + } else { + let m = concat_batches(&self.data_schema, self.output_chunks.as_slice())?; + self.output_chunks.clear(); + m + }; let n = merged.num_rows(); let with_id = append_row_id(&merged, self.next_row_id, &self.output_schema)?; diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/schema.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/schema.rs index 6c2465e426fd5..f952bad63fb83 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/schema.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/schema.rs @@ -65,33 +65,6 @@ pub fn projection_indices_excluding_row_id(schema: &ArrowSchema) -> Vec { .collect() } -/// Pads a batch to conform to the target schema by adding null-filled columns -/// for any fields present in `target_schema` but missing from the batch. -/// -/// Returns the batch unchanged (no copy) when schemas already match. -pub fn pad_batch_to_schema( - batch: &RecordBatch, - target_schema: &Arc, -) -> MergeResult { - let batch_schema = batch.schema(); - if batch_schema.fields() == target_schema.fields() { - return Ok(batch.clone()); - } - - let num_rows = batch.num_rows(); - let mut columns: Vec = Vec::with_capacity(target_schema.fields().len()); - - for field in target_schema.fields() { - match batch_schema.index_of(field.name()) { - Ok(col_idx) => columns.push(batch.column(col_idx).clone()), - Err(_) => { - columns.push(arrow::array::new_null_array(field.data_type(), num_rows)); - } - } - } - - Ok(RecordBatch::try_new(target_schema.clone(), columns)?) -} /// Appends a `___row_id` column with sequential values `[start_id, start_id + N)` /// to the given batch, producing a new batch with the output schema. @@ -107,3 +80,63 @@ pub fn append_row_id( let result = RecordBatch::try_new(output_schema.clone(), columns)?; Ok(result) } + +// ============================================================================= +// ColumnMapping — precomputed source→target index mapping +// ============================================================================= + +/// Precomputed mapping from target schema field positions to source batch +/// column indices. Built once per cursor, reused for every batch from that cursor. +/// +/// Replaces per-batch `schema.index_of(field.name())` name lookups with O(1) +/// indexed access. +pub struct ColumnMapping { + mapping: Vec>, + target_schema: Arc, + is_identity: bool, +} + +impl ColumnMapping { + /// Build a mapping from `source_schema` → `target_schema`. + pub fn new(source_schema: &ArrowSchema, target_schema: &Arc) -> Self { + let mut mapping = Vec::with_capacity(target_schema.fields().len()); + let mut is_identity = source_schema.fields().len() == target_schema.fields().len(); + + for (target_idx, field) in target_schema.fields().iter().enumerate() { + match source_schema.index_of(field.name()) { + Ok(src_idx) => { + if is_identity && src_idx != target_idx { + is_identity = false; + } + mapping.push(Some(src_idx)); + } + Err(_) => { + is_identity = false; + mapping.push(None); + } + } + } + + Self { mapping, target_schema: target_schema.clone(), is_identity } + } + + /// Remap a batch using the precomputed mapping. Zero-copy when schemas match. + #[inline] + pub fn pad_batch(&self, batch: &RecordBatch) -> MergeResult { + if self.is_identity { + return Ok(batch.clone()); + } + let num_rows = batch.num_rows(); + let mut columns: Vec = Vec::with_capacity(self.mapping.len()); + for (i, entry) in self.mapping.iter().enumerate() { + match entry { + Some(src_idx) => columns.push(batch.column(*src_idx).clone()), + None => { + let field = &self.target_schema.fields()[i]; + columns.push(arrow::array::new_null_array(field.data_type(), num_rows)); + } + } + } + Ok(RecordBatch::try_new(self.target_schema.clone(), columns)?) + } +} 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 d13230173ba8b..35a6b603565b6 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 @@ -19,7 +19,7 @@ use super::context::MergeContext; use super::cursor::FileCursor; use super::heap::{cmp_sort_values, get_sort_values, HeapItem}; use super::io_task::{get_merge_pool, BATCH_SIZE, OUTPUT_FLUSH_ROWS}; -use super::schema::pad_batch_to_schema; +use super::schema::ColumnMapping; /// Performs a streaming k-way merge with an explicit sort direction per column. pub fn merge_sorted( @@ -81,13 +81,18 @@ pub fn merge_sorted( // ── Phase 2: Create MergeContext (union schemas, writer, IO task) ─── let mut ctx = MergeContext::new( - arrow_schemas, + arrow_schemas.clone(), &parquet_descriptors, output_path, index_name, output_flush_rows, )?; + // Precompute column mappings per cursor (avoids per-batch name lookups) + let col_mappings: Vec = arrow_schemas.iter() + .map(|s| ColumnMapping::new(s, ctx.data_schema())) + .collect(); + log_info!( "[RUST] Merge initialized ({}): {} cursors", direction_label, @@ -113,12 +118,12 @@ pub fn merge_sorted( // TIER 1: Single cursor remaining — drain it if heap.is_empty() { let cursor = &mut cursors[file_id]; + let mapping = &col_mappings[file_id]; loop { let remaining = cursor.batch_height() - cursor.row_idx; if remaining > 0 { let slice = cursor.take_slice(cursor.row_idx, remaining); - let padded = pad_batch_to_schema(&slice, ctx.data_schema())?; - ctx.push_batch(padded)?; + ctx.push_batch(mapping.pad_batch(&slice)?)?; } if !cursor.advance_past_batch()? { break; @@ -129,6 +134,7 @@ pub fn merge_sorted( // TIER 2 & 3: Multiple cursors active let cursor = &mut cursors[file_id]; + let mapping = &col_mappings[file_id]; loop { let heap_top = &heap.peek().unwrap().sort_values; @@ -138,8 +144,7 @@ pub fn merge_sorted( if cmp_sort_values(&last_val, heap_top, reverse_sorts) != Ordering::Greater { let remaining = cursor.batch_height() - cursor.row_idx; let slice = cursor.take_slice(cursor.row_idx, remaining); - let padded = pad_batch_to_schema(&slice, ctx.data_schema())?; - ctx.push_batch(padded)?; + ctx.push_batch(mapping.pad_batch(&slice)?)?; if !cursor.advance_past_batch()? { break; @@ -176,8 +181,7 @@ pub fn merge_sorted( let run_len = run_end - run_start + 1; if run_len > 0 { let slice = cursor.take_slice(run_start, run_len); - let padded = pad_batch_to_schema(&slice, ctx.data_schema())?; - ctx.push_batch(padded)?; + ctx.push_batch(mapping.pad_batch(&slice)?)?; } cursor.row_idx = run_end; 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 01b06b56c2d02..89708764597aa 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 @@ -18,7 +18,7 @@ use crate::{log_debug, log_info}; use super::context::MergeContext; use super::error::MergeResult; use super::io_task::{BATCH_SIZE, OUTPUT_FLUSH_ROWS}; -use super::schema::{pad_batch_to_schema, projection_indices_excluding_row_id}; +use super::schema::{projection_indices_excluding_row_id, ColumnMapping}; /// Unsorted merge: reads each input file sequentially, pads to union schema, /// rewrites `___row_id` with globally sequential values. No sorting performed. @@ -55,13 +55,18 @@ pub fn merge_unsorted( } let mut ctx = MergeContext::new( - arrow_schemas, + arrow_schemas.clone(), &parquet_descriptors, output_path, index_name, OUTPUT_FLUSH_ROWS, )?; + // Precompute column mappings per reader + let col_mappings: Vec = arrow_schemas.iter() + .map(|s| ColumnMapping::new(s, ctx.data_schema())) + .collect(); + // Iterate readers for data. for (file_idx, reader) in readers.into_iter().enumerate() { log_debug!( @@ -70,10 +75,10 @@ pub fn merge_unsorted( input_files.len() ); + let mapping = &col_mappings[file_idx]; for batch_result in reader { let batch = batch_result?; - let padded = pad_batch_to_schema(&batch, ctx.data_schema())?; - ctx.push_batch(padded)?; + ctx.push_batch(mapping.pad_batch(&batch)?)?; } } From b999a666e816aa0575e61a90fb22ab187273c303 Mon Sep 17 00:00:00 2001 From: Shailesh-Kumar-Singh Date: Thu, 23 Apr 2026 19:05:49 +0530 Subject: [PATCH 04/10] fix sync_to_disk test Signed-off-by: Shailesh-Kumar-Singh --- .../src/main/rust/src/writer.rs | 32 +++++++++---------- 1 file changed, 15 insertions(+), 17 deletions(-) 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 214316fe16db5..c1cebd1723243 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 @@ -45,6 +45,8 @@ lazy_static! { /// Unified per-writer registry. Keyed by temp filename. static ref WRITERS: DashMap = DashMap::new(); pub static ref SETTINGS_STORE: DashMap = DashMap::new(); + /// Holds file handles for finalized files pending fsync. Removed after sync. + static ref FILE_MANAGER: DashMap = DashMap::new(); } pub struct NativeParquetWriter; @@ -186,6 +188,10 @@ impl NativeParquetWriter { let crc32 = Self::compute_file_crc32(&filename)?; log_debug!("CRC32 for file {}: {:#010x}", filename, crc32); + // Keep a handle for sync_to_disk + let file_for_sync = File::open(&filename)?; + FILE_MANAGER.insert(filename.clone(), file_for_sync); + // Read full ParquetMetaData from the final file let file = File::open(&filename)?; let reader = SerializedFileReader::new(file)?; @@ -425,23 +431,15 @@ impl NativeParquetWriter { pub fn sync_to_disk(filename: String) -> Result<(), Box> { log_debug!("sync_to_disk called for file: {}", filename); - let file = match File::open(&filename) { - Ok(f) => f, - Err(e) => { - log_error!("ERROR: Failed to open file for fsync: {}", filename); - return Err(e.into()); - } - }; - - match file.sync_all() { - Ok(_) => { - log_debug!("Successfully fsynced file: {}", filename); - Ok(()) - } - Err(e) => { - log_error!("ERROR: Failed to fsync file: {}", filename); - Err(e.into()) - } + if let Some(file) = FILE_MANAGER.get_mut(&filename) { + file.sync_all()?; + log_debug!("Successfully fsynced file: {}", filename); + drop(file); + FILE_MANAGER.remove(&filename); + Ok(()) + } else { + log_error!("ERROR: File not found for fsync: {}", filename); + Err("File not found".into()) } } From 5b9ce09acfeabed1db3bdb6c5cfb5fc1d572294e Mon Sep 17 00:00:00 2001 From: Shailesh-Kumar-Singh Date: Mon, 27 Apr 2026 00:31:26 +0530 Subject: [PATCH 05/10] refractor change, add ParquetSortConfig class Signed-off-by: Shailesh-Kumar-Singh --- .../parquet/bridge/NativeParquetWriter.java | 13 ++--- .../parquet/bridge/ParquetSortConfig.java | 53 +++++++++++++++++++ .../opensearch/parquet/bridge/RustBridge.java | 14 +++-- .../opensearch/parquet/vsr/VSRManager.java | 15 ++---- .../bridge/NativeParquetWriterTests.java | 8 +-- 5 files changed, 70 insertions(+), 33 deletions(-) create mode 100644 sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/ParquetSortConfig.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 6db3727e499ae..d9ba83f83a3d3 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 @@ -11,7 +11,6 @@ import org.opensearch.common.SetOnce; import java.io.IOException; -import java.util.List; import java.util.concurrent.atomic.AtomicBoolean; /** @@ -19,7 +18,7 @@ * *

    Wraps the stateless JNI methods in {@link RustBridge} with a file-scoped lifecycle: *

      - *
    1. {@code new NativeParquetWriter(filePath, indexName, schemaAddress, sortColumns, reverseSorts, nullsFirst)} — creates the native writer
      2. + *
    2. {@code new NativeParquetWriter(filePath, indexName, schemaAddress, sortConfig)} — creates the native writer
      3. *
    3. {@link #write(long, long)} — sends one or more Arrow batches (repeatable)
      4. *
    4. {@link #flush()} — finalizes the Parquet file and returns metadata
      5. *
    5. {@link #sync()} — fsyncs the file to durable storage (calls flush if needed)
      6. @@ -40,21 +39,17 @@ public class NativeParquetWriter { * @param filePath the path to the Parquet file to write * @param indexName the index name for settings lookup * @param schemaAddress the native memory address of the Arrow schema - * @param sortColumns the columns to sort by, or empty list for no sorting - * @param reverseSorts whether each sort column is descending, or empty list - * @param nullsFirst whether nulls sort first for each column, or empty list + * @param sortConfig the sort configuration for the Parquet file * @throws IOException if the native writer creation fails */ public NativeParquetWriter( String filePath, String indexName, long schemaAddress, - List sortColumns, - List reverseSorts, - List nullsFirst + ParquetSortConfig sortConfig ) throws IOException { this.filePath = filePath; - RustBridge.createWriter(filePath, indexName, schemaAddress, sortColumns, reverseSorts, nullsFirst); + RustBridge.createWriter(filePath, indexName, schemaAddress, sortConfig); } /** diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/ParquetSortConfig.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/ParquetSortConfig.java new file mode 100644 index 0000000000000..b6d3290e7cf50 --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/ParquetSortConfig.java @@ -0,0 +1,53 @@ +/* + * 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.opensearch.index.IndexSettings; +import org.opensearch.index.IndexSortConfig; +import org.opensearch.search.sort.SortOrder; + +import java.util.Collections; +import java.util.List; + +/** + * Encapsulates index sort configuration for the native Parquet writer. + * + *

      Extracts sort columns, sort orders, and null-handling preferences from + * {@link IndexSettings} and exposes them as typed lists ready for the native bridge. + */ +public record ParquetSortConfig(List sortColumns, List reverseSorts, List nullsFirst) { + + private static final ParquetSortConfig EMPTY = new ParquetSortConfig( + Collections.emptyList(), + Collections.emptyList(), + Collections.emptyList() + ); + + /** + * Creates a sort config from index settings. + * + * @param indexSettings the index settings to extract sort configuration from + */ + public ParquetSortConfig(IndexSettings indexSettings) { + this( + IndexSortConfig.INDEX_SORT_FIELD_SETTING.get(indexSettings.getSettings()), + IndexSortConfig.INDEX_SORT_ORDER_SETTING.get(indexSettings.getSettings()) + .stream().map(o -> o == SortOrder.DESC).toList(), + IndexSortConfig.INDEX_SORT_MISSING_SETTING.get(indexSettings.getSettings()) + .stream().map("_first"::equals).toList() + ); + } + + /** + * Returns an empty sort config (no sorting). + */ + public static ParquetSortConfig empty() { + return EMPTY; + } +} 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 27e75fdd51632..d4dbb53f269a1 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 @@ -120,24 +120,22 @@ static void createWriter( String file, String indexName, long schemaAddress, - List sortColumns, - List reverseSorts, - List nullsFirst + ParquetSortConfig sortConfig ) throws IOException { try (var call = new NativeCall()) { var f = call.str(file); var idx = call.str(indexName); - var sorts = call.strArray(sortColumns.toArray(new String[0])); - var reverseArray = marshalBoolList(call, reverseSorts); - var nullsFirstArray = marshalBoolList(call, nullsFirst); + var sorts = call.strArray(sortConfig.sortColumns().toArray(new String[0])); + var reverseArray = marshalBoolList(call, sortConfig.reverseSorts()); + var nullsFirstArray = marshalBoolList(call, sortConfig.nullsFirst()); call.invokeIO( CREATE_WRITER, f.segment(), f.len(), idx.segment(), idx.len(), schemaAddress, sorts.ptrs(), sorts.lens(), sorts.count(), - reverseArray, (long) reverseSorts.size(), - nullsFirstArray, (long) nullsFirst.size() + reverseArray, (long) sortConfig.reverseSorts().size(), + nullsFirstArray, (long) sortConfig.nullsFirst().size() ); } } 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 5557f1a34d8dd..ced6434692d50 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 @@ -14,22 +14,20 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.opensearch.index.IndexSettings; -import org.opensearch.index.IndexSortConfig; import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.nativebridge.spi.ArrowExport; import org.opensearch.parquet.ParquetDataFormatPlugin; import org.opensearch.parquet.bridge.NativeParquetWriter; import org.opensearch.parquet.bridge.ParquetFileMetadata; +import org.opensearch.parquet.bridge.ParquetSortConfig; import org.opensearch.parquet.fields.ArrowFieldRegistry; import org.opensearch.parquet.fields.ParquetField; import org.opensearch.parquet.memory.ArrowBufferPool; import org.opensearch.parquet.writer.FieldValuePair; import org.opensearch.parquet.writer.ParquetDocumentInput; -import org.opensearch.search.sort.SortOrder; import org.opensearch.threadpool.ThreadPool; import java.io.IOException; -import java.util.List; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; @@ -213,19 +211,12 @@ public void close() { } private void initializeWriter() { - // Read sort config from index settings - List sortColumns = IndexSortConfig.INDEX_SORT_FIELD_SETTING.get(indexSettings.getSettings()); - List sortOrders = IndexSortConfig.INDEX_SORT_ORDER_SETTING.get(indexSettings.getSettings()); - List reverseSorts = sortOrders.stream().map(o -> o == SortOrder.DESC).toList(); - - List missingValues = IndexSortConfig.INDEX_SORT_MISSING_SETTING.get(indexSettings.getSettings()); - List nullsFirst = missingValues.stream().map("_first"::equals).collect(java.util.stream.Collectors.toList()); - + ParquetSortConfig sortConfig = new ParquetSortConfig(indexSettings); String indexName = indexSettings.getIndex().getName(); ArrowSchema arrowSchema = managedVSR.get().exportSchema(); try { - writer = new NativeParquetWriter(fileName, indexName, arrowSchema.memoryAddress(), sortColumns, reverseSorts, nullsFirst); + writer = new NativeParquetWriter(fileName, indexName, arrowSchema.memoryAddress(), sortConfig); } catch (Exception e) { throw new RuntimeException("Failed to initialize Parquet writer: " + e.getMessage(), e); } finally { 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 20da8fe6a3f02..4ac81c3456d12 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 @@ -22,13 +22,13 @@ 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.parquet.bridge.ParquetSortConfig; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.Path; -import java.util.Collections; import java.util.List; import java.util.function.Consumer; @@ -145,14 +145,14 @@ public void testWriteAfterFlushThrows() throws Exception { public void testCreateWriterWithNonExistentDirectory() { expectThrows(IOException.class, () -> { try (ArrowExport export = exportSchema()) { - new NativeParquetWriter("/nonexistent/dir/file.parquet", "test-index", export.getSchemaAddress(), Collections.emptyList(), Collections.emptyList(), Collections.emptyList()); + new NativeParquetWriter("/nonexistent/dir/file.parquet", "test-index", export.getSchemaAddress(), ParquetSortConfig.empty()); } }); } public void testCreateWriterWithInvalidSchemaAddress() { String filePath = createTempDir().resolve("bad-schema.parquet").toString(); - expectThrows(Exception.class, () -> new NativeParquetWriter(filePath, "test-index", 0L, Collections.emptyList(), Collections.emptyList(), Collections.emptyList())); + expectThrows(Exception.class, () -> new NativeParquetWriter(filePath, "test-index", 0L, ParquetSortConfig.empty())); } public void testWriteWithSchemaMismatch() throws Exception { @@ -236,7 +236,7 @@ public void testWriteWithNullAddresses() throws Exception { private NativeParquetWriter createWriter(String filePath) throws Exception { try (ArrowExport export = exportSchema()) { - return new NativeParquetWriter(filePath, "test-index", export.getSchemaAddress(), Collections.emptyList(), Collections.emptyList(), Collections.emptyList()); + return new NativeParquetWriter(filePath, "test-index", export.getSchemaAddress(), ParquetSortConfig.empty()); } } From 9d266267ef73766333bd6982bdaf0e791cd22a17 Mon Sep 17 00:00:00 2001 From: Shailesh-Kumar-Singh Date: Mon, 27 Apr 2026 02:30:48 +0530 Subject: [PATCH 06/10] add InvokeIO in RustBridge Signed-off-by: Shailesh-Kumar-Singh --- .../opensearch/parquet/bridge/RustBridge.java | 5 +++- .../rust/tests/merge_integration_tests.rs | 28 +++++++++++++------ 2 files changed, 23 insertions(+), 10 deletions(-) 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 d4dbb53f269a1..0eb8692a1bfb4 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 @@ -12,6 +12,7 @@ import org.opensearch.nativebridge.spi.NativeLibraryLoader; import java.io.IOException; +import java.io.UncheckedIOException; import java.lang.foreign.FunctionDescriptor; import java.lang.foreign.Linker; import java.lang.foreign.SymbolLookup; @@ -245,12 +246,14 @@ public static void mergeParquetFilesInRust(List inputFiles, String outputF var inputs = call.strArray(paths); var out = call.str(outputFile); var idx = call.str(indexName); - call.invoke( + call.invokeIO( MERGE_FILES, inputs.ptrs(), inputs.lens(), inputs.count(), out.segment(), out.len(), idx.segment(), idx.len() ); + } catch (IOException e) { + throw new UncheckedIOException("Native merge failed", e); } } 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 c056071588835..12d5a859e034a 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 @@ -51,17 +51,23 @@ fn count_rows(path: &str) -> i64 { reader.metadata().file_metadata().num_rows() } -const INPUT_DIR: &str = "/Users/shaikumm/Downloads/files"; +fn input_dir() -> Option { + std::env::var("PARQUET_TEST_INPUT_DIR").ok() +} #[test] fn test_unsorted_merge_real_files() { - if !Path::new(INPUT_DIR).exists() { - eprintln!("Skipping: {} not found", INPUT_DIR); + let Some(input_dir) = input_dir() else { + eprintln!("Skipping: PARQUET_TEST_INPUT_DIR not set"); + return; + }; + if !Path::new(&input_dir).exists() { + eprintln!("Skipping: {} not found", input_dir); return; } - let files = list_parquet_files(INPUT_DIR); - assert!(!files.is_empty(), "No parquet files found in {}", INPUT_DIR); + let files = list_parquet_files(&input_dir); + assert!(!files.is_empty(), "No parquet files found in {}", input_dir); println!("Found {} input files", files.len()); let expected_rows = count_rows_in_files(&files); @@ -106,13 +112,17 @@ fn verify_row_id_order(path: &str) { #[test] fn test_sorted_merge_real_files() { - if !Path::new(INPUT_DIR).exists() { - eprintln!("Skipping: {} not found", INPUT_DIR); + let Some(input_dir) = input_dir() else { + eprintln!("Skipping: PARQUET_TEST_INPUT_DIR not set"); + return; + }; + if !Path::new(&input_dir).exists() { + eprintln!("Skipping: {} not found", input_dir); return; } - let files = list_parquet_files(INPUT_DIR); - assert!(!files.is_empty(), "No parquet files found in {}", INPUT_DIR); + let files = list_parquet_files(&input_dir); + assert!(!files.is_empty(), "No parquet files found in {}", input_dir); let expected_rows = count_rows_in_files(&files); println!("Total input rows: {}", expected_rows); From 8c2eecfde5ada48fa79b6c82bd69d098898af4c3 Mon Sep 17 00:00:00 2001 From: Shailesh-Kumar-Singh Date: Mon, 27 Apr 2026 17:58:47 +0530 Subject: [PATCH 07/10] address comments and refractor changes Signed-off-by: Shailesh-Kumar-Singh --- .../opensearch/parquet/ParquetSettings.java | 57 +++++++-- .../parquet/bridge/NativeSettings.java | 10 ++ .../opensearch/parquet/bridge/RustBridge.java | 8 +- .../parquet/engine/ParquetIndexingEngine.java | 7 +- .../src/main/rust/src/ffm.rs | 4 + .../src/main/rust/src/native_settings.rs | 10 ++ .../src/main/rust/src/writer.rs | 112 +++++++++--------- 7 files changed, 139 insertions(+), 69 deletions(-) diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetSettings.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetSettings.java index 4f1b8dfb7d4e6..fdae923b66b8c 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetSettings.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetSettings.java @@ -27,27 +27,27 @@ private ParquetSettings() {} /** Group setting prefix for all Parquet settings. */ public static final Setting PARQUET_SETTINGS = Setting.groupSetting( - "parquet.", + "index.parquet.", Setting.Property.IndexScope ); /** Maximum row group size in bytes (default 128MB). */ public static final Setting ROW_GROUP_SIZE_BYTES = Setting.byteSizeSetting( - "parquet.row_group_size_bytes", + "index.parquet.row_group_size_bytes", new ByteSizeValue(128, ByteSizeUnit.MB), Setting.Property.IndexScope ); /** Data page size limit in bytes (default 1MB). */ public static final Setting PAGE_SIZE_BYTES = Setting.byteSizeSetting( - "parquet.page_size_bytes", + "index.parquet.page_size_bytes", new ByteSizeValue(1, ByteSizeUnit.MB), Setting.Property.IndexScope ); /** Maximum number of rows per data page (default 20000). */ public static final Setting PAGE_ROW_LIMIT = Setting.intSetting( - "parquet.page_row_limit", + "index.parquet.page_row_limit", 20000, 1, Setting.Property.IndexScope @@ -55,27 +55,51 @@ private ParquetSettings() {} /** Dictionary page size limit in bytes (default 2MB). */ public static final Setting DICT_SIZE_BYTES = Setting.byteSizeSetting( - "parquet.dict_size_bytes", + "index.parquet.dict_size_bytes", new ByteSizeValue(2, ByteSizeUnit.MB), Setting.Property.IndexScope ); /** Compression codec for Parquet files, e.g. ZSTD, SNAPPY, LZ4_RAW (default LZ4_RAW). */ public static final Setting COMPRESSION_TYPE = Setting.simpleString( - "parquet.compression_type", + "index.parquet.compression_type", "LZ4_RAW", Setting.Property.IndexScope ); /** Compression level for the chosen codec (default 2, range 1–9). */ public static final Setting COMPRESSION_LEVEL = Setting.intSetting( - "parquet.compression_level", + "index.parquet.compression_level", 2, 1, 9, Setting.Property.IndexScope ); + /** Whether bloom filters are enabled for Parquet columns (default true). */ + public static final Setting BLOOM_FILTER_ENABLED = Setting.boolSetting( + "index.parquet.bloom_filter_enabled", + true, + Setting.Property.IndexScope + ); + + /** Bloom filter false positive probability (default 0.1). */ + public static final Setting BLOOM_FILTER_FPP = Setting.doubleSetting( + "index.parquet.bloom_filter_fpp", + 0.1, + 0.0, + 1.0, + Setting.Property.IndexScope + ); + + /** Bloom filter number of distinct values hint (default 100000). */ + public static final Setting BLOOM_FILTER_NDV = Setting.longSetting( + "index.parquet.bloom_filter_ndv", + 100_000L, + 1L, + Setting.Property.IndexScope + ); + /** Maximum native memory allocation for Arrow buffers, as a percentage of non-heap memory (default 10%). */ public static final Setting MAX_NATIVE_ALLOCATION = Setting.simpleString( "parquet.max_native_allocation", @@ -91,13 +115,30 @@ private ParquetSettings() {} Setting.Property.NodeScope ); + /** File size threshold for in-memory sort vs streaming merge sort (default 32MB). */ + public static final Setting SORT_IN_MEMORY_THRESHOLD = Setting.byteSizeSetting( + "index.parquet.sort_in_memory_threshold", + new ByteSizeValue(32, ByteSizeUnit.MB), + Setting.Property.IndexScope + ); + + /** Batch size for streaming merge sort (default 8192 rows). */ + public static final Setting SORT_BATCH_SIZE = Setting.intSetting( + "index.parquet.sort_batch_size", + 8192, + 1, + Setting.Property.IndexScope + ); + /** Returns all settings defined by the Parquet plugin. */ public static List> getSettings() { return List.of( PARQUET_SETTINGS, ROW_GROUP_SIZE_BYTES, PAGE_SIZE_BYTES, PAGE_ROW_LIMIT, DICT_SIZE_BYTES, COMPRESSION_TYPE, COMPRESSION_LEVEL, - MAX_NATIVE_ALLOCATION, MAX_ROWS_PER_VSR + BLOOM_FILTER_ENABLED, BLOOM_FILTER_FPP, BLOOM_FILTER_NDV, + MAX_NATIVE_ALLOCATION, MAX_ROWS_PER_VSR, + SORT_IN_MEMORY_THRESHOLD, SORT_BATCH_SIZE ); } } diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeSettings.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeSettings.java index a8ae0a7b677da..0e636fac7317c 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeSettings.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeSettings.java @@ -25,6 +25,8 @@ public class NativeSettings { private final Boolean bloomFilterEnabled; private final Double bloomFilterFpp; private final Long bloomFilterNdv; + private final Long sortInMemoryThresholdBytes; + private final Integer sortBatchSize; private NativeSettings(Builder builder) { this.indexName = builder.indexName; @@ -37,6 +39,8 @@ private NativeSettings(Builder builder) { this.bloomFilterEnabled = builder.bloomFilterEnabled; this.bloomFilterFpp = builder.bloomFilterFpp; this.bloomFilterNdv = builder.bloomFilterNdv; + this.sortInMemoryThresholdBytes = builder.sortInMemoryThresholdBytes; + this.sortBatchSize = builder.sortBatchSize; } public String getIndexName() { return indexName; } @@ -49,6 +53,8 @@ private NativeSettings(Builder builder) { public Boolean getBloomFilterEnabled() { return bloomFilterEnabled; } public Double getBloomFilterFpp() { return bloomFilterFpp; } public Long getBloomFilterNdv() { return bloomFilterNdv; } + public Long getSortInMemoryThresholdBytes() { return sortInMemoryThresholdBytes; } + public Integer getSortBatchSize() { return sortBatchSize; } public static Builder builder() { return new Builder(); } @@ -63,6 +69,8 @@ public static class Builder { private Boolean bloomFilterEnabled; private Double bloomFilterFpp; private Long bloomFilterNdv; + private Long sortInMemoryThresholdBytes; + private Integer sortBatchSize; public Builder indexName(String v) { this.indexName = v; return this; } public Builder compressionType(String v) { this.compressionType = v; return this; } @@ -74,6 +82,8 @@ public static class Builder { public Builder bloomFilterEnabled(Boolean v) { this.bloomFilterEnabled = v; return this; } public Builder bloomFilterFpp(Double v) { this.bloomFilterFpp = v; return this; } public Builder bloomFilterNdv(Long v) { this.bloomFilterNdv = v; return this; } + public Builder sortInMemoryThresholdBytes(Long v) { this.sortInMemoryThresholdBytes = v; return this; } + public Builder sortBatchSize(Integer v) { this.sortBatchSize = v; return this; } public NativeSettings build() { return new NativeSettings(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 0eb8692a1bfb4..6cc5cd203b51e 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 @@ -97,7 +97,9 @@ public class RustBridge { ValueLayout.JAVA_LONG, // row_group_size_bytes ValueLayout.JAVA_LONG, // bloom_filter_enabled ValueLayout.JAVA_DOUBLE, // bloom_filter_fpp - ValueLayout.JAVA_LONG // bloom_filter_ndv + ValueLayout.JAVA_LONG, // bloom_filter_ndv + ValueLayout.JAVA_LONG, // sort_in_memory_threshold_bytes + ValueLayout.JAVA_LONG // sort_batch_size ) ); REMOVE_SETTINGS = linker.downcallHandle( @@ -228,7 +230,9 @@ public static void onSettingsUpdate(NativeSettings nativeSettings) throws IOExce nativeSettings.getRowGroupSizeBytes() != null ? nativeSettings.getRowGroupSizeBytes() : -1L, nativeSettings.getBloomFilterEnabled() != null ? (nativeSettings.getBloomFilterEnabled() ? 1L : 0L) : -1L, nativeSettings.getBloomFilterFpp() != null ? nativeSettings.getBloomFilterFpp() : -1.0, - nativeSettings.getBloomFilterNdv() != null ? nativeSettings.getBloomFilterNdv() : -1L + nativeSettings.getBloomFilterNdv() != null ? nativeSettings.getBloomFilterNdv() : -1L, + nativeSettings.getSortInMemoryThresholdBytes() != null ? nativeSettings.getSortInMemoryThresholdBytes() : -1L, + nativeSettings.getSortBatchSize() != null ? (long) nativeSettings.getSortBatchSize() : -1L ); } } 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 778a368a89f88..5670c56a2edb9 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 @@ -159,6 +159,11 @@ private void pushSettingsToRust() { .pageRowLimit(ParquetSettings.PAGE_ROW_LIMIT.get(settings)) .dictSizeBytes(ParquetSettings.DICT_SIZE_BYTES.get(settings).getBytes()) .rowGroupSizeBytes(ParquetSettings.ROW_GROUP_SIZE_BYTES.get(settings).getBytes()) + .bloomFilterEnabled(ParquetSettings.BLOOM_FILTER_ENABLED.get(settings)) + .bloomFilterFpp(ParquetSettings.BLOOM_FILTER_FPP.get(settings)) + .bloomFilterNdv(ParquetSettings.BLOOM_FILTER_NDV.get(settings)) + .sortInMemoryThresholdBytes(ParquetSettings.SORT_IN_MEMORY_THRESHOLD.get(settings).getBytes()) + .sortBatchSize(ParquetSettings.SORT_BATCH_SIZE.get(settings)) .build(); try { RustBridge.onSettingsUpdate(config); @@ -250,7 +255,7 @@ public void close() throws IOException { try { RustBridge.removeSettings(indexSettings.getIndex().getName()); } catch (Exception e) { - logger.warn("Failed to remove Parquet settings from Rust store for index [{}]", indexSettings.getIndex().getName(), e); + logger.warn("Failed to remove Parquet settings from Rust store for index [{}]: {}", indexSettings.getIndex().getName(), e.getMessage()); } bufferPool.close(); } 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 3e6a430c456fc..cb31f768f6dfe 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 @@ -219,6 +219,8 @@ pub unsafe extern "C" fn parquet_on_settings_update( bloom_filter_enabled: i64, bloom_filter_fpp: f64, bloom_filter_ndv: i64, + sort_in_memory_threshold_bytes: i64, + sort_batch_size: i64, ) -> i64 { let index_name = str_from_raw(index_name_ptr, index_name_len) .map_err(|e| format!("parquet_on_settings_update index_name: {}", e))?.to_string(); @@ -247,6 +249,8 @@ pub unsafe extern "C" fn parquet_on_settings_update( bloom_filter_enabled: opt_bool(bloom_filter_enabled), bloom_filter_fpp: opt_f64(bloom_filter_fpp), bloom_filter_ndv: opt_u64(bloom_filter_ndv), + sort_in_memory_threshold_bytes: opt_u64(sort_in_memory_threshold_bytes), + sort_batch_size: opt_usize(sort_batch_size), ..Default::default() }; diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/native_settings.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/native_settings.rs index 293b5cda4ebcc..bc548c6129773 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/native_settings.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/native_settings.rs @@ -27,6 +27,8 @@ pub struct NativeSettings { pub sort_columns: Vec, pub reverse_sorts: Vec, pub nulls_first: Vec, + pub sort_in_memory_threshold_bytes: Option, + pub sort_batch_size: Option, } impl NativeSettings { @@ -77,6 +79,14 @@ impl NativeSettings { pub fn has_field_configs(&self) -> bool { self.field_configs.as_ref().map_or(false, |configs| !configs.is_empty()) } + + pub fn get_sort_in_memory_threshold_bytes(&self) -> u64 { + self.sort_in_memory_threshold_bytes.unwrap_or(32 * 1024 * 1024) + } + + pub fn get_sort_batch_size(&self) -> usize { + self.sort_batch_size.unwrap_or(8192) + } } #[cfg(test)] 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 c1cebd1723243..1ed22ba6bb2d2 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 @@ -8,7 +8,7 @@ use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; use arrow::record_batch::RecordBatch; -use arrow::compute::{concat_batches, lexsort_to_indices, take, SortColumn}; +use arrow::compute::{lexsort_to_indices, take, SortColumn}; use dashmap::DashMap; use lazy_static::lazy_static; use parquet::arrow::{arrow_reader::ParquetRecordBatchReaderBuilder, ArrowWriter}; @@ -18,8 +18,8 @@ use std::io::Read; use std::path::Path; use std::sync::{Arc, Mutex}; -use crate::{log_info, log_error, log_debug}; -use crate::merge::schema::ROW_ID_COLUMN_NAME; +use crate::{log_error, log_debug}; +use crate::merge::{merge_sorted, schema::ROW_ID_COLUMN_NAME}; use crate::native_settings::NativeSettings; use crate::writer_properties_builder::WriterPropertiesBuilder; @@ -35,10 +35,7 @@ pub struct FinalizeResult { struct WriterState { writer: Arc>>, file_handle: File, - index_name: String, - sort_columns: Vec, - reverse_sorts: Vec, - nulls_first: Vec, + settings: NativeSettings, } lazy_static! { @@ -75,7 +72,7 @@ impl NativeParquetWriter { reverse_sorts: Vec, nulls_first: Vec, ) -> Result<(), Box> { - log_info!( + log_debug!( "create_writer called for file: {}, index: {}, schema_address: {}, sort_columns: {:?}, reverse_sorts: {:?}, nulls_first: {:?}", filename, index_name, schema_address, sort_columns, reverse_sorts, nulls_first ); @@ -99,33 +96,25 @@ impl NativeParquetWriter { let file = File::create(&temp_filename)?; let file_clone = file.try_clone()?; - let config: NativeSettings = SETTINGS_STORE + let mut settings: NativeSettings = SETTINGS_STORE .get(&index_name) .map(|r| r.clone()) .unwrap_or_default(); - let props = WriterPropertiesBuilder::build(&config); + settings.index_name = Some(index_name.clone()); + settings.sort_columns = sort_columns; + settings.reverse_sorts = reverse_sorts; + settings.nulls_first = nulls_first; - SETTINGS_STORE.entry(index_name.clone()).and_modify(|s| { - s.sort_columns = sort_columns.clone(); - s.reverse_sorts = reverse_sorts.clone(); - s.nulls_first = nulls_first.clone(); - }).or_insert_with(|| { - let mut s = NativeSettings::default(); - s.sort_columns = sort_columns.clone(); - s.reverse_sorts = reverse_sorts.clone(); - s.nulls_first = nulls_first.clone(); - s - }); + let props = WriterPropertiesBuilder::build(&settings); + + SETTINGS_STORE.insert(index_name, settings.clone()); let writer = ArrowWriter::try_new(file, schema, Some(props))?; WRITERS.insert(temp_filename, WriterState { writer: Arc::new(Mutex::new(writer)), file_handle: file_clone, - index_name, - sort_columns, - reverse_sorts, - nulls_first, + settings, }); Ok(()) @@ -168,21 +157,24 @@ impl NativeParquetWriter { pub fn finalize_writer(filename: String) -> Result, Box> { let temp_filename = Self::temp_filename(&filename); - log_info!("finalize_writer called for file: {} (temp: {})", filename, temp_filename); + log_debug!("finalize_writer called for file: {} (temp: {})", filename, temp_filename); if let Some((_, state)) = WRITERS.remove(&temp_filename) { - let WriterState { writer: writer_arc, file_handle: _file, index_name, sort_columns, reverse_sorts, nulls_first } = state; + let WriterState { writer: writer_arc, file_handle: _file, settings } = state; + let index_name = settings.index_name.as_deref().unwrap_or(""); match Arc::try_unwrap(writer_arc) { Ok(mutex) => { let writer = mutex.into_inner().unwrap(); match writer.close() { Ok(_) => { - log_info!("Successfully closed temp writer for: {}", temp_filename); + log_debug!("Successfully closed temp writer for: {}", temp_filename); // _file is dropped here, closing the file handle - Self::sort_and_rewrite_parquet(&temp_filename, &filename, &index_name, &sort_columns, &reverse_sorts, &nulls_first)?; + Self::sort_and_rewrite_parquet(&temp_filename, &filename, index_name, &settings.sort_columns, &settings.reverse_sorts, &settings.nulls_first)?; - let _ = std::fs::remove_file(&temp_filename); + if let Err(e) = std::fs::remove_file(&temp_filename) { + log_error!("Failed to remove temp file {}: {}", temp_filename, e); + } // Compute CRC32 by reading the final sorted file let crc32 = Self::compute_file_crc32(&filename)?; @@ -236,24 +228,28 @@ impl NativeParquetWriter { reverse_sorts: &[bool], nulls_first: &[bool], ) -> Result<(), Box> { - log_info!( + log_debug!( "sort_and_rewrite_parquet: temp={}, output={}, sort_columns={:?}, reverse_sorts={:?}, nulls_first={:?}", temp_filename, output_filename, sort_columns, reverse_sorts, nulls_first ); if sort_columns.is_empty() { - log_info!("No sort columns specified, renaming temp file to final"); + log_debug!("No sort columns specified, renaming temp file to final"); std::fs::rename(temp_filename, output_filename)?; return Ok(()); } + let config = SETTINGS_STORE + .get(index_name) + .map(|r| r.clone()) + .unwrap_or_default(); + let file_size = std::fs::metadata(temp_filename)?.len(); - const MAX_MEMORY_SIZE: u64 = 32 * 1024 * 1024; - if file_size <= MAX_MEMORY_SIZE { + if file_size <= config.get_sort_in_memory_threshold_bytes() { Self::sort_small_file(temp_filename, output_filename, index_name, sort_columns, reverse_sorts, nulls_first) } else { - Self::sort_large_file(temp_filename, output_filename, index_name, sort_columns, reverse_sorts, nulls_first) + Self::sort_large_file(temp_filename, output_filename, index_name, sort_columns, reverse_sorts, nulls_first, config.get_sort_batch_size()) } } @@ -266,26 +262,25 @@ impl NativeParquetWriter { reverse_sorts: &[bool], nulls_first: &[bool], ) -> Result<(), Box> { - log_info!("Using in-memory sort for small file: {}", temp_filename); + log_debug!("Using in-memory sort for small file: {}", temp_filename); let file = File::open(temp_filename)?; let builder = ParquetRecordBatchReaderBuilder::try_new(file)?; - let arrow_reader = builder.with_batch_size(2048).build()?; - - let mut batches = Vec::new(); - for batch_result in arrow_reader { - batches.push(batch_result?); - } - - if batches.is_empty() { - log_info!("No data to sort in file: {}", temp_filename); - std::fs::rename(temp_filename, output_filename)?; - return Ok(()); - } + let row_count = builder.metadata().file_metadata().num_rows() as usize; + let arrow_reader = builder.with_batch_size(row_count).build()?; + + let batch = match arrow_reader.into_iter().next() { + Some(Ok(b)) if b.num_rows() > 0 => b, + Some(Err(e)) => return Err(e.into()), + _ => { + log_debug!("No data to sort in file: {}", temp_filename); + std::fs::rename(temp_filename, output_filename)?; + return Ok(()); + } + }; - let schema = batches[0].schema(); - let combined_batch = concat_batches(&schema, &batches)?; - let sorted_batch = Self::sort_batch(&combined_batch, sort_columns, reverse_sorts, nulls_first)?; + let schema = batch.schema(); + let sorted_batch = Self::sort_batch(&batch, sort_columns, reverse_sorts, nulls_first)?; let final_batch = Self::rewrite_row_ids(&sorted_batch, &schema)?; Self::write_final_file(output_filename, index_name, &final_batch, schema)?; @@ -302,12 +297,13 @@ impl NativeParquetWriter { sort_columns: &[String], reverse_sorts: &[bool], nulls_first: &[bool], + batch_size: usize, ) -> Result<(), Box> { - log_info!("Using streaming merge sort for large file: {}", temp_filename); + log_debug!("Using streaming merge sort for large file: {}", temp_filename); let file = File::open(temp_filename)?; let builder = ParquetRecordBatchReaderBuilder::try_new(file)?; - let arrow_reader = builder.with_batch_size(8192).build()?; + let arrow_reader = builder.with_batch_size(batch_size).build()?; let mut chunk_paths: Vec = Vec::new(); let mut batch_count = 0; @@ -329,15 +325,15 @@ impl NativeParquetWriter { } if chunk_paths.is_empty() { - log_info!("No data to sort in file: {}", temp_filename); + log_debug!("No data to sort in file: {}", temp_filename); std::fs::rename(temp_filename, output_filename)?; return Ok(()); } - log_info!("Created {} sorted chunks, merging via streaming k-way merge", batch_count); + log_debug!("Created {} sorted chunks, merging via streaming k-way merge", batch_count); // Use the streaming merge to produce the final sorted file - crate::merge::merge_sorted( + merge_sorted( &chunk_paths, output_filename, index_name, @@ -397,7 +393,7 @@ impl NativeParquetWriter { use arrow::array::Int64Array; if let Some(row_id_idx) = schema.fields().iter().position(|f| f.name() == ROW_ID_COLUMN_NAME) { - log_info!("Rewriting ___row_id column with sequential values 0..{}", batch.num_rows()); + log_debug!("Rewriting ___row_id column with sequential values 0..{}", batch.num_rows()); let sequential_ids = Int64Array::from_iter_values( (0..batch.num_rows() as u64).map(|x| x as i64) ); @@ -424,7 +420,7 @@ impl NativeParquetWriter { let mut writer = ArrowWriter::try_new(file, schema, Some(props))?; writer.write(batch)?; writer.close()?; - log_info!("Successfully wrote final file: {}", output_filename); + log_debug!("Successfully wrote final file: {}", output_filename); Ok(()) } From ea501f928d1ded831484606c6e83e15c5e7fc096 Mon Sep 17 00:00:00 2001 From: Shailesh-Kumar-Singh Date: Mon, 27 Apr 2026 22:53:25 +0530 Subject: [PATCH 08/10] run spotlessApply Signed-off-by: Shailesh-Kumar-Singh --- .../opensearch/parquet/ParquetSettings.java | 23 +-- .../parquet/bridge/NativeParquetWriter.java | 7 +- .../parquet/bridge/NativeSettings.java | 142 ++++++++++++++---- .../parquet/bridge/ParquetSortConfig.java | 6 +- .../opensearch/parquet/bridge/RustBridge.java | 111 ++++++++------ .../parquet/engine/ParquetIndexingEngine.java | 6 +- .../parquet/merge/ParquetMergeExecutor.java | 1 - .../parquet/merge/ParquetMergeStrategy.java | 1 - .../merge/StreamingParquetMergeStrategy.java | 39 ++--- .../opensearch/parquet/vsr/VSRManager.java | 10 +- .../parquet/writer/ParquetWriter.java | 8 +- .../bridge/NativeParquetWriterTests.java | 8 +- .../engine/ParquetIndexingEngineTests.java | 2 +- .../parquet/writer/ParquetWriterTests.java | 23 ++- 14 files changed, 256 insertions(+), 131 deletions(-) diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetSettings.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetSettings.java index fdae923b66b8c..79dbb9089d9e4 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetSettings.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetSettings.java @@ -26,10 +26,7 @@ private ParquetSettings() {} public static final int DEFAULT_MAX_ROWS_PER_VSR = 50000; /** Group setting prefix for all Parquet settings. */ - public static final Setting PARQUET_SETTINGS = Setting.groupSetting( - "index.parquet.", - Setting.Property.IndexScope - ); + public static final Setting PARQUET_SETTINGS = Setting.groupSetting("index.parquet.", Setting.Property.IndexScope); /** Maximum row group size in bytes (default 128MB). */ public static final Setting ROW_GROUP_SIZE_BYTES = Setting.byteSizeSetting( @@ -134,11 +131,19 @@ private ParquetSettings() {} public static List> getSettings() { return List.of( PARQUET_SETTINGS, - ROW_GROUP_SIZE_BYTES, PAGE_SIZE_BYTES, PAGE_ROW_LIMIT, DICT_SIZE_BYTES, - COMPRESSION_TYPE, COMPRESSION_LEVEL, - BLOOM_FILTER_ENABLED, BLOOM_FILTER_FPP, BLOOM_FILTER_NDV, - MAX_NATIVE_ALLOCATION, MAX_ROWS_PER_VSR, - SORT_IN_MEMORY_THRESHOLD, SORT_BATCH_SIZE + ROW_GROUP_SIZE_BYTES, + PAGE_SIZE_BYTES, + PAGE_ROW_LIMIT, + DICT_SIZE_BYTES, + COMPRESSION_TYPE, + COMPRESSION_LEVEL, + BLOOM_FILTER_ENABLED, + BLOOM_FILTER_FPP, + BLOOM_FILTER_NDV, + MAX_NATIVE_ALLOCATION, + MAX_ROWS_PER_VSR, + SORT_IN_MEMORY_THRESHOLD, + SORT_BATCH_SIZE ); } } 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 d9ba83f83a3d3..7230c2e79c752 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 @@ -42,12 +42,7 @@ public class NativeParquetWriter { * @param sortConfig the sort configuration for the Parquet file * @throws IOException if the native writer creation fails */ - public NativeParquetWriter( - String filePath, - String indexName, - long schemaAddress, - ParquetSortConfig sortConfig - ) throws IOException { + public NativeParquetWriter(String filePath, String indexName, long schemaAddress, ParquetSortConfig sortConfig) throws IOException { this.filePath = filePath; RustBridge.createWriter(filePath, indexName, schemaAddress, sortConfig); } diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeSettings.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeSettings.java index 0e636fac7317c..c17c24668e67f 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeSettings.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeSettings.java @@ -43,20 +43,57 @@ private NativeSettings(Builder builder) { this.sortBatchSize = builder.sortBatchSize; } - public String getIndexName() { return indexName; } - public String getCompressionType() { return compressionType; } - public Integer getCompressionLevel() { return compressionLevel; } - public Long getPageSizeBytes() { return pageSizeBytes; } - public Integer getPageRowLimit() { return pageRowLimit; } - public Long getDictSizeBytes() { return dictSizeBytes; } - public Long getRowGroupSizeBytes() { return rowGroupSizeBytes; } - public Boolean getBloomFilterEnabled() { return bloomFilterEnabled; } - public Double getBloomFilterFpp() { return bloomFilterFpp; } - public Long getBloomFilterNdv() { return bloomFilterNdv; } - public Long getSortInMemoryThresholdBytes() { return sortInMemoryThresholdBytes; } - public Integer getSortBatchSize() { return sortBatchSize; } - - public static Builder builder() { return new Builder(); } + public String getIndexName() { + return indexName; + } + + public String getCompressionType() { + return compressionType; + } + + public Integer getCompressionLevel() { + return compressionLevel; + } + + public Long getPageSizeBytes() { + return pageSizeBytes; + } + + public Integer getPageRowLimit() { + return pageRowLimit; + } + + public Long getDictSizeBytes() { + return dictSizeBytes; + } + + public Long getRowGroupSizeBytes() { + return rowGroupSizeBytes; + } + + public Boolean getBloomFilterEnabled() { + return bloomFilterEnabled; + } + + public Double getBloomFilterFpp() { + return bloomFilterFpp; + } + + public Long getBloomFilterNdv() { + return bloomFilterNdv; + } + + public Long getSortInMemoryThresholdBytes() { + return sortInMemoryThresholdBytes; + } + + public Integer getSortBatchSize() { + return sortBatchSize; + } + + public static Builder builder() { + return new Builder(); + } public static class Builder { private String indexName; @@ -72,19 +109,68 @@ public static class Builder { private Long sortInMemoryThresholdBytes; private Integer sortBatchSize; - public Builder indexName(String v) { this.indexName = v; return this; } - public Builder compressionType(String v) { this.compressionType = v; return this; } - public Builder compressionLevel(Integer v) { this.compressionLevel = v; return this; } - public Builder pageSizeBytes(Long v) { this.pageSizeBytes = v; return this; } - public Builder pageRowLimit(Integer v) { this.pageRowLimit = v; return this; } - public Builder dictSizeBytes(Long v) { this.dictSizeBytes = v; return this; } - public Builder rowGroupSizeBytes(Long v) { this.rowGroupSizeBytes = v; return this; } - public Builder bloomFilterEnabled(Boolean v) { this.bloomFilterEnabled = v; return this; } - public Builder bloomFilterFpp(Double v) { this.bloomFilterFpp = v; return this; } - public Builder bloomFilterNdv(Long v) { this.bloomFilterNdv = v; return this; } - public Builder sortInMemoryThresholdBytes(Long v) { this.sortInMemoryThresholdBytes = v; return this; } - public Builder sortBatchSize(Integer v) { this.sortBatchSize = v; return this; } - - public NativeSettings build() { return new NativeSettings(this); } + public Builder indexName(String v) { + this.indexName = v; + return this; + } + + public Builder compressionType(String v) { + this.compressionType = v; + return this; + } + + public Builder compressionLevel(Integer v) { + this.compressionLevel = v; + return this; + } + + public Builder pageSizeBytes(Long v) { + this.pageSizeBytes = v; + return this; + } + + public Builder pageRowLimit(Integer v) { + this.pageRowLimit = v; + return this; + } + + public Builder dictSizeBytes(Long v) { + this.dictSizeBytes = v; + return this; + } + + public Builder rowGroupSizeBytes(Long v) { + this.rowGroupSizeBytes = v; + return this; + } + + public Builder bloomFilterEnabled(Boolean v) { + this.bloomFilterEnabled = v; + return this; + } + + public Builder bloomFilterFpp(Double v) { + this.bloomFilterFpp = v; + return this; + } + + public Builder bloomFilterNdv(Long v) { + this.bloomFilterNdv = v; + return this; + } + + public Builder sortInMemoryThresholdBytes(Long v) { + this.sortInMemoryThresholdBytes = v; + return this; + } + + public Builder sortBatchSize(Integer v) { + this.sortBatchSize = v; + return this; + } + + public NativeSettings build() { + return new NativeSettings(this); + } } } diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/ParquetSortConfig.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/ParquetSortConfig.java index b6d3290e7cf50..7d86ac3365f04 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/ParquetSortConfig.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/ParquetSortConfig.java @@ -37,10 +37,8 @@ public record ParquetSortConfig(List sortColumns, List reverseS public ParquetSortConfig(IndexSettings indexSettings) { this( IndexSortConfig.INDEX_SORT_FIELD_SETTING.get(indexSettings.getSettings()), - IndexSortConfig.INDEX_SORT_ORDER_SETTING.get(indexSettings.getSettings()) - .stream().map(o -> o == SortOrder.DESC).toList(), - IndexSortConfig.INDEX_SORT_MISSING_SETTING.get(indexSettings.getSettings()) - .stream().map("_first"::equals).toList() + IndexSortConfig.INDEX_SORT_ORDER_SETTING.get(indexSettings.getSettings()).stream().map(o -> o == SortOrder.DESC).toList(), + IndexSortConfig.INDEX_SORT_MISSING_SETTING.get(indexSettings.getSettings()).stream().map("_first"::equals).toList() ); } 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 6cc5cd203b51e..9590460862b23 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 @@ -41,29 +41,41 @@ public class RustBridge { lib.find("parquet_create_writer").orElseThrow(), FunctionDescriptor.of( ValueLayout.JAVA_LONG, - ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, // file - ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, // index_name + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG, // file + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG, // index_name ValueLayout.JAVA_LONG, // schema_address - ValueLayout.ADDRESS, ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, // sort_columns (ptrs, lens, count) - ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, // reverse_sorts (vals, count) - ValueLayout.ADDRESS, ValueLayout.JAVA_LONG // nulls_first (vals, count) + ValueLayout.ADDRESS, + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG, // sort_columns (ptrs, lens, count) + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG, // reverse_sorts (vals, count) + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG // nulls_first (vals, count) ) ); WRITE = linker.downcallHandle( lib.find("parquet_write").orElseThrow(), FunctionDescriptor.of( ValueLayout.JAVA_LONG, - ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, - ValueLayout.JAVA_LONG, ValueLayout.JAVA_LONG + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG, + ValueLayout.JAVA_LONG, + ValueLayout.JAVA_LONG ) ); FINALIZE_WRITER = linker.downcallHandle( lib.find("parquet_finalize_writer").orElseThrow(), FunctionDescriptor.of( ValueLayout.JAVA_LONG, - ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, - ValueLayout.ADDRESS, ValueLayout.ADDRESS, - ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, ValueLayout.ADDRESS, + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS, + ValueLayout.ADDRESS, + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS, ValueLayout.ADDRESS ) ); @@ -75,9 +87,13 @@ public class RustBridge { lib.find("parquet_get_file_metadata").orElseThrow(), FunctionDescriptor.of( ValueLayout.JAVA_LONG, - ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, - ValueLayout.ADDRESS, ValueLayout.ADDRESS, - ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, ValueLayout.ADDRESS + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS, + ValueLayout.ADDRESS, + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS ) ); GET_FILTERED_BYTES = linker.downcallHandle( @@ -88,8 +104,10 @@ public class RustBridge { lib.find("parquet_on_settings_update").orElseThrow(), FunctionDescriptor.of( ValueLayout.JAVA_LONG, - ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, // index_name - ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, // compression_type + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG, // index_name + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG, // compression_type ValueLayout.JAVA_LONG, // compression_level ValueLayout.JAVA_LONG, // page_size_bytes ValueLayout.JAVA_LONG, // page_row_limit @@ -110,21 +128,20 @@ public class RustBridge { lib.find("parquet_merge_files").orElseThrow(), FunctionDescriptor.of( ValueLayout.JAVA_LONG, - ValueLayout.ADDRESS, ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, // input files (ptrs, lens, count) - ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, // output file - ValueLayout.ADDRESS, ValueLayout.JAVA_LONG // index_name + ValueLayout.ADDRESS, + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG, // input files (ptrs, lens, count) + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG, // output file + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG // index_name ) ); } public static void initLogger() {} - static void createWriter( - String file, - String indexName, - long schemaAddress, - ParquetSortConfig sortConfig - ) throws IOException { + static void createWriter(String file, String indexName, long schemaAddress, ParquetSortConfig sortConfig) throws IOException { try (var call = new NativeCall()) { var f = call.str(file); var idx = call.str(indexName); @@ -133,12 +150,18 @@ static void createWriter( var nullsFirstArray = marshalBoolList(call, sortConfig.nullsFirst()); call.invokeIO( CREATE_WRITER, - f.segment(), f.len(), - idx.segment(), idx.len(), + f.segment(), + f.len(), + idx.segment(), + idx.len(), schemaAddress, - sorts.ptrs(), sorts.lens(), sorts.count(), - reverseArray, (long) sortConfig.reverseSorts().size(), - nullsFirstArray, (long) sortConfig.nullsFirst().size() + sorts.ptrs(), + sorts.lens(), + sorts.count(), + reverseArray, + (long) sortConfig.reverseSorts().size(), + nullsFirstArray, + (long) sortConfig.nullsFirst().size() ); } } @@ -159,9 +182,13 @@ static ParquetFileMetadata finalizeWriter(String file) throws IOException { var out = call.outBuffer(1024); long rc = call.invokeIO( FINALIZE_WRITER, - f.segment(), f.len(), - versionOut, numRowsOut, - out.data(), (long) out.capacity(), out.lenOut(), + f.segment(), + f.len(), + versionOut, + numRowsOut, + out.data(), + (long) out.capacity(), + out.lenOut(), crc32Out ); if (rc == 1) return null; @@ -190,12 +217,7 @@ public static ParquetFileMetadata getFileMetadata(String file) throws IOExceptio var versionOut = call.intOut(); var numRowsOut = call.longOut(); var out = call.outBuffer(1024); - call.invokeIO( - GET_FILE_METADATA, - f.segment(), f.len(), - versionOut, numRowsOut, - out.data(), (long) out.capacity(), out.lenOut() - ); + call.invokeIO(GET_FILE_METADATA, f.segment(), f.len(), versionOut, numRowsOut, out.data(), (long) out.capacity(), out.lenOut()); int createdByLen = out.actualLength(); return new ParquetFileMetadata( versionOut.get(ValueLayout.JAVA_INT, 0), @@ -221,8 +243,10 @@ public static void onSettingsUpdate(NativeSettings nativeSettings) throws IOExce var ct = nativeSettings.getCompressionType() != null ? call.str(nativeSettings.getCompressionType()) : null; call.invokeIO( ON_SETTINGS_UPDATE, - idx.segment(), idx.len(), - ct != null ? ct.segment() : java.lang.foreign.MemorySegment.NULL, ct != null ? ct.len() : -1L, + idx.segment(), + idx.len(), + ct != null ? ct.segment() : java.lang.foreign.MemorySegment.NULL, + ct != null ? ct.len() : -1L, nativeSettings.getCompressionLevel() != null ? (long) nativeSettings.getCompressionLevel() : -1L, nativeSettings.getPageSizeBytes() != null ? nativeSettings.getPageSizeBytes() : -1L, nativeSettings.getPageRowLimit() != null ? (long) nativeSettings.getPageRowLimit() : -1L, @@ -250,12 +274,7 @@ public static void mergeParquetFilesInRust(List inputFiles, String outputF var inputs = call.strArray(paths); var out = call.str(outputFile); var idx = call.str(indexName); - call.invokeIO( - MERGE_FILES, - inputs.ptrs(), inputs.lens(), inputs.count(), - out.segment(), out.len(), - idx.segment(), idx.len() - ); + call.invokeIO(MERGE_FILES, inputs.ptrs(), inputs.lens(), inputs.count(), out.segment(), out.len(), idx.segment(), idx.len()); } catch (IOException e) { throw new UncheckedIOException("Native merge failed", e); } 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 5670c56a2edb9..cdd1ec1172c2f 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 @@ -255,7 +255,11 @@ public void close() throws IOException { try { RustBridge.removeSettings(indexSettings.getIndex().getName()); } catch (Exception e) { - logger.warn("Failed to remove Parquet settings from Rust store for index [{}]: {}", indexSettings.getIndex().getName(), e.getMessage()); + logger.warn( + "Failed to remove Parquet settings from Rust store for index [{}]: {}", + indexSettings.getIndex().getName(), + e.getMessage() + ); } bufferPool.close(); } diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/ParquetMergeExecutor.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/ParquetMergeExecutor.java index 52d2c2c462d07..98a2269e7e4fa 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/ParquetMergeExecutor.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/ParquetMergeExecutor.java @@ -8,7 +8,6 @@ package org.opensearch.parquet.merge; - import org.opensearch.index.engine.dataformat.MergeInput; import org.opensearch.index.engine.dataformat.MergeResult; import org.opensearch.index.engine.dataformat.Merger; diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/ParquetMergeStrategy.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/ParquetMergeStrategy.java index e9bb508152ec7..fe3c13c61e94d 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/ParquetMergeStrategy.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/ParquetMergeStrategy.java @@ -8,7 +8,6 @@ package org.opensearch.parquet.merge; - import org.opensearch.index.engine.dataformat.MergeInput; import org.opensearch.index.engine.dataformat.MergeResult; diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/StreamingParquetMergeStrategy.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/StreamingParquetMergeStrategy.java index 539840d48e5ff..4e98d1c91ac95 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/StreamingParquetMergeStrategy.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/StreamingParquetMergeStrategy.java @@ -21,15 +21,17 @@ import java.nio.file.Files; import java.nio.file.Path; -import java.util.*; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; /** * Implements merging of Parquet files. */ public class StreamingParquetMergeStrategy implements ParquetMergeStrategy { - private static final Logger logger = - LogManager.getLogger(StreamingParquetMergeStrategy.class); + private static final Logger logger = LogManager.getLogger(StreamingParquetMergeStrategy.class); @Override public MergeResult mergeParquetFiles(MergeInput mergeInput) { @@ -41,8 +43,7 @@ public MergeResult mergeParquetFiles(MergeInput mergeInput) { } List filePaths = new ArrayList<>(); - files.forEach(writerFileSet -> writerFileSet.files().forEach( - file -> filePaths.add(Path.of(writerFileSet.directory(), file)))); + files.forEach(writerFileSet -> writerFileSet.files().forEach(file -> filePaths.add(Path.of(writerFileSet.directory(), file)))); String outputDirectory = files.getFirst().directory(); String mergedFilePath = getMergedFilePath(writerGeneration, outputDirectory); @@ -52,35 +53,23 @@ public MergeResult mergeParquetFiles(MergeInput mergeInput) { // Merge files in Rust RustBridge.mergeParquetFilesInRust(filePaths, mergedFilePath, mergeInput.indexName()); - WriterFileSet mergedWriterFileSet = - WriterFileSet.builder().directory(Path.of(outputDirectory)).addFile(mergedFileName).writerGeneration(writerGeneration).build(); + WriterFileSet mergedWriterFileSet = WriterFileSet.builder() + .directory(Path.of(outputDirectory)) + .addFile(mergedFileName) + .writerGeneration(writerGeneration) + .build(); - Map mergedWriterFileSetMap = Collections.singletonMap( - new ParquetDataFormat(), - mergedWriterFileSet - ); + Map mergedWriterFileSetMap = Collections.singletonMap(new ParquetDataFormat(), mergedWriterFileSet); return new MergeResult(mergedWriterFileSetMap); } catch (Exception exception) { - logger.error( - () -> new ParameterizedMessage( - "Merge failed while creating merged file [{}]", - mergedFilePath - ), - exception - ); + logger.error(() -> new ParameterizedMessage("Merge failed while creating merged file [{}]", mergedFilePath), exception); try { Files.deleteIfExists(Path.of(mergedFilePath)); logger.info("Stale Merged File Deleted at : [{}]", mergedFilePath); } catch (Exception innerException) { - logger.error( - () -> new ParameterizedMessage( - "Failed to delete stale merged file [{}]", - mergedFilePath - ), - innerException - ); + logger.error(() -> new ParameterizedMessage("Failed to delete stale merged file [{}]", mergedFilePath), innerException); } throw exception; 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 ced6434692d50..354f735e6c616 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 @@ -70,8 +70,14 @@ public class VSRManager implements AutoCloseable { /** * Creates a new VSRManager with asynchronous background writes (production default). */ - public VSRManager(String fileName, IndexSettings indexSettings, Schema schema, ArrowBufferPool bufferPool, - int maxRowsPerVSR, ThreadPool threadPool) { + public VSRManager( + String fileName, + IndexSettings indexSettings, + Schema schema, + ArrowBufferPool bufferPool, + int maxRowsPerVSR, + ThreadPool threadPool + ) { this(fileName, indexSettings, schema, bufferPool, maxRowsPerVSR, threadPool, true); } 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 66ff27b22c2ef..09d02c7340567 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 @@ -73,8 +73,12 @@ public ParquetWriter( this.dataFormat = dataFormat; this.checksumStrategy = checksumStrategy; this.vsrManager = new VSRManager( - file, indexSettings, schema, bufferPool, - ParquetSettings.MAX_ROWS_PER_VSR.get(indexSettings.getSettings()), threadPool + file, + indexSettings, + schema, + bufferPool, + ParquetSettings.MAX_ROWS_PER_VSR.get(indexSettings.getSettings()), + threadPool ); } 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 4ac81c3456d12..e0795008dde03 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 @@ -22,7 +22,6 @@ 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.parquet.bridge.ParquetSortConfig; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; @@ -145,7 +144,12 @@ public void testWriteAfterFlushThrows() throws Exception { public void testCreateWriterWithNonExistentDirectory() { expectThrows(IOException.class, () -> { try (ArrowExport export = exportSchema()) { - new NativeParquetWriter("/nonexistent/dir/file.parquet", "test-index", export.getSchemaAddress(), ParquetSortConfig.empty()); + new NativeParquetWriter( + "/nonexistent/dir/file.parquet", + "test-index", + export.getSchemaAddress(), + ParquetSortConfig.empty() + ); } }); } diff --git a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/engine/ParquetIndexingEngineTests.java b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/engine/ParquetIndexingEngineTests.java index 215e5bd5655bb..2061d614a86c4 100644 --- a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/engine/ParquetIndexingEngineTests.java +++ b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/engine/ParquetIndexingEngineTests.java @@ -10,9 +10,9 @@ import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; -import org.opensearch.common.settings.Settings; import org.opensearch.Version; import org.opensearch.cluster.metadata.IndexMetadata; +import org.opensearch.common.settings.Settings; import org.opensearch.core.index.shard.ShardId; import org.opensearch.index.IndexSettings; import org.opensearch.index.engine.dataformat.FileInfos; diff --git a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/writer/ParquetWriterTests.java b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/writer/ParquetWriterTests.java index 6f9089889bc07..d61ec4936c475 100644 --- a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/writer/ParquetWriterTests.java +++ b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/writer/ParquetWriterTests.java @@ -15,7 +15,6 @@ import org.opensearch.common.settings.Settings; import org.opensearch.index.IndexSettings; import org.opensearch.index.engine.dataformat.FileInfos; -import org.opensearch.index.engine.dataformat.WriteResult; import org.opensearch.index.mapper.KeywordFieldMapper; import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.index.mapper.NumberFieldMapper; @@ -82,7 +81,16 @@ public void tearDown() throws Exception { public void testAddDocReturnsSuccess() throws Exception { String filePath = createTempDir().resolve("success.parquet").toString(); - ParquetWriter writer = new ParquetWriter(filePath, 1L, new ParquetDataFormat(), schema, bufferPool, indexSettings, threadPool, null); + ParquetWriter writer = new ParquetWriter( + filePath, + 1L, + new ParquetDataFormat(), + schema, + bufferPool, + indexSettings, + threadPool, + null + ); for (int i = 0; i < 10; i++) { ParquetDocumentInput doc = new ParquetDocumentInput(); @@ -101,7 +109,16 @@ public void testAddDocReturnsSuccess() throws Exception { public void testFlushWithNoDocuments() throws Exception { String filePath = createTempDir().resolve("empty.parquet").toString(); - ParquetWriter writer = new ParquetWriter(filePath, 1L, new ParquetDataFormat(), schema, bufferPool, indexSettings, threadPool, null); + ParquetWriter writer = new ParquetWriter( + filePath, + 1L, + new ParquetDataFormat(), + schema, + bufferPool, + indexSettings, + threadPool, + null + ); ParquetDocumentInput doc = new ParquetDocumentInput(); doc.addField(idField, 1); From d47427d8b445b43552d45aaefb7e4b6d418dd2a5 Mon Sep 17 00:00:00 2001 From: Shailesh-Kumar-Singh Date: Wed, 29 Apr 2026 02:01:30 +0530 Subject: [PATCH 09/10] do spotlessApply Signed-off-by: Shailesh-Kumar-Singh --- .../opensearch/parquet/engine/ParquetIndexingEngine.java | 1 - .../parquet/merge/StreamingParquetMergeStrategy.java | 7 ++++--- .../org/opensearch/index/engine/dataformat/MergeInput.java | 6 +----- 3 files changed, 5 insertions(+), 9 deletions(-) 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 e5bac9890f714..7b37ff1c05b0a 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 @@ -14,7 +14,6 @@ import org.opensearch.common.settings.Settings; import org.opensearch.index.IndexSettings; import org.opensearch.index.engine.dataformat.IndexingExecutionEngine; -import org.opensearch.index.engine.dataformat.MergeResult; import org.opensearch.index.engine.dataformat.Merger; import org.opensearch.index.engine.dataformat.RefreshInput; import org.opensearch.index.engine.dataformat.RefreshResult; diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/StreamingParquetMergeStrategy.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/StreamingParquetMergeStrategy.java index f45a10bc7456c..d909cf1defb71 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/StreamingParquetMergeStrategy.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/StreamingParquetMergeStrategy.java @@ -52,9 +52,10 @@ public MergeResult mergeParquetFiles(MergeInput mergeInput) { } List filePaths = new ArrayList<>(); - files.forEach(writerFileSet -> writerFileSet.files().forEach( - file -> filePaths.add(shardDataPath.resolve(writerFileSet.directory()).resolve(file)) - )); + files.forEach( + writerFileSet -> writerFileSet.files() + .forEach(file -> filePaths.add(shardDataPath.resolve(writerFileSet.directory()).resolve(file))) + ); String outputDirectory = shardDataPath.resolve(files.getFirst().directory()).toString(); String mergedFilePath = getMergedFilePath(writerGeneration, outputDirectory); diff --git a/server/src/main/java/org/opensearch/index/engine/dataformat/MergeInput.java b/server/src/main/java/org/opensearch/index/engine/dataformat/MergeInput.java index aed2da2f32c6c..961b532d2ea1d 100644 --- a/server/src/main/java/org/opensearch/index/engine/dataformat/MergeInput.java +++ b/server/src/main/java/org/opensearch/index/engine/dataformat/MergeInput.java @@ -10,7 +10,6 @@ import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.index.engine.exec.Segment; - import org.opensearch.index.engine.exec.WriterFileSet; import java.util.ArrayList; @@ -41,10 +40,7 @@ private MergeInput(Builder builder) { * @return list of writer file sets for the format across all segments */ public List getFilesForFormat(String formatName) { - return segments.stream() - .map(seg -> seg.dfGroupedSearchableFiles().get(formatName)) - .filter(Objects::nonNull) - .toList(); + return segments.stream().map(seg -> seg.dfGroupedSearchableFiles().get(formatName)).filter(Objects::nonNull).toList(); } /** From d97d3ce2319cd4e7ecb0e2c6d9f0e5c6692bdd98 Mon Sep 17 00:00:00 2001 From: Shailesh-Kumar-Singh Date: Thu, 30 Apr 2026 18:06:18 +0530 Subject: [PATCH 10/10] add IntegTests, CRC in merge and address comments Signed-off-by: Shailesh-Kumar-Singh --- gradle.properties | 1 - .../composite/CompositeMergeIT.java | 277 +++++++++++++----- .../opensearch/parquet/ParquetSettings.java | 46 ++- .../parquet/bridge/NativeSettings.java | 60 +++- .../opensearch/parquet/bridge/RustBridge.java | 41 ++- .../parquet/engine/ParquetIndexingEngine.java | 11 +- .../parquet/fields/ArrowSchemaBuilder.java | 3 +- ...y.java => NativeParquetMergeStrategy.java} | 6 +- .../opensearch/parquet/vsr/VSRManager.java | 3 +- .../src/main/rust/Cargo.toml | 1 + .../src/main/rust/src/crc_writer.rs | 49 ++++ .../src/main/rust/src/ffm.rs | 91 +++++- .../src/main/rust/src/lib.rs | 1 + .../src/main/rust/src/merge/context.rs | 68 ++++- .../src/main/rust/src/merge/cursor.rs | 2 +- .../src/main/rust/src/merge/io_task.rs | 48 +-- .../src/main/rust/src/merge/schema.rs | 10 +- .../src/main/rust/src/merge/sorted.rs | 35 ++- .../src/main/rust/src/merge/unsorted.rs | 34 ++- .../src/main/rust/src/native_settings.rs | 25 +- .../src/main/rust/src/writer.rs | 72 +++-- .../rust/src/writer_properties_builder.rs | 5 +- .../rust/tests/merge_integration_tests.rs | 14 +- 23 files changed, 680 insertions(+), 223 deletions(-) rename sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/{StreamingParquetMergeStrategy.java => NativeParquetMergeStrategy.java} (92%) create mode 100644 sandbox/plugins/parquet-data-format/src/main/rust/src/crc_writer.rs diff --git a/gradle.properties b/gradle.properties index 0fa2b072de8f4..47c3efdfbd2a0 100644 --- a/gradle.properties +++ b/gradle.properties @@ -32,4 +32,3 @@ systemProp.jdk.tls.client.protocols=TLSv1.2,TLSv1.3 # jvm args for faster test execution by default systemProp.tests.jvm.argline=-XX:TieredStopAtLevel=1 -XX:ReservedCodeCacheSize=64m -systemProp.sandbox.enabled=true diff --git a/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/CompositeMergeIT.java b/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/CompositeMergeIT.java index 4faa5948d7c4e..d965df923473d 100644 --- a/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/CompositeMergeIT.java +++ b/sandbox/plugins/composite-engine/src/internalClusterTest/java/org/opensearch/composite/CompositeMergeIT.java @@ -12,29 +12,34 @@ import org.opensearch.action.admin.indices.stats.IndicesStatsResponse; import org.opensearch.action.admin.indices.stats.ShardStats; import org.opensearch.action.index.IndexResponse; +import org.opensearch.be.datafusion.DataFusionPlugin; import org.opensearch.be.lucene.LucenePlugin; import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.common.SuppressForbidden; import org.opensearch.common.settings.Settings; import org.opensearch.common.util.FeatureFlags; +import org.opensearch.common.xcontent.json.JsonXContent; import org.opensearch.core.rest.RestStatus; -import org.opensearch.index.IndexSettings; +import org.opensearch.core.xcontent.DeprecationHandler; +import org.opensearch.core.xcontent.NamedXContentRegistry; +import org.opensearch.core.xcontent.XContentParser; +import org.opensearch.index.IndexService; import org.opensearch.index.engine.CommitStats; -import org.opensearch.index.engine.dataformat.DataFormatDescriptor; -import org.opensearch.index.engine.dataformat.DataFormatRegistry; -import org.opensearch.index.engine.dataformat.ReaderManagerConfig; -import org.opensearch.index.engine.dataformat.stub.MockDataFormat; -import org.opensearch.index.engine.dataformat.stub.MockDataFormatPlugin; -import org.opensearch.index.engine.dataformat.stub.MockReaderManager; -import org.opensearch.index.engine.exec.EngineReaderManager; +import org.opensearch.index.engine.exec.Segment; +import org.opensearch.index.engine.exec.WriterFileSet; import org.opensearch.index.engine.exec.coord.DataformatAwareCatalogSnapshot; import org.opensearch.index.merge.MergeStats; -import org.opensearch.index.store.PrecomputedChecksumStrategy; +import org.opensearch.index.shard.IndexShard; +import org.opensearch.indices.IndicesService; +import org.opensearch.parquet.ParquetDataFormatPlugin; +import org.opensearch.parquet.bridge.ParquetFileMetadata; +import org.opensearch.parquet.bridge.RustBridge; import org.opensearch.plugins.Plugin; -import org.opensearch.plugins.SearchBackEndPlugin; import org.opensearch.test.OpenSearchIntegTestCase; import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; import java.util.Arrays; import java.util.Collection; import java.util.List; @@ -43,12 +48,11 @@ import java.util.function.Function; /** - * Integration tests for composite merge operations across single and multiple data format engines. + * Integration tests for composite merge with real Parquet backend. * - * Requires JDK 25 and sandbox enabled. Run with: + * Run with: * ./gradlew :sandbox:plugins:composite-engine:internalClusterTest \ - * --tests "*.CompositeMergeIT" \ - * -Dsandbox.enabled=true + * --tests "*.CompositeMergeIT" -Dsandbox.enabled=true */ @OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, numDataNodes = 1) public class CompositeMergeIT extends OpenSearchIntegTestCase { @@ -56,38 +60,6 @@ public class CompositeMergeIT extends OpenSearchIntegTestCase { private static final String INDEX_NAME = "test-composite-merge"; private static final String MERGE_ENABLED_PROPERTY = "opensearch.pluggable.dataformat.merge.enabled"; - // ── Mock DataFormatPlugin using test framework stubs ── - - public static class MockParquetDataFormatPlugin extends MockDataFormatPlugin implements SearchBackEndPlugin { - private static final MockDataFormat PARQUET_FORMAT = new MockDataFormat("parquet", 0L, Set.of()); - - public MockParquetDataFormatPlugin() { - super(PARQUET_FORMAT); - } - - @Override - public Map getFormatDescriptors(IndexSettings indexSettings, DataFormatRegistry registry) { - return Map.of("parquet", new DataFormatDescriptor("parquet", new PrecomputedChecksumStrategy())); - } - - @Override - public String name() { - return "mock-parquet-backend"; - } - - @Override - public List getSupportedFormats() { - return List.of("parquet"); - } - - @Override - public EngineReaderManager createReaderManager(ReaderManagerConfig settings) { - return new MockReaderManager("parquet"); - } - } - - // ── Test setup ── - @Override public void setUp() throws Exception { enableMerge(); @@ -117,7 +89,7 @@ private static void disableMerge() { @Override protected Collection> nodePlugins() { - return Arrays.asList(MockParquetDataFormatPlugin.class, CompositeDataFormatPlugin.class, LucenePlugin.class); + return Arrays.asList(ParquetDataFormatPlugin.class, CompositeDataFormatPlugin.class, LucenePlugin.class, DataFusionPlugin.class); } @Override @@ -128,30 +100,65 @@ protected Settings nodeSettings(int nodeOrdinal) { .build(); } - // ── Tests ── + /** + * Verifies background merge produces a valid merged parquet file + * with correct row count and source files cleaned up. + */ + public void testBackgroundMerge() throws Exception { + client().admin() + .indices() + .prepareCreate(INDEX_NAME) + .setSettings(unsortedSettings()) + .setMapping("name", "type=keyword", "age", "type=integer") + .get(); + ensureGreen(INDEX_NAME); + + int docsPerCycle = 5; + int refreshCycles = 15; + indexDocsAcrossMultipleRefreshes(refreshCycles, docsPerCycle); + int totalDocs = refreshCycles * docsPerCycle; + + assertBusy(() -> { + flush(INDEX_NAME); + DataformatAwareCatalogSnapshot snapshot = getCatalogSnapshot(); + assertTrue( + "Expected merges to reduce segment count below " + refreshCycles + ", but got: " + snapshot.getSegments().size(), + snapshot.getSegments().size() < refreshCycles + ); + }); + + MergeStats mergeStats = getMergeStats(); + assertTrue("Expected at least one merge to have occurred", mergeStats.getTotal() > 0); + + DataformatAwareCatalogSnapshot snapshot = getCatalogSnapshot(); + assertEquals(Set.of("parquet"), snapshot.getDataFormats()); + + verifyRowCount(snapshot, totalDocs); + } /** - * Verifies that background merges are triggered automatically after refresh - * when enough segments accumulate to exceed the TieredMergePolicy threshold. - *

      - * Flow: index docs across many refresh cycles → each refresh calls - * triggerPossibleMerges() → MergeScheduler picks up merge candidates - * asynchronously → segment count decreases. + * Verifies sorted merge with age DESC (nulls first), name ASC (nulls last). */ - public void testBackgroundMergeSingleEngine() throws Exception { - createIndex(INDEX_NAME, singleEngineSettings()); + public void testSortedMerge() throws Exception { + client().admin() + .indices() + .prepareCreate(INDEX_NAME) + .setSettings(sortedSettings()) + .setMapping("name", "type=keyword", "age", "type=integer") + .get(); ensureGreen(INDEX_NAME); - // Create enough segments to exceed TieredMergePolicy's default threshold (~10) - int totalSegmentsCreated = indexDocsAcrossMultipleRefreshes(15, 5); + int docsPerCycle = 10; + int refreshCycles = 15; + indexDocsWithNullsAcrossRefreshes(refreshCycles, docsPerCycle); + int totalDocs = refreshCycles * docsPerCycle; - // Wait for async background merges to complete assertBusy(() -> { flush(INDEX_NAME); DataformatAwareCatalogSnapshot snapshot = getCatalogSnapshot(); assertTrue( - "Expected merges to reduce segment count below " + totalSegmentsCreated + ", but got: " + snapshot.getSegments().size(), - snapshot.getSegments().size() < totalSegmentsCreated + "Expected merges to reduce segment count below " + refreshCycles + ", but got: " + snapshot.getSegments().size(), + snapshot.getSegments().size() < refreshCycles ); }); @@ -160,14 +167,18 @@ public void testBackgroundMergeSingleEngine() throws Exception { DataformatAwareCatalogSnapshot snapshot = getCatalogSnapshot(); assertEquals(Set.of("parquet"), snapshot.getDataFormats()); + + verifyRowCount(snapshot, totalDocs); + verifySortOrder(snapshot); } - // ── Helpers ── + // ── Settings ── - private Settings singleEngineSettings() { + private Settings unsortedSettings() { return 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") @@ -175,19 +186,153 @@ private Settings singleEngineSettings() { .build(); } - private int indexDocsAcrossMultipleRefreshes(int refreshCycles, int docsPerCycle) { + private Settings sortedSettings() { + return 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") + .putList("index.sort.field", "age", "name") + .putList("index.sort.order", "desc", "asc") + .putList("index.sort.missing", "_first", "_last") + .build(); + } + + // ── Indexing ── + + private void indexDocsAcrossMultipleRefreshes(int refreshCycles, int docsPerCycle) { for (int cycle = 0; cycle < refreshCycles; cycle++) { for (int i = 0; i < docsPerCycle; i++) { IndexResponse response = client().prepareIndex() .setIndex(INDEX_NAME) - .setSource("field_text", randomAlphaOfLength(10), "field_number", randomIntBetween(1, 1000)) + .setSource("name", randomAlphaOfLength(10), "age", randomIntBetween(1, 1000)) .get(); assertEquals(RestStatus.CREATED, response.status()); } RefreshResponse refreshResponse = client().admin().indices().prepareRefresh(INDEX_NAME).get(); assertEquals(RestStatus.OK, refreshResponse.getStatus()); } - return refreshCycles; + } + + private void indexDocsWithNullsAcrossRefreshes(int refreshCycles, int docsPerCycle) { + for (int cycle = 0; cycle < refreshCycles; cycle++) { + for (int i = 0; i < docsPerCycle; i++) { + IndexResponse response; + if (i % 5 == 0) { + response = client().prepareIndex().setIndex(INDEX_NAME).setSource("name", randomAlphaOfLength(10)).get(); + } else { + response = client().prepareIndex() + .setIndex(INDEX_NAME) + .setSource("name", randomAlphaOfLength(10), "age", randomIntBetween(0, 100)) + .get(); + } + assertEquals(RestStatus.CREATED, response.status()); + } + RefreshResponse refreshResponse = client().admin().indices().prepareRefresh(INDEX_NAME).get(); + assertEquals(RestStatus.OK, refreshResponse.getStatus()); + } + } + + // ── Verification ── + + private void verifyRowCount(DataformatAwareCatalogSnapshot snapshot, int expectedTotalDocs) throws IOException { + Path parquetDir = getParquetDir(); + long totalRows = 0; + for (Segment segment : snapshot.getSegments()) { + WriterFileSet wfs = segment.dfGroupedSearchableFiles().get("parquet"); + assertNotNull("Segment should have parquet files", wfs); + for (String file : wfs.files()) { + Path filePath = parquetDir.resolve(file); + assertTrue("Parquet file should exist: " + filePath, Files.exists(filePath)); + ParquetFileMetadata metadata = RustBridge.getFileMetadata(filePath.toString()); + totalRows += metadata.numRows(); + } + } + assertEquals("Total rows across all segments should match ingested docs", expectedTotalDocs, totalRows); + } + + /** + * Verifies that merged parquet files have age in DESC order with nulls first, + * and within same age, name in ASC order with nulls last. + */ + @SuppressForbidden(reason = "JSON parsing for test verification of parquet output") + private void verifySortOrder(DataformatAwareCatalogSnapshot snapshot) throws Exception { + Path parquetDir = getParquetDir(); + for (Segment segment : snapshot.getSegments()) { + WriterFileSet wfs = segment.dfGroupedSearchableFiles().get("parquet"); + for (String file : wfs.files()) { + Path filePath = parquetDir.resolve(file); + String json = RustBridge.readAsJson(filePath.toString()); + List> rows; + try ( + XContentParser parser = JsonXContent.jsonXContent.createParser( + NamedXContentRegistry.EMPTY, + DeprecationHandler.THROW_UNSUPPORTED_OPERATION, + json + ) + ) { + rows = parser.list().stream().map(o -> { + @SuppressWarnings("unchecked") + Map m = (Map) o; + return m; + }).toList(); + } + if (rows.size() <= 1) continue; + + for (int i = 1; i < rows.size(); i++) { + Object prevAge = rows.get(i - 1).get("age"); + Object currAge = rows.get(i).get("age"); + + // nulls first for age + if (prevAge == null && currAge == null) continue; + if (prevAge == null) continue; // null before non-null is correct + if (currAge == null) { + fail("age null should come before non-null, but found non-null at " + (i - 1) + " and null at " + i); + } + + int prevAgeVal = ((Number) prevAge).intValue(); + int currAgeVal = ((Number) currAge).intValue(); + + assertTrue( + "age should be DESC but found " + prevAgeVal + " before " + currAgeVal + " at row " + i, + prevAgeVal >= currAgeVal + ); + + // When age is equal, verify name ASC (nulls last) + if (prevAgeVal == currAgeVal) { + Object prevName = rows.get(i - 1).get("name"); + Object currName = rows.get(i).get("name"); + + if (prevName != null && currName == null) continue; // non-null before null is correct for nulls last + if (prevName == null && currName != null) { + fail("name nulls should be last, but found null at " + (i - 1) + " and non-null at " + i); + } + if (prevName != null && currName != null) { + assertTrue( + "name should be ASC but found '" + prevName + "' before '" + currName + "' at row " + i, + ((String) prevName).compareTo((String) currName) <= 0 + ); + } + } + } + } + } + } + + private Path getParquetDir() { + IndexShard shard = getPrimaryShard(); + return shard.shardPath().getDataPath().resolve("parquet"); + } + + private IndexShard getPrimaryShard() { + String nodeName = getClusterState().routingTable().index(INDEX_NAME).shard(0).primaryShard().currentNodeId(); + String nodeNameResolved = getClusterState().nodes().get(nodeName).getName(); + IndicesService indicesService = internalCluster().getInstance(IndicesService.class, nodeNameResolved); + IndexService indexService = indicesService.indexServiceSafe(resolveIndex(INDEX_NAME)); + return indexService.getShard(0); } private DataformatAwareCatalogSnapshot getCatalogSnapshot() throws IOException { diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetSettings.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetSettings.java index 79dbb9089d9e4..ab58d0bfdf11c 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetSettings.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/ParquetSettings.java @@ -28,13 +28,6 @@ private ParquetSettings() {} /** Group setting prefix for all Parquet settings. */ public static final Setting PARQUET_SETTINGS = Setting.groupSetting("index.parquet.", Setting.Property.IndexScope); - /** Maximum row group size in bytes (default 128MB). */ - public static final Setting ROW_GROUP_SIZE_BYTES = Setting.byteSizeSetting( - "index.parquet.row_group_size_bytes", - new ByteSizeValue(128, ByteSizeUnit.MB), - Setting.Property.IndexScope - ); - /** Data page size limit in bytes (default 1MB). */ public static final Setting PAGE_SIZE_BYTES = Setting.byteSizeSetting( "index.parquet.page_size_bytes", @@ -127,11 +120,42 @@ private ParquetSettings() {} Setting.Property.IndexScope ); + /** Maximum number of rows per row group (default 1000000). */ + public static final Setting ROW_GROUP_MAX_ROWS = Setting.intSetting( + "index.parquet.row_group_max_rows", + 1_000_000, + 1, + Setting.Property.IndexScope + ); + + /** Batch size for reading records during merge (default 100000 rows). */ + public static final Setting MERGE_BATCH_SIZE = Setting.intSetting( + "index.parquet.merge_batch_size", + 100_000, + 1, + Setting.Property.IndexScope + ); + + /** Number of Rayon threads for parallel column encoding during merge (default num_cores/8, min 1). */ + public static final Setting MERGE_RAYON_THREADS = Setting.intSetting( + "parquet.merge_rayon_threads", + Math.max(1, Runtime.getRuntime().availableProcessors() / 8), + 1, + Setting.Property.NodeScope + ); + + /** Number of Tokio IO threads for async disk writes during merge (default num_cores/8, min 1). */ + public static final Setting MERGE_IO_THREADS = Setting.intSetting( + "parquet.merge_io_threads", + Math.max(1, Runtime.getRuntime().availableProcessors() / 8), + 1, + Setting.Property.NodeScope + ); + /** Returns all settings defined by the Parquet plugin. */ public static List> getSettings() { return List.of( PARQUET_SETTINGS, - ROW_GROUP_SIZE_BYTES, PAGE_SIZE_BYTES, PAGE_ROW_LIMIT, DICT_SIZE_BYTES, @@ -143,7 +167,11 @@ public static List> getSettings() { MAX_NATIVE_ALLOCATION, MAX_ROWS_PER_VSR, SORT_IN_MEMORY_THRESHOLD, - SORT_BATCH_SIZE + SORT_BATCH_SIZE, + ROW_GROUP_MAX_ROWS, + MERGE_BATCH_SIZE, + MERGE_RAYON_THREADS, + MERGE_IO_THREADS ); } } diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeSettings.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeSettings.java index c17c24668e67f..db940828424d3 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeSettings.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/NativeSettings.java @@ -21,12 +21,15 @@ public class NativeSettings { private final Long pageSizeBytes; private final Integer pageRowLimit; private final Long dictSizeBytes; - private final Long rowGroupSizeBytes; private final Boolean bloomFilterEnabled; private final Double bloomFilterFpp; private final Long bloomFilterNdv; private final Long sortInMemoryThresholdBytes; private final Integer sortBatchSize; + private final Integer rowGroupMaxRows; + private final Integer mergeBatchSize; + private final Integer mergeRayonThreads; + private final Integer mergeIoThreads; private NativeSettings(Builder builder) { this.indexName = builder.indexName; @@ -35,12 +38,15 @@ private NativeSettings(Builder builder) { this.pageSizeBytes = builder.pageSizeBytes; this.pageRowLimit = builder.pageRowLimit; this.dictSizeBytes = builder.dictSizeBytes; - this.rowGroupSizeBytes = builder.rowGroupSizeBytes; this.bloomFilterEnabled = builder.bloomFilterEnabled; this.bloomFilterFpp = builder.bloomFilterFpp; this.bloomFilterNdv = builder.bloomFilterNdv; this.sortInMemoryThresholdBytes = builder.sortInMemoryThresholdBytes; this.sortBatchSize = builder.sortBatchSize; + this.rowGroupMaxRows = builder.rowGroupMaxRows; + this.mergeBatchSize = builder.mergeBatchSize; + this.mergeRayonThreads = builder.mergeRayonThreads; + this.mergeIoThreads = builder.mergeIoThreads; } public String getIndexName() { @@ -67,10 +73,6 @@ public Long getDictSizeBytes() { return dictSizeBytes; } - public Long getRowGroupSizeBytes() { - return rowGroupSizeBytes; - } - public Boolean getBloomFilterEnabled() { return bloomFilterEnabled; } @@ -91,6 +93,22 @@ public Integer getSortBatchSize() { return sortBatchSize; } + public Integer getRowGroupMaxRows() { + return rowGroupMaxRows; + } + + public Integer getMergeBatchSize() { + return mergeBatchSize; + } + + public Integer getMergeRayonThreads() { + return mergeRayonThreads; + } + + public Integer getMergeIoThreads() { + return mergeIoThreads; + } + public static Builder builder() { return new Builder(); } @@ -102,12 +120,15 @@ public static class Builder { private Long pageSizeBytes; private Integer pageRowLimit; private Long dictSizeBytes; - private Long rowGroupSizeBytes; private Boolean bloomFilterEnabled; private Double bloomFilterFpp; private Long bloomFilterNdv; private Long sortInMemoryThresholdBytes; private Integer sortBatchSize; + private Integer rowGroupMaxRows; + private Integer mergeBatchSize; + private Integer mergeRayonThreads; + private Integer mergeIoThreads; public Builder indexName(String v) { this.indexName = v; @@ -139,11 +160,6 @@ public Builder dictSizeBytes(Long v) { return this; } - public Builder rowGroupSizeBytes(Long v) { - this.rowGroupSizeBytes = v; - return this; - } - public Builder bloomFilterEnabled(Boolean v) { this.bloomFilterEnabled = v; return this; @@ -169,6 +185,26 @@ public Builder sortBatchSize(Integer v) { return this; } + public Builder rowGroupMaxRows(Integer v) { + this.rowGroupMaxRows = v; + return this; + } + + public Builder mergeBatchSize(Integer v) { + this.mergeBatchSize = v; + return this; + } + + public Builder mergeRayonThreads(Integer v) { + this.mergeRayonThreads = v; + return this; + } + + public Builder mergeIoThreads(Integer v) { + this.mergeIoThreads = v; + return this; + } + public NativeSettings build() { return new NativeSettings(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 9590460862b23..ca4ffa0c9c7d9 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 @@ -33,6 +33,7 @@ public class RustBridge { private static final MethodHandle ON_SETTINGS_UPDATE; private static final MethodHandle REMOVE_SETTINGS; private static final MethodHandle MERGE_FILES; + private static final MethodHandle READ_AS_JSON; static { SymbolLookup lib = NativeLibraryLoader.symbolLookup(); @@ -112,12 +113,15 @@ public class RustBridge { ValueLayout.JAVA_LONG, // page_size_bytes ValueLayout.JAVA_LONG, // page_row_limit ValueLayout.JAVA_LONG, // dict_size_bytes - ValueLayout.JAVA_LONG, // row_group_size_bytes ValueLayout.JAVA_LONG, // bloom_filter_enabled ValueLayout.JAVA_DOUBLE, // bloom_filter_fpp ValueLayout.JAVA_LONG, // bloom_filter_ndv ValueLayout.JAVA_LONG, // sort_in_memory_threshold_bytes - ValueLayout.JAVA_LONG // sort_batch_size + ValueLayout.JAVA_LONG, // sort_batch_size + ValueLayout.JAVA_LONG, // row_group_max_rows + ValueLayout.JAVA_LONG, // merge_batch_size + ValueLayout.JAVA_LONG, // merge_rayon_threads + ValueLayout.JAVA_LONG // merge_io_threads ) ); REMOVE_SETTINGS = linker.downcallHandle( @@ -137,6 +141,17 @@ public class RustBridge { ValueLayout.JAVA_LONG // index_name ) ); + READ_AS_JSON = linker.downcallHandle( + lib.find("parquet_read_as_json").orElseThrow(), + FunctionDescriptor.of( + ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG, // file + ValueLayout.ADDRESS, // out_buf + ValueLayout.JAVA_LONG, // buf_capacity + ValueLayout.ADDRESS // out_len + ) + ); } public static void initLogger() {} @@ -251,12 +266,15 @@ public static void onSettingsUpdate(NativeSettings nativeSettings) throws IOExce nativeSettings.getPageSizeBytes() != null ? nativeSettings.getPageSizeBytes() : -1L, nativeSettings.getPageRowLimit() != null ? (long) nativeSettings.getPageRowLimit() : -1L, nativeSettings.getDictSizeBytes() != null ? nativeSettings.getDictSizeBytes() : -1L, - nativeSettings.getRowGroupSizeBytes() != null ? nativeSettings.getRowGroupSizeBytes() : -1L, nativeSettings.getBloomFilterEnabled() != null ? (nativeSettings.getBloomFilterEnabled() ? 1L : 0L) : -1L, nativeSettings.getBloomFilterFpp() != null ? nativeSettings.getBloomFilterFpp() : -1.0, nativeSettings.getBloomFilterNdv() != null ? nativeSettings.getBloomFilterNdv() : -1L, nativeSettings.getSortInMemoryThresholdBytes() != null ? nativeSettings.getSortInMemoryThresholdBytes() : -1L, - nativeSettings.getSortBatchSize() != null ? (long) nativeSettings.getSortBatchSize() : -1L + nativeSettings.getSortBatchSize() != null ? (long) nativeSettings.getSortBatchSize() : -1L, + nativeSettings.getRowGroupMaxRows() != null ? (long) nativeSettings.getRowGroupMaxRows() : -1L, + nativeSettings.getMergeBatchSize() != null ? (long) nativeSettings.getMergeBatchSize() : -1L, + nativeSettings.getMergeRayonThreads() != null ? (long) nativeSettings.getMergeRayonThreads() : -1L, + nativeSettings.getMergeIoThreads() != null ? (long) nativeSettings.getMergeIoThreads() : -1L ); } } @@ -291,5 +309,20 @@ private static java.lang.foreign.MemorySegment marshalBoolList(NativeCall call, return seg; } + /** + * Reads a parquet file and returns its contents as a JSON string. + */ + public static String readAsJson(String file) throws IOException { + try (var call = new NativeCall()) { + var f = call.str(file); + int bufSize = 10 * 1024 * 1024; // 10MB + var outBuf = call.buf(bufSize); + var outLen = call.longOut(); + call.invokeIO(READ_AS_JSON, f.segment(), f.len(), outBuf, (long) bufSize, outLen); + int len = (int) outLen.get(ValueLayout.JAVA_LONG, 0); + return new String(outBuf.asSlice(0, len).toArray(ValueLayout.JAVA_BYTE), StandardCharsets.UTF_8); + } + } + private RustBridge() {} } 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 7b37ff1c05b0a..0b4cb025e8463 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 @@ -27,8 +27,8 @@ import org.opensearch.parquet.bridge.NativeSettings; import org.opensearch.parquet.bridge.RustBridge; import org.opensearch.parquet.memory.ArrowBufferPool; +import org.opensearch.parquet.merge.NativeParquetMergeStrategy; import org.opensearch.parquet.merge.ParquetMergeExecutor; -import org.opensearch.parquet.merge.StreamingParquetMergeStrategy; import org.opensearch.parquet.writer.ParquetDocumentInput; import org.opensearch.parquet.writer.ParquetWriter; import org.opensearch.threadpool.ThreadPool; @@ -73,6 +73,7 @@ public class ParquetIndexingEngine implements IndexingExecutionEngine schemaSupplier; private final ArrowBufferPool bufferPool; private final IndexSettings indexSettings; + private final Settings nodeSettings; private final ThreadPool threadPool; private final FormatChecksumStrategy checksumStrategy; private final Merger parquetMerger; @@ -127,6 +128,7 @@ public ParquetIndexingEngine( this.schemaSupplier = schemaSupplier; this.bufferPool = new ArrowBufferPool(settings); this.indexSettings = indexSettings; + this.nodeSettings = settings; this.threadPool = threadPool; this.checksumStrategy = checksumStrategy; try { @@ -137,7 +139,7 @@ public ParquetIndexingEngine( throw new RuntimeException(e); } this.parquetMerger = new ParquetMergeExecutor( - new StreamingParquetMergeStrategy(dataFormat, indexSettings.getIndex().getName(), shardPath.getDataPath()) + new NativeParquetMergeStrategy(dataFormat, indexSettings.getIndex().getName(), shardPath.getDataPath()) ); pushSettingsToRust(); } @@ -160,12 +162,15 @@ private void pushSettingsToRust() { .pageSizeBytes(ParquetSettings.PAGE_SIZE_BYTES.get(settings).getBytes()) .pageRowLimit(ParquetSettings.PAGE_ROW_LIMIT.get(settings)) .dictSizeBytes(ParquetSettings.DICT_SIZE_BYTES.get(settings).getBytes()) - .rowGroupSizeBytes(ParquetSettings.ROW_GROUP_SIZE_BYTES.get(settings).getBytes()) .bloomFilterEnabled(ParquetSettings.BLOOM_FILTER_ENABLED.get(settings)) .bloomFilterFpp(ParquetSettings.BLOOM_FILTER_FPP.get(settings)) .bloomFilterNdv(ParquetSettings.BLOOM_FILTER_NDV.get(settings)) .sortInMemoryThresholdBytes(ParquetSettings.SORT_IN_MEMORY_THRESHOLD.get(settings).getBytes()) .sortBatchSize(ParquetSettings.SORT_BATCH_SIZE.get(settings)) + .rowGroupMaxRows(ParquetSettings.ROW_GROUP_MAX_ROWS.get(settings)) + .mergeBatchSize(ParquetSettings.MERGE_BATCH_SIZE.get(settings)) + .mergeRayonThreads(ParquetSettings.MERGE_RAYON_THREADS.get(nodeSettings)) + .mergeIoThreads(ParquetSettings.MERGE_IO_THREADS.get(nodeSettings)) .build(); try { RustBridge.onSettingsUpdate(config); diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/fields/ArrowSchemaBuilder.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/fields/ArrowSchemaBuilder.java index 49c1d86b5742d..84b2b21712fa3 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/fields/ArrowSchemaBuilder.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/fields/ArrowSchemaBuilder.java @@ -12,6 +12,7 @@ import org.apache.arrow.vector.types.pojo.Schema; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.opensearch.index.engine.dataformat.DocumentInput; import org.opensearch.index.mapper.FieldNamesFieldMapper; import org.opensearch.index.mapper.IndexFieldMapper; import org.opensearch.index.mapper.Mapper; @@ -57,7 +58,7 @@ public static Schema getSchema(MapperService mapperService) { } // Add row ID field (long) LongParquetField longField = new LongParquetField(); - fields.add(new Field("_row_id", longField.getFieldType(), null)); + fields.add(new Field(DocumentInput.ROW_ID_FIELD, longField.getFieldType(), null)); return new Schema(fields); } diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/StreamingParquetMergeStrategy.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/NativeParquetMergeStrategy.java similarity index 92% rename from sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/StreamingParquetMergeStrategy.java rename to sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/NativeParquetMergeStrategy.java index d909cf1defb71..a3bbdee35c6f8 100644 --- a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/StreamingParquetMergeStrategy.java +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/merge/NativeParquetMergeStrategy.java @@ -28,15 +28,15 @@ /** * Implements merging of Parquet files. */ -public class StreamingParquetMergeStrategy implements ParquetMergeStrategy { +public class NativeParquetMergeStrategy implements ParquetMergeStrategy { - private static final Logger logger = LogManager.getLogger(StreamingParquetMergeStrategy.class); + private static final Logger logger = LogManager.getLogger(NativeParquetMergeStrategy.class); private final DataFormat dataFormat; private final String indexName; private final Path shardDataPath; - public StreamingParquetMergeStrategy(DataFormat dataFormat, String indexName, Path shardDataPath) { + public NativeParquetMergeStrategy(DataFormat dataFormat, String indexName, Path shardDataPath) { this.dataFormat = dataFormat; this.indexName = indexName; this.shardDataPath = shardDataPath; 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 354f735e6c616..668d142d6aae1 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 @@ -14,6 +14,7 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.opensearch.index.IndexSettings; +import org.opensearch.index.engine.dataformat.DocumentInput; import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.nativebridge.spi.ArrowExport; import org.opensearch.parquet.ParquetDataFormatPlugin; @@ -130,7 +131,7 @@ public void addDocument(ParquetDocumentInput doc) throws IOException { parquetField.createField(fieldType, activeVSR, pair.getValue()); } int rowIndex = activeVSR.getRowCount(); - BigIntVector rowIdVector = (BigIntVector) activeVSR.getVector("___row_id"); + BigIntVector rowIdVector = (BigIntVector) activeVSR.getVector(DocumentInput.ROW_ID_FIELD); if (rowIdVector != null) { rowIdVector.setSafe(rowIndex, doc.getRowId()); } 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 9b40ecb1c21b5..379ec6aea4149 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/Cargo.toml +++ b/sandbox/plugins/parquet-data-format/src/main/rust/Cargo.toml @@ -22,6 +22,7 @@ native-bridge-common = { workspace = true } rayon = { workspace = true } tokio = { workspace = true } crc32fast = { workspace = true } +serde_json = { workspace = true } [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 new file mode 100644 index 0000000000000..7ae7c436e9477 --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/crc_writer.rs @@ -0,0 +1,49 @@ +/* + * 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. + */ + +use std::io::{Result, Write}; +use std::sync::{Arc, Mutex}; + +/// Shared CRC32 handle that can be cloned and read independently of the writer. +#[derive(Clone)] +pub struct CrcHandle { + hasher: Arc>, +} + +impl CrcHandle { + pub fn crc32(&self) -> u32 { + self.hasher.lock().unwrap().clone().finalize() + } +} + +/// A writer wrapper that computes CRC32 incrementally on every write. +/// The CRC can be read via a `CrcHandle` without consuming the writer. +pub struct CrcWriter { + inner: W, + hasher: Arc>, +} + +impl CrcWriter { + pub fn new(inner: W) -> (Self, CrcHandle) { + let hasher = Arc::new(Mutex::new(crc32fast::Hasher::new())); + let handle = CrcHandle { hasher: hasher.clone() }; + (Self { inner, hasher }, handle) + } +} + +impl Write for CrcWriter { + fn write(&mut self, buf: &[u8]) -> Result { + let n = self.inner.write(buf)?; + self.hasher.lock().unwrap().update(&buf[..n]); + Ok(n) + } + + fn flush(&mut self) -> Result<()> { + self.inner.flush() + } +} 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 cb31f768f6dfe..233c7aef36923 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 @@ -215,12 +215,15 @@ pub unsafe extern "C" fn parquet_on_settings_update( page_size_bytes: i64, page_row_limit: i64, dict_size_bytes: i64, - row_group_size_bytes: i64, bloom_filter_enabled: i64, bloom_filter_fpp: f64, bloom_filter_ndv: i64, sort_in_memory_threshold_bytes: i64, sort_batch_size: i64, + row_group_max_rows: i64, + merge_batch_size: i64, + merge_rayon_threads: i64, + merge_io_threads: i64, ) -> i64 { let index_name = str_from_raw(index_name_ptr, index_name_len) .map_err(|e| format!("parquet_on_settings_update index_name: {}", e))?.to_string(); @@ -245,12 +248,15 @@ pub unsafe extern "C" fn parquet_on_settings_update( page_size_bytes: opt_usize(page_size_bytes), page_row_limit: opt_usize(page_row_limit), dict_size_bytes: opt_usize(dict_size_bytes), - row_group_size_bytes: opt_usize(row_group_size_bytes), bloom_filter_enabled: opt_bool(bloom_filter_enabled), bloom_filter_fpp: opt_f64(bloom_filter_fpp), bloom_filter_ndv: opt_u64(bloom_filter_ndv), sort_in_memory_threshold_bytes: opt_u64(sort_in_memory_threshold_bytes), sort_batch_size: opt_usize(sort_batch_size), + row_group_max_rows: opt_usize(row_group_max_rows), + merge_batch_size: opt_usize(merge_batch_size), + merge_rayon_threads: opt_usize(merge_rayon_threads), + merge_io_threads: opt_usize(merge_io_threads), ..Default::default() }; @@ -326,3 +332,84 @@ pub unsafe extern "C" fn parquet_merge_files( .map(|_| 0) .map_err(|e| format!("{}", e)) } + +// --------------------------------------------------------------------------- +// Parquet reader (for test verification) +// --------------------------------------------------------------------------- + +/// Reads a parquet file and returns its contents as a JSON string. +/// Each row is a JSON object. The result is a JSON array of objects. +/// The JSON bytes are written into `out_buf`, actual length into `out_len`. +/// Returns 0 on success. +#[ffm_safe] +#[no_mangle] +pub unsafe extern "C" fn parquet_read_as_json( + file_ptr: *const u8, + file_len: i64, + out_buf: *mut u8, + buf_capacity: i64, + out_len: *mut i64, +) -> i64 { + use arrow::array::Array; + + let filename = str_from_raw(file_ptr, file_len) + .map_err(|e| format!("parquet_read_as_json: {}", e))?.to_string(); + + let file = std::fs::File::open(&filename) + .map_err(|e| format!("Failed to open {}: {}", filename, e))?; + let builder = parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder::try_new(file) + .map_err(|e| format!("Failed to read parquet: {}", e))?; + let reader = builder.with_batch_size(8192).build() + .map_err(|e| format!("Failed to build reader: {}", e))?; + + let mut rows: Vec = Vec::new(); + for batch_result in reader { + let batch = batch_result.map_err(|e| format!("Read error: {}", e))?; + let schema = batch.schema(); + for row_idx in 0..batch.num_rows() { + let mut obj = serde_json::Map::new(); + for (col_idx, field) in schema.fields().iter().enumerate() { + let col = batch.column(col_idx); + let val = if col.is_null(row_idx) { + serde_json::Value::Null + } else { + match col.data_type() { + arrow::datatypes::DataType::Int32 => { + let arr = col.as_any().downcast_ref::().unwrap(); + serde_json::Value::Number(arr.value(row_idx).into()) + } + arrow::datatypes::DataType::Int64 => { + let arr = col.as_any().downcast_ref::().unwrap(); + serde_json::Value::Number(arr.value(row_idx).into()) + } + arrow::datatypes::DataType::Utf8 => { + let arr = col.as_any().downcast_ref::().unwrap(); + serde_json::Value::String(arr.value(row_idx).to_string()) + } + arrow::datatypes::DataType::Boolean => { + let arr = col.as_any().downcast_ref::().unwrap(); + serde_json::Value::Bool(arr.value(row_idx)) + } + arrow::datatypes::DataType::Float64 => { + let arr = col.as_any().downcast_ref::().unwrap(); + serde_json::json!(arr.value(row_idx)) + } + _ => serde_json::Value::String(format!("", col.data_type())), + } + }; + obj.insert(field.name().clone(), val); + } + rows.push(serde_json::Value::Object(obj)); + } + } + + let json_str = serde_json::to_string(&rows) + .map_err(|e| format!("JSON serialization failed: {}", e))?; + let bytes = json_str.as_bytes(); + if bytes.len() > buf_capacity as usize { + return Err(format!("JSON output ({} bytes) exceeds buffer capacity ({})", bytes.len(), buf_capacity)); + } + std::ptr::copy_nonoverlapping(bytes.as_ptr(), out_buf, bytes.len()); + *out_len = bytes.len() as i64; + Ok(0) +} 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 ff62af3543296..2ce15506f12c4 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 @@ -18,6 +18,7 @@ pub mod native_settings; pub mod field_config; pub mod writer_properties_builder; pub mod rate_limited_writer; +pub mod crc_writer; pub mod merge; pub use native_settings::NativeSettings; 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 8eb8782f2076e..e2a07c2efffeb 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 @@ -19,6 +19,7 @@ use parquet::schema::types::SchemaDescriptor; use rayon::prelude::*; use tokio::sync::{mpsc as tokio_mpsc, oneshot}; +use crate::crc_writer::CrcWriter; use crate::rate_limited_writer::RateLimitedWriter; use crate::writer_properties_builder::WriterPropertiesBuilder; use crate::{log_debug, SETTINGS_STORE}; @@ -43,6 +44,7 @@ pub struct MergeContext { row_group_index: usize, next_row_id: i64, total_rows_written: usize, + rayon_threads: Option, } impl MergeContext { @@ -54,6 +56,8 @@ impl MergeContext { output_path: &str, index_name: &str, output_flush_rows: usize, + rayon_threads: Option, + io_threads: Option, ) -> MergeResult { if let Some(parent) = Path::new(output_path).parent() { if !parent.exists() { @@ -90,15 +94,17 @@ impl MergeContext { let throttled_writer = RateLimitedWriter::new(output_file, RATE_LIMIT_MB_PER_SEC).map_err(MergeError::Io)?; + let (crc_writer, crc_handle) = CrcWriter::new(throttled_writer); + let config = SETTINGS_STORE .get(index_name) .map(|r| r.clone()) .unwrap_or_default(); let writer_props = Arc::new(WriterPropertiesBuilder::build(&config)); - let writer = SerializedFileWriter::new(throttled_writer, parquet_root, writer_props)?; + 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); + let io_tx = spawn_io_task(writer, crc_handle, io_threads); Ok(Self { data_schema, @@ -111,6 +117,7 @@ impl MergeContext { row_group_index: 0, next_row_id: 0, total_rows_written: 0, + rayon_threads, }) } @@ -152,22 +159,19 @@ impl MergeContext { .rg_writer_factory .create_column_writers(self.row_group_index)?; - let mut leaves_and_writers = Vec::new(); - { - let mut writer_iter = col_writers.into_iter(); - for (arr, field) in with_id.columns().iter().zip(self.output_schema.fields()) { - for leaf in compute_leaves(field, arr)? { - let col_writer = writer_iter.next().ok_or_else(|| { - MergeError::Logic("Fewer column writers than leaf columns".into()) - })?; - leaves_and_writers.push((leaf, col_writer)); + let leaves_and_writers = match Self::pair_leaves_with_writers(&with_id, &self.output_schema, col_writers) { + Ok(paired) => paired, + Err((err, remaining)) => { + for w in remaining { + let _ = w.close(); } + return Err(err); } - } + }; let chunk_results: Vec< Result, - > = get_merge_pool().install(|| { + > = get_merge_pool(self.rayon_threads).install(|| { leaves_and_writers .into_par_iter() .map(|(leaf, mut col_writer)| { @@ -201,12 +205,44 @@ impl MergeContext { Ok(()) } - /// Final flush + close the IO task. Returns Parquet metadata. - pub fn finish(mut self) -> MergeResult { + /// Pairs leaf arrays with column writers, returning unconsumed writers on error + /// so the caller can close them. + fn pair_leaves_with_writers( + batch: &RecordBatch, + schema: &Arc, + col_writers: Vec, + ) -> Result< + Vec<(parquet::arrow::arrow_writer::ArrowLeafColumn, parquet::arrow::arrow_writer::ArrowColumnWriter)>, + (MergeError, Vec), + > { + let mut writer_iter = col_writers.into_iter(); + let mut paired = Vec::new(); + for (arr, field) in batch.columns().iter().zip(schema.fields()) { + let leaves = match compute_leaves(field, arr) { + Ok(l) => l, + Err(e) => return Err((e.into(), writer_iter.collect())), + }; + for leaf in leaves { + match writer_iter.next() { + Some(w) => paired.push((leaf, w)), + None => { + return Err(( + MergeError::Logic("Fewer column writers than leaf columns".into()), + Vec::new(), + )) + } + } + } + } + Ok(paired) + } + + /// Final flush + close the IO task. Returns Parquet metadata and CRC32. + pub fn finish(mut self) -> MergeResult<(parquet::file::metadata::ParquetMetaData, u32)> { self.flush()?; let (reply_tx, reply_rx) = - oneshot::channel::>(); + oneshot::channel::>(); self.io_tx .blocking_send(IoCommand::Close(reply_tx)) 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 04a8e32223514..b530820ef3c98 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 @@ -139,7 +139,7 @@ impl FileCursor { let reader = Arc::clone(&self.reader); let tx = self.prefetch_tx.clone(); - get_merge_pool().spawn(move || { + get_merge_pool(None).spawn(move || { let mut reader = reader.lock().unwrap(); let result = match reader.next() { Some(Ok(batch)) if batch.num_rows() > 0 => Some(Ok(batch)), 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 6253d85d358df..1cbd4f41aa4ca 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 @@ -18,6 +18,7 @@ use tokio::runtime::Runtime; use tokio::sync::{mpsc as tokio_mpsc, oneshot}; use tokio::task::JoinHandle; +use crate::crc_writer::CrcWriter; use crate::rate_limited_writer::RateLimitedWriter; use crate::log_error; @@ -26,20 +27,16 @@ use super::error::{MergeError, MergeResult}; // Constants // ============================================================================= -/// Number of rows to request per Parquet read batch. -pub const BATCH_SIZE: usize = 100_000; - -/// Approximate number of rows to buffer before flushing a row group. -pub const OUTPUT_FLUSH_ROWS: usize = 1_000_000; - /// Disk write rate limit in MB/s. pub const RATE_LIMIT_MB_PER_SEC: f64 = 20.0; -/// Number of threads in the shared Rayon pool for parallel column encoding. -const RAYON_NUM_THREADS: usize = 4; - -/// Number of Tokio worker threads for async IO. -const TOKIO_WORKER_THREADS: usize = 4; +/// Default thread count for merge pools: max(1, num_cpus / 8). +fn default_merge_threads() -> usize { + std::thread::available_parallelism() + .map(|n| n.get() / 8) + .unwrap_or(1) + .max(1) +} /// Bounded channel capacity between the merge loop and the IO task. const IO_CHANNEL_BUFFER: usize = 2; @@ -50,10 +47,11 @@ const IO_CHANNEL_BUFFER: usize = 2; static MERGE_POOL: OnceLock = OnceLock::new(); -pub fn get_merge_pool() -> &'static ThreadPool { +pub fn get_merge_pool(num_threads: Option) -> &'static ThreadPool { MERGE_POOL.get_or_init(|| { + let n = num_threads.unwrap_or_else(default_merge_threads); rayon::ThreadPoolBuilder::new() - .num_threads(RAYON_NUM_THREADS) + .num_threads(n) .thread_name(|idx| format!("parquet-merge-{}", idx)) .build() .expect("Failed to build parquet-merge Rayon thread pool") @@ -66,10 +64,11 @@ pub fn get_merge_pool() -> &'static ThreadPool { static IO_RUNTIME: OnceLock = OnceLock::new(); -fn get_io_runtime() -> &'static Runtime { +fn get_io_runtime(num_threads: Option) -> &'static Runtime { IO_RUNTIME.get_or_init(|| { + let n = num_threads.unwrap_or_else(default_merge_threads); tokio::runtime::Builder::new_multi_thread() - .worker_threads(TOKIO_WORKER_THREADS) + .worker_threads(n) .thread_name("parquet-io") .enable_all() .build() @@ -81,10 +80,13 @@ fn get_io_runtime() -> &'static Runtime { // IO task protocol // ============================================================================= +/// Writer type used by the IO task: CRC → rate-limit → file. +pub type MergeWriter = CrcWriter>; + /// Commands sent from the merge loop to the background IO task. pub enum IoCommand { WriteRowGroup(Vec), - Close(oneshot::Sender>), + Close(oneshot::Sender>), } async fn drain_on_error(rx: &mut tokio_mpsc::Receiver, msg: &str) { @@ -104,14 +106,16 @@ async fn drain_on_error(rx: &mut tokio_mpsc::Receiver, msg: &str) { /// but is **not** awaited immediately — this allows the merge loop to prepare /// the next row group while the current one is still being flushed to disk. pub fn spawn_io_task( - writer: SerializedFileWriter>, + writer: SerializedFileWriter, + crc_handle: crate::crc_writer::CrcHandle, + io_threads: Option, ) -> tokio_mpsc::Sender { let (tx, mut rx) = tokio_mpsc::channel::(IO_CHANNEL_BUFFER); - get_io_runtime().spawn(async move { - let mut writer: Option>> = Some(writer); + get_io_runtime(io_threads).spawn(async move { + let mut writer: Option> = Some(writer); let mut in_flight: Option< - JoinHandle>>>, + JoinHandle>>, > = None; while let Some(cmd) = rx.recv().await { @@ -165,8 +169,10 @@ pub fn spawn_io_task( } let w = writer.take().unwrap(); + let crc = crc_handle.clone(); let result = tokio::task::spawn_blocking(move || { - w.close().map_err(MergeError::from) + let metadata = w.close().map_err(MergeError::from)?; + Ok((metadata, crc.crc32())) }) .await; diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/schema.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/schema.rs index f952bad63fb83..9cfbb10fd8c7d 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/schema.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/merge/schema.rs @@ -17,13 +17,13 @@ use parquet::schema::types::Type; use super::error::MergeResult; /// Reserved column name for the synthetic row identifier added during merge. -pub const ROW_ID_COLUMN_NAME: &str = "___row_id"; +pub const ROW_ID_COLUMN_NAME: &str = "__row_id__"; /// Builds the output Parquet schema as the union of pre-read schema descriptors. /// /// The output schema contains every column seen across all inputs, except: -/// - Any existing `___row_id` column is removed. -/// - A fresh `___row_id` INT64 REQUIRED column is appended at the end. +/// - Any existing `__row_id__` column is removed. +/// - A fresh `__row_id__` INT64 REQUIRED column is appended at the end. pub fn build_parquet_root_schema( schema_descriptors: &[parquet::schema::types::SchemaDescriptor], ) -> MergeResult> { @@ -54,7 +54,7 @@ pub fn build_parquet_root_schema( Ok(Arc::new(parquet_root)) } -/// Returns column indices that exclude `___row_id`, for use as a projection mask. +/// Returns column indices that exclude `__row_id__`, for use as a projection mask. pub fn projection_indices_excluding_row_id(schema: &ArrowSchema) -> Vec { schema .fields() @@ -66,7 +66,7 @@ pub fn projection_indices_excluding_row_id(schema: &ArrowSchema) -> Vec { } -/// Appends a `___row_id` column with sequential values `[start_id, start_id + N)` +/// Appends a `__row_id__` column with sequential values `[start_id, start_id + N)` /// to the given batch, producing a new batch with the output schema. pub fn append_row_id( batch: &RecordBatch, 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 35a6b603565b6..23f9121cb35c3 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 @@ -18,7 +18,7 @@ use crate::{log_debug, log_info}; use super::context::MergeContext; use super::cursor::FileCursor; use super::heap::{cmp_sort_values, get_sort_values, HeapItem}; -use super::io_task::{get_merge_pool, BATCH_SIZE, OUTPUT_FLUSH_ROWS}; +use super::io_task::get_merge_pool; use super::schema::ColumnMapping; /// Performs a streaming k-way merge with an explicit sort direction per column. @@ -29,11 +29,17 @@ pub fn merge_sorted( sort_columns: &[String], reverse_sorts: &[bool], nulls_first: &[bool], -) -> super::MergeResult<()> { - let batch_size = BATCH_SIZE; - let output_flush_rows = OUTPUT_FLUSH_ROWS; +) -> super::MergeResult { + let config = crate::writer::SETTINGS_STORE + .get(index_name) + .map(|r| r.clone()) + .unwrap_or_default(); + let batch_size = config.get_merge_batch_size(); + let output_flush_rows = config.get_row_group_max_rows(); + let rayon_threads = config.get_merge_rayon_threads(); + let io_threads = config.get_merge_io_threads(); if input_files.is_empty() { - return Ok(()); + return Ok(0); } if sort_columns.is_empty() { @@ -42,7 +48,7 @@ pub fn merge_sorted( )); } - let pool = get_merge_pool(); + let pool = get_merge_pool(rayon_threads); let direction_label = if reverse_sorts.iter().all(|&r| !r) { "ascending" } else if reverse_sorts.iter().all(|&r| r) { @@ -51,7 +57,7 @@ pub fn merge_sorted( "mixed" }; - log_info!( + log_debug!( "[RUST] Starting streaming merge ({}): {} input files, sort_columns={:?}, \ batch_size={}, flush_rows={}, merge_threads={}, output='{}'", direction_label, @@ -86,6 +92,8 @@ pub fn merge_sorted( output_path, index_name, output_flush_rows, + rayon_threads, + io_threads, )?; // Precompute column mappings per cursor (avoids per-batch name lookups) @@ -93,7 +101,7 @@ pub fn merge_sorted( .map(|s| ColumnMapping::new(s, ctx.data_schema())) .collect(); - log_info!( + log_debug!( "[RUST] Merge initialized ({}): {} cursors", direction_label, num_cursors @@ -202,15 +210,16 @@ pub fn merge_sorted( } // ── Phase 5: Close ────────────────────────────────────────────────── - let _metadata = ctx.finish()?; + let (_metadata, crc32) = ctx.finish()?; - log_info!( - "[RUST] Merge complete ({}): {} total rows written to '{}' in {} row groups", + log_debug!( + "[RUST] Merge complete ({}): {} total rows written to '{}' in {} row groups, crc32={:#010x}", direction_label, _metadata.file_metadata().num_rows(), output_path, - _metadata.num_row_groups() + _metadata.num_row_groups(), + crc32 ); - Ok(()) + Ok(crc32) } 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 89708764597aa..6618cb4477cad 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 @@ -17,17 +17,24 @@ use crate::{log_debug, log_info}; use super::context::MergeContext; use super::error::MergeResult; -use super::io_task::{BATCH_SIZE, OUTPUT_FLUSH_ROWS}; use super::schema::{projection_indices_excluding_row_id, ColumnMapping}; /// Unsorted merge: reads each input file sequentially, pads to union schema, -/// rewrites `___row_id` with globally sequential values. No sorting performed. +/// rewrites `__row_id__` with globally sequential values. No sorting performed. pub fn merge_unsorted( input_files: &[String], output_path: &str, index_name: &str, -) -> MergeResult<()> { - log_info!( +) -> MergeResult { + let config = crate::writer::SETTINGS_STORE + .get(index_name) + .map(|r| r.clone()) + .unwrap_or_default(); + let batch_size = config.get_merge_batch_size(); + let output_flush_rows = config.get_row_group_max_rows(); + let rayon_threads = config.get_merge_rayon_threads(); + let io_threads = config.get_merge_io_threads(); + log_debug!( "[RUST] Starting unsorted merge: {} input files, output='{}'", input_files.len(), output_path @@ -46,9 +53,9 @@ pub fn merge_unsorted( 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()?; + let reader = builder.with_batch_size(batch_size).with_projection(projection).build()?; - // The reader's schema is the projected schema (___row_id excluded). + // The reader's schema is the projected schema (__row_id__ excluded). arrow_schemas.push(reader.schema().as_ref().clone()); parquet_descriptors.push(parquet_descr); readers.push(reader); @@ -59,7 +66,9 @@ pub fn merge_unsorted( &parquet_descriptors, output_path, index_name, - OUTPUT_FLUSH_ROWS, + output_flush_rows, + rayon_threads, + io_threads, )?; // Precompute column mappings per reader @@ -82,14 +91,15 @@ pub fn merge_unsorted( } } - let _metadata = ctx.finish()?; + let (_metadata, crc32) = ctx.finish()?; - log_info!( - "[RUST] Unsorted merge complete: {} total rows written to '{}' in {} row groups", + log_debug!( + "[RUST] Unsorted merge complete: {} total rows written to '{}' in {} row groups, crc32={:#010x}", _metadata.file_metadata().num_rows(), output_path, - _metadata.num_row_groups() + _metadata.num_row_groups(), + crc32 ); - Ok(()) + Ok(crc32) } diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/native_settings.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/native_settings.rs index bc548c6129773..49e68b58437dc 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/native_settings.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/native_settings.rs @@ -18,7 +18,6 @@ pub struct NativeSettings { pub page_size_bytes: Option, pub page_row_limit: Option, pub dict_size_bytes: Option, - pub row_group_size_bytes: Option, pub field_configs: Option>, pub custom_settings: Option>, pub bloom_filter_enabled: Option, @@ -29,6 +28,10 @@ pub struct NativeSettings { pub nulls_first: Vec, pub sort_in_memory_threshold_bytes: Option, pub sort_batch_size: Option, + pub merge_batch_size: Option, + pub row_group_max_rows: Option, + pub merge_rayon_threads: Option, + pub merge_io_threads: Option, } impl NativeSettings { @@ -56,10 +59,6 @@ impl NativeSettings { self.dict_size_bytes.unwrap_or(2 * 1024 * 1024) } - pub fn get_row_group_size_bytes(&self) -> usize { - self.row_group_size_bytes.unwrap_or(128 * 1024 * 1024) - } - pub fn get_bloom_filter_enabled(&self) -> bool { self.bloom_filter_enabled.unwrap_or(true) } @@ -87,6 +86,22 @@ impl NativeSettings { pub fn get_sort_batch_size(&self) -> usize { self.sort_batch_size.unwrap_or(8192) } + + pub fn get_merge_batch_size(&self) -> usize { + self.merge_batch_size.unwrap_or(100_000) + } + + pub fn get_row_group_max_rows(&self) -> usize { + self.row_group_max_rows.unwrap_or(1_000_000) + } + + pub fn get_merge_rayon_threads(&self) -> Option { + self.merge_rayon_threads + } + + pub fn get_merge_io_threads(&self) -> Option { + self.merge_io_threads + } } #[cfg(test)] 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 284e1e4201db0..d244e306adec6 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 @@ -14,11 +14,11 @@ use lazy_static::lazy_static; use parquet::arrow::{arrow_reader::ParquetRecordBatchReaderBuilder, ArrowWriter}; use parquet::file::reader::{FileReader, SerializedFileReader}; use std::fs::File; -use std::io::Read; use std::path::Path; use std::sync::{Arc, Mutex}; use crate::{log_error, log_debug}; +use crate::crc_writer::CrcWriter; use crate::merge::{merge_sorted, schema::ROW_ID_COLUMN_NAME}; use crate::native_settings::NativeSettings; use crate::writer_properties_builder::WriterPropertiesBuilder; @@ -33,8 +33,9 @@ pub struct FinalizeResult { /// 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 { - writer: Arc>>, + writer: Arc>>>, settings: NativeSettings, + crc_handle: crate::crc_writer::CrcHandle, } lazy_static! { @@ -93,6 +94,7 @@ impl NativeParquetWriter { log_debug!("Schema created with {} fields", schema.fields().len()); let file = File::create(&temp_filename)?; + let (crc_file, crc_handle) = CrcWriter::new(file); let mut settings: NativeSettings = SETTINGS_STORE .get(&index_name) @@ -107,11 +109,12 @@ impl NativeParquetWriter { SETTINGS_STORE.insert(index_name, settings.clone()); - let writer = ArrowWriter::try_new(file, schema, Some(props))?; + let writer = ArrowWriter::try_new(crc_file, schema, Some(props))?; WRITERS.insert(temp_filename, WriterState { writer: Arc::new(Mutex::new(writer)), settings, + crc_handle, }); Ok(()) @@ -157,17 +160,17 @@ impl NativeParquetWriter { log_debug!("finalize_writer called for file: {} (temp: {})", filename, temp_filename); if let Some((_, state)) = WRITERS.remove(&temp_filename) { - let WriterState { writer: writer_arc, settings } = state; + let WriterState { writer: writer_arc, settings, crc_handle } = state; let index_name = settings.index_name.as_deref().unwrap_or(""); match Arc::try_unwrap(writer_arc) { Ok(mutex) => { let writer = mutex.into_inner().unwrap(); match writer.close() { Ok(_) => { + let temp_crc32 = crc_handle.crc32(); log_debug!("Successfully closed temp writer for: {}", temp_filename); - // _file is dropped here, closing the file handle - Self::sort_and_rewrite_parquet(&temp_filename, &filename, index_name, &settings.sort_columns, &settings.reverse_sorts, &settings.nulls_first)?; + let crc32 = Self::sort_and_rewrite_parquet(&temp_filename, &filename, index_name, &settings.sort_columns, &settings.reverse_sorts, &settings.nulls_first, temp_crc32)?; if Path::new(&temp_filename).exists() { if let Err(e) = std::fs::remove_file(&temp_filename) { @@ -175,8 +178,6 @@ impl NativeParquetWriter { } } - // Compute CRC32 by reading the final sorted file - let crc32 = Self::compute_file_crc32(&filename)?; log_debug!("CRC32 for file {}: {:#010x}", filename, crc32); // Keep a handle for sync_to_disk @@ -207,17 +208,6 @@ impl NativeParquetWriter { } } - fn compute_file_crc32(path: &str) -> Result> { - let mut file = File::open(path)?; - let mut hasher = crc32fast::Hasher::new(); - let mut buf = [0u8; 64 * 1024]; - loop { - let n = file.read(&mut buf)?; - if n == 0 { break; } - hasher.update(&buf[..n]); - } - Ok(hasher.finalize()) - } fn sort_and_rewrite_parquet( temp_filename: &str, @@ -226,7 +216,8 @@ impl NativeParquetWriter { sort_columns: &[String], reverse_sorts: &[bool], nulls_first: &[bool], - ) -> Result<(), Box> { + temp_crc32: u32, + ) -> Result> { log_debug!( "sort_and_rewrite_parquet: temp={}, output={}, sort_columns={:?}, reverse_sorts={:?}, nulls_first={:?}", temp_filename, output_filename, sort_columns, reverse_sorts, nulls_first @@ -235,7 +226,7 @@ impl NativeParquetWriter { if sort_columns.is_empty() { log_debug!("No sort columns specified, renaming temp file to final"); std::fs::rename(temp_filename, output_filename)?; - return Ok(()); + return Ok(temp_crc32); } let config = SETTINGS_STORE @@ -260,7 +251,7 @@ impl NativeParquetWriter { sort_columns: &[String], reverse_sorts: &[bool], nulls_first: &[bool], - ) -> Result<(), Box> { + ) -> Result> { log_debug!("Using in-memory sort for small file: {}", temp_filename); let file = File::open(temp_filename)?; @@ -274,7 +265,7 @@ impl NativeParquetWriter { _ => { log_debug!("No data to sort in file: {}", temp_filename); std::fs::rename(temp_filename, output_filename)?; - return Ok(()); + return Ok(0); } }; @@ -282,8 +273,8 @@ impl NativeParquetWriter { let sorted_batch = Self::sort_batch(&batch, sort_columns, reverse_sorts, nulls_first)?; let final_batch = Self::rewrite_row_ids(&sorted_batch, &schema)?; - Self::write_final_file(output_filename, index_name, &final_batch, schema)?; - Ok(()) + let crc32 = Self::write_final_file(output_filename, index_name, &final_batch, schema)?; + Ok(crc32) } /// For large files: read in batches, sort each batch individually, write each @@ -297,7 +288,7 @@ impl NativeParquetWriter { reverse_sorts: &[bool], nulls_first: &[bool], batch_size: usize, - ) -> Result<(), Box> { + ) -> Result> { log_debug!("Using streaming merge sort for large file: {}", temp_filename); let file = File::open(temp_filename)?; @@ -306,17 +297,18 @@ impl NativeParquetWriter { let mut chunk_paths: Vec = Vec::new(); let mut batch_count = 0; - let temp_dir = std::env::temp_dir(); + let chunk_dir = Path::new(output_filename).parent().unwrap_or_else(|| Path::new(".")); for batch_result in arrow_reader { let batch = batch_result?; let schema = batch.schema(); let sorted_batch = Self::sort_batch(&batch, sort_columns, reverse_sorts, nulls_first)?; - let chunk_filename = temp_dir - .join(format!("sort_chunk_{}_{}.parquet", batch_count, std::process::id())) + let chunk_filename = chunk_dir + .join(format!("temp_sort_chunk_{}_{}.parquet", batch_count, std::process::id())) .to_string_lossy() .to_string(); + // CRC for temp chunks is not needed, discard it Self::write_final_file(&chunk_filename, index_name, &sorted_batch, schema)?; chunk_paths.push(chunk_filename); @@ -326,13 +318,13 @@ impl NativeParquetWriter { if chunk_paths.is_empty() { log_debug!("No data to sort in file: {}", temp_filename); std::fs::rename(temp_filename, output_filename)?; - return Ok(()); + return Ok(0); } log_debug!("Created {} sorted chunks, merging via streaming k-way merge", batch_count); - // Use the streaming merge to produce the final sorted file - merge_sorted( + // merge_sorted returns CRC32 of the final merged output + let crc32 = merge_sorted( &chunk_paths, output_filename, index_name, @@ -346,7 +338,7 @@ impl NativeParquetWriter { let _ = std::fs::remove_file(path); } - Ok(()) + Ok(crc32) } fn sort_batch( @@ -384,7 +376,7 @@ impl NativeParquetWriter { Ok(RecordBatch::try_new(batch.schema(), sorted_columns?)?) } - /// If a ___row_id column exists, rewrite it with sequential values 0..N. + /// If a __row_id__ column exists, rewrite it with sequential values 0..N. fn rewrite_row_ids( batch: &RecordBatch, schema: &Arc, @@ -392,7 +384,7 @@ impl NativeParquetWriter { use arrow::array::Int64Array; if let Some(row_id_idx) = schema.fields().iter().position(|f| f.name() == ROW_ID_COLUMN_NAME) { - log_debug!("Rewriting ___row_id column with sequential values 0..{}", batch.num_rows()); + log_debug!("Rewriting __row_id__ column with sequential values 0..{}", batch.num_rows()); let sequential_ids = Int64Array::from_iter_values( (0..batch.num_rows() as u64).map(|x| x as i64) ); @@ -409,18 +401,20 @@ impl NativeParquetWriter { index_name: &str, batch: &RecordBatch, schema: Arc, - ) -> Result<(), Box> { + ) -> Result> { let config = SETTINGS_STORE .get(index_name) .map(|r| r.clone()) .unwrap_or_default(); let props = WriterPropertiesBuilder::build(&config); let file = File::create(output_filename)?; - let mut writer = ArrowWriter::try_new(file, schema, Some(props))?; + let (crc_file, crc_handle) = CrcWriter::new(file); + let mut writer = ArrowWriter::try_new(crc_file, schema, Some(props))?; writer.write(batch)?; writer.close()?; - log_debug!("Successfully wrote final file: {}", output_filename); - Ok(()) + let crc32 = crc_handle.crc32(); + log_debug!("Successfully wrote final file: {} (crc32={:#010x})", output_filename, crc32); + Ok(crc32) } pub fn sync_to_disk(filename: String) -> Result<(), Box> { diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/writer_properties_builder.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/writer_properties_builder.rs index 8af048cb2290a..d5676a33c9e0f 100644 --- a/sandbox/plugins/parquet-data-format/src/main/rust/src/writer_properties_builder.rs +++ b/sandbox/plugins/parquet-data-format/src/main/rust/src/writer_properties_builder.rs @@ -83,13 +83,14 @@ impl WriterPropertiesBuilder { builder } - /// Applies row group size and row count settings. + /// Applies row group row count limit. + /// In parquet-rs 57.x, `set_max_row_group_size` is a row count limit (not bytes). fn apply_row_group_settings( builder: parquet::file::properties::WriterPropertiesBuilder, config: &NativeSettings ) -> parquet::file::properties::WriterPropertiesBuilder { builder - .set_max_row_group_size(config.get_row_group_size_bytes()) + .set_max_row_group_size(config.get_row_group_max_rows()) } /// Applies dictionary encoding settings. 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 12d5a859e034a..c05f865381991 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 @@ -86,12 +86,12 @@ fn test_unsorted_merge_real_files() { assert_eq!(actual_rows, expected_rows, "Row count mismatch"); } -/// Verify that ___row_id in the output is monotonically increasing (0, 1, 2, ...). +/// Verify that __row_id__ in the output is monotonically increasing (0, 1, 2, ...). fn verify_row_id_order(path: &str) { let file = File::open(path).unwrap(); let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); let schema = builder.schema().clone(); - let col_idx = schema.index_of("___row_id").expect("___row_id not in output"); + let col_idx = schema.index_of("__row_id__").expect("__row_id__ not in output"); let reader = builder.build().unwrap(); let mut expected: i64 = 0; @@ -99,14 +99,14 @@ fn verify_row_id_order(path: &str) { let batch = batch.unwrap(); let col = batch.column(col_idx).as_any() .downcast_ref::() - .expect("___row_id should be Int64"); + .expect("__row_id__ should be Int64"); for i in 0..col.len() { - assert!(!col.is_null(i), "___row_id should never be null"); - assert_eq!(col.value(i), expected, "___row_id gap at row {}", expected); + assert!(!col.is_null(i), "__row_id__ should never be null"); + assert_eq!(col.value(i), expected, "__row_id__ gap at row {}", expected); expected += 1; } } - println!("Verified ___row_id is sequential 0..{}", expected); + println!("Verified __row_id__ is sequential 0..{}", expected); } @@ -144,7 +144,7 @@ fn test_sorted_merge_real_files() { println!("Output rows: {}", actual_rows); assert_eq!(actual_rows, expected_rows, "Row count mismatch"); - // Verify ___row_id is sequential 0..N + // Verify __row_id__ is sequential 0..N verify_row_id_order(&output_str); // Verify EventDate is non-decreasing in the merged output