diff --git a/sandbox/libs/dataformat-native/src/main/java/org/opensearch/nativebridge/spi/NativeCall.java b/sandbox/libs/dataformat-native/src/main/java/org/opensearch/nativebridge/spi/NativeCall.java index ca52c3bc5e643..f8b617046ed81 100644 --- a/sandbox/libs/dataformat-native/src/main/java/org/opensearch/nativebridge/spi/NativeCall.java +++ b/sandbox/libs/dataformat-native/src/main/java/org/opensearch/nativebridge/spi/NativeCall.java @@ -268,6 +268,49 @@ public static void invokeVoid(MethodHandle handle, Object... args) { } } + // ---- Static invocation (no Arena) ---- + // + // The instance invoke / invokeIO above run inside a NativeCall so short-lived arguments + // (marshalled strings, out-pointers) can be allocated in the instance's Arena and freed on + // close. The two static variants below are for callers that already own every argument in + // their own memory and so need no Arena — chiefly the column-reader read path, where + // ParquetColumnReader keeps its buffers in a pooled Arena reused across many reads. Using the + // instance form there would allocate and immediately close a throwaway Arena on every read. + // Being static, these allocate nothing; they only run the handle and check the status. They are + // also stateless (they touch only their parameters and locals), so concurrent calls from + // different query threads share no mutable state. + + /** + * Invokes a {@code long}-returning native call and checks its status: returns the value + * when {@code >= 0} (positive status codes are passed through), throws {@link IOException} + * when {@code < 0}. Static (no {@link Arena}), so the caller must own every argument. + */ + public static long invokeIOStatic(MethodHandle handle, Object... args) throws IOException { + try { + long result = (long) handle.invokeWithArguments(args); + return NativeLibraryLoader.checkResultIO(result); + } catch (IOException | RuntimeException e) { + throw e; + } catch (Throwable t) { + throw new IOException(t); + } + } + + /** + * Same as {@link #invokeIOStatic} but throws {@link RuntimeException} on a {@code < 0} + * status instead of {@link IOException}. + */ + public static long invokeStatic(MethodHandle handle, Object... args) { + try { + long result = (long) handle.invokeWithArguments(args); + return NativeLibraryLoader.checkResult(result); + } catch (RuntimeException e) { + throw e; + } catch (Throwable t) { + throw new RuntimeException(t); + } + } + @Override public void close() { closed = true; diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/RustBridge.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/bridge/RustBridge.java index a60d0e758c63e..27edad2011e7c 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 @@ -34,7 +34,8 @@ import java.util.function.LongSupplier; /** - * FFM bridge to the native Rust parquet writer library. + * FFM bridge to the native Rust parquet library: the writer/merge path plus the DocValues codec's + * column-reader read path. */ public class RustBridge { @@ -57,6 +58,25 @@ public class RustBridge { private static final MethodHandle GET_POOL_STATS; private static final MethodHandle REGISTER_OVERCOMMIT_CALLBACKS; + // DocValues codec — Parquet column-reader FFI handles + private static final MethodHandle OPEN_COLUMN_READER; + private static final MethodHandle CLOSE_COLUMN_READER; + private static final MethodHandle OPEN_COLUMN_READER_COUNT; + private static final MethodHandle READ_VALUE_AT_ROW; + private static final MethodHandle READ_REPEATED_AT_ROW; + private static final MethodHandle GET_COLUMN_NUM_PAGES; + private static final MethodHandle GET_COLUMN_PAGE_INDEX; + private static final MethodHandle DECODE_PAGE_AT_ROW; + + /** + * Positive status code returned by the column-reader read/decode functions + * when a caller out-buffer was too small. The required sizes are written to + * the out-parameters so the caller can grow its buffers and retry once. This + * is distinct from the {@code < 0} error-pointer convention (a small negative + * constant would be dereferenced as an error pointer and corrupt memory). + */ + public static final long RC_OVERFLOW = 1L; + static { SymbolLookup lib = NativeLibraryLoader.symbolLookup(); Linker linker = Linker.nativeLinker(); @@ -288,6 +308,90 @@ public class RustBridge { lib.find("parquet_register_overcommit_callbacks").orElseThrow(), FunctionDescriptor.ofVoid(ValueLayout.ADDRESS, ValueLayout.ADDRESS) ); + + // ── DocValues codec column-reader functions ── + OPEN_COLUMN_READER = linker.downcallHandle( + lib.find("parquet_open_column_reader").orElseThrow(), + FunctionDescriptor.of( + ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS, // file_ptr + ValueLayout.JAVA_LONG, // file_len + ValueLayout.ADDRESS, // col_ptr + ValueLayout.JAVA_LONG, // col_len + ValueLayout.JAVA_INT // expected_type + ) + ); + CLOSE_COLUMN_READER = linker.downcallHandle( + lib.find("parquet_close_column_reader").orElseThrow(), + FunctionDescriptor.of(ValueLayout.JAVA_LONG, ValueLayout.JAVA_LONG) + ); + OPEN_COLUMN_READER_COUNT = linker.downcallHandle( + lib.find("parquet_open_column_reader_count").orElseThrow(), + FunctionDescriptor.of(ValueLayout.JAVA_LONG) + ); + READ_VALUE_AT_ROW = linker.downcallHandle( + lib.find("parquet_read_value_at_row").orElseThrow(), + FunctionDescriptor.of( + ValueLayout.JAVA_LONG, + ValueLayout.JAVA_LONG, // handle + ValueLayout.JAVA_LONG, // row + ValueLayout.ADDRESS, // out_present + ValueLayout.ADDRESS, // out_long + ValueLayout.ADDRESS, // out_buf + ValueLayout.JAVA_LONG, // out_buf_cap + ValueLayout.ADDRESS // out_len + ) + ); + READ_REPEATED_AT_ROW = linker.downcallHandle( + lib.find("parquet_read_repeated_at_row").orElseThrow(), + FunctionDescriptor.of( + ValueLayout.JAVA_LONG, + ValueLayout.JAVA_LONG, // handle + ValueLayout.JAVA_LONG, // row + ValueLayout.ADDRESS, // out_count + ValueLayout.ADDRESS, // out_longs + ValueLayout.JAVA_LONG, // out_long_cap + ValueLayout.ADDRESS, // out_byte_buf + ValueLayout.ADDRESS, // out_byte_offsets + ValueLayout.JAVA_LONG // out_byte_buf_cap + ) + ); + GET_COLUMN_NUM_PAGES = linker.downcallHandle( + lib.find("parquet_get_column_num_pages").orElseThrow(), + FunctionDescriptor.of(ValueLayout.JAVA_LONG, ValueLayout.JAVA_LONG) + ); + GET_COLUMN_PAGE_INDEX = linker.downcallHandle( + lib.find("parquet_get_column_page_index").orElseThrow(), + FunctionDescriptor.of( + ValueLayout.JAVA_LONG, + ValueLayout.JAVA_LONG, // handle + ValueLayout.ADDRESS, // out_first_row + ValueLayout.ADDRESS, // out_file_offset + ValueLayout.ADDRESS, // out_compressed_size + ValueLayout.ADDRESS, // out_null_count + ValueLayout.ADDRESS, // out_min_long + ValueLayout.ADDRESS, // out_max_long + ValueLayout.JAVA_LONG, // out_buf_capacity + ValueLayout.ADDRESS // out_actual_pages + ) + ); + DECODE_PAGE_AT_ROW = linker.downcallHandle( + lib.find("parquet_decode_page_at_row").orElseThrow(), + FunctionDescriptor.of( + ValueLayout.JAVA_LONG, + ValueLayout.JAVA_LONG, // handle + ValueLayout.JAVA_LONG, // row + ValueLayout.ADDRESS, // out_first_row + ValueLayout.ADDRESS, // out_last_row + ValueLayout.ADDRESS, // out_value_buf + ValueLayout.JAVA_LONG, // out_value_buf_cap + ValueLayout.ADDRESS, // out_value_actual_len + ValueLayout.ADDRESS, // out_byte_offsets + ValueLayout.JAVA_LONG, // out_byte_offsets_cap + ValueLayout.ADDRESS, // out_presence_bitset + ValueLayout.JAVA_LONG // out_presence_bits_cap + ) + ); } public static void initLogger() {} @@ -699,6 +803,174 @@ public static String readAsJson(String file) throws IOException { } } + // ──────────────────────────────────────────────────────────────────────── + // DocValues codec — Parquet column-reader bridge methods + // + // These methods own the MethodHandle invocations; the MemorySegment scratch + // buffers are supplied by the caller (ParquetColumnReader), which owns their + // lifecycle via its own Arena. Read/decode methods return the raw native + // status code so the caller can implement the grow-and-retry protocol on + // RC_OVERFLOW; a {@code < 0} result is decoded into an IOException here. + // ──────────────────────────────────────────────────────────────────────── + + /** + * Invokes a {@code long}-returning handle on caller-owned segments and checks the + * result, without allocating a per-call {@link NativeCall} arena. A {@code < 0} + * result is decoded into an {@link IOException}; {@code >= 0} (including + * {@link #RC_OVERFLOW}) is returned as-is. + * + *

Delegates to {@link NativeCall#invokeIOStatic} so the underlying MethodHandle + * invocation lives in the native-bridge library (which permits it), keeping this plugin + * free of the banned {@code MethodHandle#invokeWithArguments} call. Used on the column-reader + * read path. + */ + private static long invokeChecked(MethodHandle handle, Object... args) throws IOException { + return NativeCall.invokeIOStatic(handle, args); + } + + /** + * Opens a per-column reader over {@code file} for {@code column}, validating that the + * column exists and its physical type matches {@code expectedTypeCode} + * (0=INT32, 1=INT64, 2=FLOAT, 3=DOUBLE, 4=BOOL, 5=BYTE_ARRAY). + * + * @return a {@code >= 0} opaque reader handle (a Rust-side i64; lives until {@link #closeColumnReader}) + * @throws IOException if the column is missing, the type mismatches, or the file cannot be read + */ + public static long openColumnReader(String file, String column, int expectedTypeCode) throws IOException { + try (var call = new NativeCall()) { + var f = call.str(file); + var c = call.str(column); + return call.invokeIO(OPEN_COLUMN_READER, f.segment(), f.len(), c.segment(), c.len(), expectedTypeCode); + } + } + + /** + * Closes a column reader handle, releasing its native file handle and buffers. The native side + * treats an unknown handle as an error, so callers must not double-close; {@code ParquetColumnReader} + * enforces this by tracking a closed sentinel and calling here at most once per handle. + * + * @throws IOException if the handle is unknown + */ + public static void closeColumnReader(long handle) throws IOException { + invokeChecked(CLOSE_COLUMN_READER, handle); + } + + /** + * Returns the number of currently open native column-reader handles. Test-only: used to assert + * that readers do not leak their native handles (every open is matched by a close). + */ + public static long openColumnReaderCount() { + return NativeCall.invokeStatic(OPEN_COLUMN_READER_COUNT); + } + + /** + * Reads the single value at {@code row} into caller-provided out-segments. + * The caller (ParquetColumnReader) owns the segments. Returns the native status: + * {@code 0} on success, {@link #RC_OVERFLOW} when {@code outBuf} was too small + * (required length written to {@code outLen}); throws on a native error. + */ + static long readValueAtRow( + long handle, + long row, + MemorySegment outPresent, + MemorySegment outLong, + MemorySegment outBuf, + long outBufCap, + MemorySegment outLen + ) throws IOException { + return invokeChecked(READ_VALUE_AT_ROW, handle, row, outPresent, outLong, outBuf, outBufCap, outLen); + } + + /** + * Reads all values at {@code row} for a repeated (multi-valued) column into caller-provided + * out-segments. Returns the native status: {@code 0} on success, {@link #RC_OVERFLOW} when a + * buffer was too small (required element count in {@code outCount}; required byte + * size in {@code outByteOffsets[count]} when only the byte buffer overflowed). + */ + static long readRepeatedAtRow( + long handle, + long row, + MemorySegment outCount, + MemorySegment outLongs, + long outLongCap, + MemorySegment outByteBuf, + MemorySegment outByteOffsets, + long outByteBufCap + ) throws IOException { + return invokeChecked(READ_REPEATED_AT_ROW, handle, row, outCount, outLongs, outLongCap, outByteBuf, outByteOffsets, outByteBufCap); + } + + /** Returns the number of pages in the column (used to pre-size the page-index arrays). */ + static long getColumnNumPages(long handle) throws IOException { + return invokeChecked(GET_COLUMN_NUM_PAGES, handle); + } + + /** + * Loads the column's per-page row-range jump table + stats into caller-provided + * parallel out-segments, each of capacity {@code outBufCapacity} pages. Returns the + * native status: {@code 0} on success, {@link #RC_OVERFLOW} when capacity is too + * small (true page count written to {@code outActualPages}). + */ + static long getColumnPageIndex( + long handle, + MemorySegment outFirstRow, + MemorySegment outFileOffset, + MemorySegment outCompressedSize, + MemorySegment outNullCount, + MemorySegment outMinLong, + MemorySegment outMaxLong, + long outBufCapacity, + MemorySegment outActualPages + ) throws IOException { + return invokeChecked( + GET_COLUMN_PAGE_INDEX, + handle, + outFirstRow, + outFileOffset, + outCompressedSize, + outNullCount, + outMinLong, + outMaxLong, + outBufCapacity, + outActualPages + ); + } + + /** + * Decodes the page containing {@code row} (values + presence bitset) + * into caller-provided out-segments. Returns the native status: {@code 0} on success, + * {@link #RC_OVERFLOW} when a buffer was too small (page row range and required value + * byte length are still written so the caller can size every buffer and retry once). + */ + static long decodePageAtRow( + long handle, + long row, + MemorySegment outFirstRow, + MemorySegment outLastRow, + MemorySegment outValueBuf, + long outValueBufCap, + MemorySegment outValueActualLen, + MemorySegment outByteOffsets, + long outByteOffsetsCap, + MemorySegment outPresenceBitset, + long outPresenceBitsCap + ) throws IOException { + return invokeChecked( + DECODE_PAGE_AT_ROW, + handle, + row, + outFirstRow, + outLastRow, + outValueBuf, + outValueBufCap, + outValueActualLen, + outByteOffsets, + outByteOffsetsCap, + outPresenceBitset, + outPresenceBitsCap + ); + } + private record MapArrays(NativeCall.StrArray keys, NativeCall.StrArray values) { } diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/codec/FieldTypeMapping.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/codec/FieldTypeMapping.java new file mode 100644 index 0000000000000..23d708230792c --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/codec/FieldTypeMapping.java @@ -0,0 +1,121 @@ +/* + * 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.codec; + +import org.apache.lucene.index.DocValuesType; + +import java.util.Locale; +import java.util.Map; + +/** + * Deterministic mapping from OpenSearch field mapping types to the Lucene DocValues type and + * the Parquet physical type the codec reads + * + *

Every Parquet column is written as {@code optional repeated }; whether a field + * is single- or multi-valued is decided by the OpenSearch mapping, not by this table. The DV + * type recorded here is the single-valued form; the matching repeated form + * ({@code SORTED_NUMERIC} for numerics, {@code SORTED_SET} for keyword/ip) is selected by the + * producer when a multi-valued iterator is requested. + */ +public final class FieldTypeMapping { + + /** + * The resolved DocValues + physical type for one OpenSearch mapping type. + * + * @param singleValued Lucene DV type served when the field holds one value per document + * (e.g. {@code NUMERIC} for {@code long}, {@code SORTED} for {@code keyword}) + * @param multiValued Lucene DV type served when the field is an array + * ({@code SORTED_NUMERIC} for numerics, {@code SORTED_SET} for keyword/ip; + * {@code NONE} for types with no array form, e.g. {@code text}) + * @param physical the Parquet physical type the column is stored as, read back by the codec + * + *

{@code boolean} maps to a numeric DV type ({@code NUMERIC} / {@code SORTED_NUMERIC}) because + * Lucene has no boolean DV type; it is stored as {@code 0}/{@code 1}, matching how core + * OpenSearch indexes boolean doc values. + */ + public record Mapping(DocValuesType singleValued, DocValuesType multiValued, ParquetPhysicalType physical) { + } + + // TODO(type-coverage): these writable OpenSearch types are not yet mapped, so the codec serves no + // doc values for them: half_float, unsigned_long, scaled_float, token_count, version, wildcard, + // constant_keyword, flat_object, match_only_text. Note scaled_float is stored as a raw scaled + // long, so it is not a plain INT64 mapping. forType and validate throw for unmapped types. + private static final Map BY_TYPE = Map.ofEntries( + Map.entry("boolean", new Mapping(DocValuesType.NUMERIC, DocValuesType.SORTED_NUMERIC, ParquetPhysicalType.BOOL)), + Map.entry("byte", new Mapping(DocValuesType.NUMERIC, DocValuesType.SORTED_NUMERIC, ParquetPhysicalType.INT32)), + Map.entry("short", new Mapping(DocValuesType.NUMERIC, DocValuesType.SORTED_NUMERIC, ParquetPhysicalType.INT32)), + Map.entry("integer", new Mapping(DocValuesType.NUMERIC, DocValuesType.SORTED_NUMERIC, ParquetPhysicalType.INT32)), + Map.entry("long", new Mapping(DocValuesType.NUMERIC, DocValuesType.SORTED_NUMERIC, ParquetPhysicalType.INT64)), + Map.entry("float", new Mapping(DocValuesType.NUMERIC, DocValuesType.SORTED_NUMERIC, ParquetPhysicalType.FLOAT)), + Map.entry("double", new Mapping(DocValuesType.NUMERIC, DocValuesType.SORTED_NUMERIC, ParquetPhysicalType.DOUBLE)), + Map.entry("date", new Mapping(DocValuesType.NUMERIC, DocValuesType.SORTED_NUMERIC, ParquetPhysicalType.INT64)), + Map.entry("date_nanos", new Mapping(DocValuesType.NUMERIC, DocValuesType.SORTED_NUMERIC, ParquetPhysicalType.INT64)), + Map.entry("keyword", new Mapping(DocValuesType.SORTED, DocValuesType.SORTED_SET, ParquetPhysicalType.BYTE_ARRAY)), + Map.entry("ip", new Mapping(DocValuesType.SORTED, DocValuesType.SORTED_SET, ParquetPhysicalType.BYTE_ARRAY)), + Map.entry("text", new Mapping(DocValuesType.BINARY, DocValuesType.NONE, ParquetPhysicalType.BYTE_ARRAY)), + Map.entry("binary", new Mapping(DocValuesType.BINARY, DocValuesType.NONE, ParquetPhysicalType.BYTE_ARRAY)) + ); + + private FieldTypeMapping() {} + + /** True if the codec has a Parquet DocValues mapping for the given OpenSearch mapping type. */ + public static boolean isSupported(String mappingType) { + return BY_TYPE.containsKey(mappingType); + } + + /** + * Returns the mapping for {@code mappingType}. + * + * @throws IllegalArgumentException if the mapping type has no Parquet DocValues mapping + */ + public static Mapping forType(String mappingType) { + Mapping m = BY_TYPE.get(mappingType); + if (m == null) { + throw new IllegalArgumentException( + String.format(Locale.ROOT, "Parquet DocValues codec has no mapping for OpenSearch type '%s'", mappingType) + ); + } + return m; + } + + /** + * Validates that the field's mapping type supports the requested Lucene DV type, throwing + * {@link IllegalArgumentException} naming the field and mapping type when incompatible. + * + *

The requested type may be the single- or multi-valued form of the mapping's DV type + * (e.g. requesting {@code SORTED_NUMERIC} for a {@code long} field, whose single-valued + * form is {@code NUMERIC}, is valid). + */ + public static void validate(String field, String mappingType, DocValuesType requested) { + Mapping m = BY_TYPE.get(mappingType); + if (m == null) { + throw new IllegalArgumentException( + String.format( + Locale.ROOT, + "field '%s' has mapping type '%s', which the Parquet DocValues codec does not support", + field, + mappingType + ) + ); + } + if (requested != m.singleValued() && requested != m.multiValued()) { + throw new IllegalArgumentException( + String.format( + Locale.ROOT, + "field '%s' (mapping type '%s') supports DocValues type %s/%s but %s was requested", + field, + mappingType, + m.singleValued(), + m.multiValued(), + requested + ) + ); + } + } +} diff --git a/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/codec/ParquetPhysicalType.java b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/codec/ParquetPhysicalType.java new file mode 100644 index 0000000000000..f0a9c0a9ba3e8 --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/main/java/org/opensearch/parquet/codec/ParquetPhysicalType.java @@ -0,0 +1,55 @@ +/* + * 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.codec; + +/** + * The Parquet physical type of a column, as understood by the native column reader. + * + *

The {@link #code()} values are the discriminants exchanged with the Rust FFM + * bridge ({@code parquet_open_column_reader}'s {@code expected_type} argument and the + * physical-type validation it performs). They must stay in lock-step with the + * {@code TYPE_*} constants in {@code ffm.rs}. + * + *

Parquet's {@code FIXED_LEN_BYTE_ARRAY} is folded into {@link #BYTE_ARRAY} because + * the codec treats both as opaque byte sequences for {@code BinaryDocValues} purposes. + */ +public enum ParquetPhysicalType { + /** 32-bit signed integer; raw bits sign-extended to a {@code long}. */ + INT32(0), + /** 64-bit signed integer; raw bits verbatim. */ + INT64(1), + /** 32-bit IEEE-754 float; raw bits via {@link Float#floatToRawIntBits}. */ + FLOAT(2), + /** 64-bit IEEE-754 double; raw bits via {@link Double#doubleToRawLongBits}. */ + DOUBLE(3), + /** Boolean; encoded as 0 or 1. */ + BOOL(4), + /** Variable- or fixed-length byte array (UTF-8 strings, IP, binary). */ + BYTE_ARRAY(5); + + private final int code; + + ParquetPhysicalType(int code) { + this.code = code; + } + + /** The native discriminant exchanged with the Rust FFM bridge. */ + public int code() { + return code; + } + + /** + * True for the fixed-width types whose values are exchanged as raw {@code long} bits (everything + * except {@code BYTE_ARRAY}, which is variable-width and exchanged as bytes + length). Named for + * width rather than "primitive" because Parquet's own taxonomy counts BYTE_ARRAY as a primitive. + */ + public boolean isFixedWidth() { + return this != BYTE_ARRAY; + } +} diff --git a/sandbox/plugins/parquet-data-format/src/main/rust/src/ffm.rs b/sandbox/plugins/parquet-data-format/src/main/rust/src/ffm.rs index 36f94820a37ff..a8cea7446fe00 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 @@ -6,15 +6,27 @@ * compatible open source license. */ -//! FFM bridge for the Parquet writer. +//! FFM bridge for the Parquet data format: the writer/merge path plus the DocValues codec's +//! column-reader read path. //! //! Return convention: `>= 0` success, `< 0` error pointer (negate to get ptr, //! call `native_error_message`/`native_error_free`). +use std::collections::HashMap; +use std::fs::File; use std::slice; use std::str; +use std::sync::atomic::{AtomicI64, Ordering}; +use std::sync::{Mutex, MutexGuard}; +use lazy_static::lazy_static; use native_bridge_common::{ffm_safe, log_debug}; +use parquet::basic::Type as PhysicalType; +use parquet::column::reader::{ColumnReader, ColumnReaderImpl}; +use parquet::data_type::DataType as ParquetDataType; +use parquet::file::page_index::column_index::ColumnIndexMetaData; +use parquet::file::reader::{FileReader, SerializedFileReader}; +use parquet::file::serialized_reader::ReadOptionsBuilder; use crate::field_config::FieldConfig; use crate::merge; @@ -1005,3 +1017,1447 @@ pub unsafe extern "C" fn parquet_get_pool_stats(out_buf: *mut i64) { *out_buf.add(i) = *val as i64; } } + +// =========================================================================== +// Parquet column reader +// --------------------------------------------------------------------------- +// +// Read-only per-column random-access reader used by the Lucene +// `ParquetDocValuesProducer`. The reader exposes the column's physical values +// by row position so the Java side can materialise per-document doc values. +// +// Return convention (consistent with the rest of this file): +// - `>= 0` success. For reads, `0` (`RC_OK`) means "value(s) written"; the +// positive sentinel `RC_OVERFLOW` (1) means "caller buffer too small — +// required sizes were written to the out-parameters, retry once". +// - `< 0` error pointer (negate, then `native_error_message`/`native_error_free`). +// +// Overflow is a positive status, not a negative one: negative returns are all +// reserved for error pointers, so a small negative constant would be +// dereferenced as an error pointer. This mirrors `parquet_finalize_writer`, +// which already returns `Ok(1)` for "no writer". + +/// Read succeeded; value(s) written to the caller buffers. +pub(crate) const RC_OK: i64 = 0; +/// A caller buffer was too small. Required sizes were written to the +/// out-parameters; the caller should grow its buffers and retry once. +pub(crate) const RC_OVERFLOW: i64 = 1; + +/// `expected_type` discriminants exchanged with Java +/// (matches `ParquetPhysicalType` on the Java side). +const TYPE_INT32: i32 = 0; +const TYPE_INT64: i32 = 1; +const TYPE_FLOAT: i32 = 2; +const TYPE_DOUBLE: i32 = 3; +const TYPE_BOOL: i32 = 4; +const TYPE_BYTE_ARRAY: i32 = 5; + +/// Per-column reader state, one instance per open handle in the registry. Owns the +/// file handle (via `SerializedFileReader`) and the row-group layout that maps a global +/// row position to a `(row_group, local_offset)` pair. See the per-field comments for +/// what each read function uses. +struct ColumnReaderState { + reader: SerializedFileReader, + /// Leaf column index within the Parquet schema descriptor. + leaf_idx: usize, + /// Physical type of the column (validated against the caller's expectation). + physical_type: PhysicalType, + /// True when the column has a repetition level > 0 (multi-valued). + repeated: bool, + /// Max definition level of the column (0 = required; >0 = optional/nested). + max_def_level: i16, + /// Total number of rows (records) in the file. + row_count: i64, + /// Global row index of the first row in each row group. + rg_first_row: Vec, + /// Number of rows in each row group. + rg_num_rows: Vec, + /// Per-page layout (row-range jump table + page stats), ascending by + /// `global_first_row`. Built once at `open()` from the Parquet OffsetIndex + + /// ColumnIndex when present, else one entry per row group as a fallback. + pages: Vec, + /// Decode cursor retained by `parquet_decode_page_at_row` across calls. `None` + /// until the first page decode and whenever it has been invalidated (row-group + /// change, backwards seek, or a decode error). The single-row readers + /// (`parquet_read_value_at_row`, `parquet_read_repeated_at_row`) do not use it. + cursor: Option, + /// Pre-allocated scratch buffers for page decoding, reused across + /// `parquet_decode_page_at_row` calls to avoid per-call heap allocation. The + /// Vecs keep their capacity across `.clear()`, so only the first decode allocates. + scratch: DecodeScratch, +} + +/// Retained cursor for `parquet_decode_page_at_row`. Keeps the typed column +/// reader open across consecutive page decodes within the same row group, so an +/// ascending scan skips forward from the current position instead of reopening +/// the row group (which re-reads row-group metadata and the dictionary page) for +/// every page. +struct CursorState { + /// The row group index this cursor was opened for. + rg_idx: usize, + /// The typed column reader, retained across pages within the same row group. + col_reader: ColumnReader, + /// Global row position: the next record read from this cursor would start at this row. + position: i64, +} + +/// Pre-allocated scratch buffers for page decoding. Kept on `ColumnReaderState` and cleared (not +/// reallocated) between calls, so steady-state decoding is allocation-free after the first miss. +struct DecodeScratch { + /// Definition levels returned by `read_records` (used to derive per-row presence). + def_levels: Vec, + /// Dense non-null values read from the page. One typed buffer per primitive type; a decode + /// uses only the buffer matching the column's type. Kept as separate Vecs (not one reused + /// byte buffer) so decoding is type-safe with no transmute. + values_i64: Vec, + values_i32: Vec, + values_f32: Vec, + values_f64: Vec, + values_bool: Vec, +} + +impl DecodeScratch { + fn new() -> Self { + DecodeScratch { + def_levels: Vec::new(), + values_i64: Vec::new(), + values_i32: Vec::new(), + values_f32: Vec::new(), + values_f64: Vec::new(), + values_bool: Vec::new(), + } + } +} + +/// One row-aligned page in the column (or one row group, in the no-page-index +/// fallback). All row indices are global (file-relative). +#[derive(Clone)] +struct PageEntry { + /// Global index of the first row in the page. + global_first_row: i64, + /// Number of rows in the page. + num_rows: i64, + /// Byte offset of the page in the file. 0 when unknown. + file_offset: i64, + /// Compressed page size in bytes. 0 when unknown. + compressed_size: i32, + /// Number of nulls in the page. -1 when unknown. + null_count: i64, + /// Min value raw bits; meaningful only for numeric columns with a + /// page index, else 0. + min_long: i64, + /// Max value raw bits; meaningful only for numeric columns with a + /// page index, else 0. + max_long: i64, + /// Row group containing the page. + rg_idx: usize, + /// Index of the page's first row within its row group. + local_first_row: i64, +} + +/// Gets or builds the per-file metadata cache entry. First call for a given path parses the +/// footer + page index and pre-computes every column's page layout; subsequent calls (other +/// columns, other queries) get an Arc clone in O(1) — no re-parse, no per-column layout work. +fn get_or_build_file_metadata(filename: &str) -> Result, String> { + let mut cache = FILE_METADATA_CACHE + .lock() + .map_err(|_| "file metadata cache mutex poisoned".to_string())?; + if let Some(entry) = cache.get(filename) { + return Ok(std::sync::Arc::clone(entry)); + } + + // First access to this file: parse footer + page index. + let file = File::open(filename).map_err(|e| format!("Failed to open '{}': {}", filename, e))?; + let options = ReadOptionsBuilder::new().with_page_index().build(); + let reader = SerializedFileReader::new_with_options(file, options) + .map_err(|e| format!("Failed to read parquet metadata '{}': {}", filename, e))?; + + let metadata = reader.metadata(); + let schema = metadata.file_metadata().schema_descr_ptr(); + let row_count = metadata.file_metadata().num_rows(); + + let n_rg = metadata.num_row_groups(); + let mut rg_first_row = Vec::with_capacity(n_rg); + let mut rg_num_rows = Vec::with_capacity(n_rg); + let mut acc = 0i64; + for i in 0..n_rg { + let rn = metadata.row_group(i).num_rows(); + rg_first_row.push(acc); + rg_num_rows.push(rn); + acc += rn; + } + + // Pre-compute per-column descriptors. + let columns: Vec<_> = (0..schema.num_columns()) + .map(|i| { + let d = schema.column(i); + (i, d.physical_type(), d.max_rep_level(), d.max_def_level()) + }) + .collect(); + + // Pre-compute page layouts for ALL columns (small per column; avoids per-column re-parse). + let mut column_pages = HashMap::new(); + for &(leaf_idx, phys, _, _) in &columns { + let pages = build_page_layout(metadata, leaf_idx, phys, &rg_first_row, &rg_num_rows); + column_pages.insert(leaf_idx, pages); + } + + let entry = std::sync::Arc::new(FileMetadataCache { + schema, + row_count, + rg_first_row, + rg_num_rows, + column_pages, + }); + cache.insert(filename.to_string(), std::sync::Arc::clone(&entry)); + Ok(entry) +} + +impl ColumnReaderState { + fn open(filename: &str, column: &str, expected_type: i32) -> Result { + // Use the node-level metadata cache — first call parses; subsequent calls are O(1). + let fmc = get_or_build_file_metadata(filename)?; + + // Resolve column from cached schema descriptor (no file I/O, no re-parse). + let mut found: Option<(usize, PhysicalType, i16, i16)> = None; + for i in 0..fmc.schema.num_columns() { + let descr = fmc.schema.column(i); + if descr.name() == column || descr.path().string() == column { + found = Some((i, descr.physical_type(), descr.max_rep_level(), descr.max_def_level())); + break; + } + } + let (leaf_idx, phys, max_rep, max_def) = found.ok_or_else(|| { + format!("Column '{}' not found in parquet file '{}'", column, filename) + })?; + + let actual = physical_type_code(phys); + if actual != expected_type { + return Err(format!( + "Column '{}' physical type mismatch in '{}': expected type code {}, found {:?} (code {})", + column, filename, expected_type, phys, actual + )); + } + + // Get the pre-computed page layout from the cache. + let pages = fmc.column_pages.get(&leaf_idx) + .cloned() + .unwrap_or_default(); + + // Open a file handle for this column reader WITH the page index loaded: the retained + // cursor's skip_records uses the page index internally for efficient page-hopping. The + // per-file metadata cache above saves the JAVA-FACING costs (schema resolution, page + // layout computation, ColumnPageIndex FFM marshal) — this reader's page-index load is + // cheap (already in OS page cache from the first parse) and necessary for decode perf. + // + // TODO(consider optimisation): this re-opens the file and re-reads the footer, which + // get_or_build_file_metadata already parsed above. Only the cheap footer read + syscall are + // duplicated (the page-index parse is cached, and the bytes are warm), but the reader could + // be constructed from the already-parsed ParquetMetaData to skip the second footer read. + let file = File::open(filename).map_err(|e| format!("Failed to open '{}': {}", filename, e))?; + let options = ReadOptionsBuilder::new().with_page_index().build(); + let reader = SerializedFileReader::new_with_options(file, options) + .map_err(|e| format!("Failed to read parquet '{}': {}", filename, e))?; + + Ok(ColumnReaderState { + reader, + leaf_idx, + physical_type: phys, + repeated: max_rep > 0, + max_def_level: max_def, + row_count: fmc.row_count, + rg_first_row: fmc.rg_first_row.clone(), + rg_num_rows: fmc.rg_num_rows.clone(), + pages, + cursor: None, + scratch: DecodeScratch::new(), + }) + } + + /// Translate a global row position into `(row_group_index, local_offset)`. + /// + /// TODO(consider optimisation): linear scan over row groups. Row groups are few (tens) and this + /// is only on the single-row read path (not the page-decode aggregation path), so it is low + /// priority, but it could be a binary search over `rg_first_row` like `page_for_row`. + fn locate(&self, row: i64) -> Result<(usize, i64), String> { + for i in 0..self.rg_first_row.len() { + let start = self.rg_first_row[i]; + let end = start + self.rg_num_rows[i]; + if row >= start && row < end { + return Ok((i, row - start)); + } + } + Err(format!("Row {} not found in any row group (row count {})", row, self.row_count)) + } + + /// Find the index of the page containing global row `row` (binary search over + /// the ascending page layout). + fn page_for_row(&self, row: i64) -> Result { + // partition_point finds the first page whose global_first_row > row; the + // page we want is the one immediately before it. + let p = self.pages.partition_point(|e| e.global_first_row <= row); + if p == 0 { + return Err(format!("Row {} precedes the first page (row count {})", row, self.row_count)); + } + let idx = p - 1; + let entry = &self.pages[idx]; + if row >= entry.global_first_row && row < entry.global_first_row + entry.num_rows { + Ok(idx) + } else { + Err(format!("Row {} not found in any page (row count {})", row, self.row_count)) + } + } +} + +/// Builds the per-page layout for a column. Prefers the Parquet OffsetIndex + +/// ColumnIndex (true page granularity); falls back to one entry per row group +/// when the file has no page index. +fn build_page_layout( + metadata: &parquet::file::metadata::ParquetMetaData, + leaf_idx: usize, + phys: PhysicalType, + rg_first_row: &[i64], + rg_num_rows: &[i64], +) -> Vec { + let n_rg = metadata.num_row_groups(); + let offset_index = metadata.offset_index(); + let column_index = metadata.column_index(); + + let mut pages: Vec = Vec::new(); + + for rg in 0..n_rg { + let oi_pages = offset_index + .and_then(|oi| oi.get(rg)) + .and_then(|cols| cols.get(leaf_idx)); + let ci = column_index + .and_then(|ci| ci.get(rg)) + .and_then(|cols| cols.get(leaf_idx)); + + match oi_pages { + Some(oi) => { + let locations = oi.page_locations(); + let rg_rows = rg_num_rows[rg]; + for (p, loc) in locations.iter().enumerate() { + let local_first = loc.first_row_index; + let next_local = if p + 1 < locations.len() { + locations[p + 1].first_row_index + } else { + rg_rows + }; + let num_rows = next_local - local_first; + let null_count = ci.and_then(|c| c.null_count(p)).unwrap_or(-1); + let (min_long, max_long) = ci + .map(|c| page_min_max(c, p, phys)) + .unwrap_or(MINMAX_UNKNOWN); + pages.push(PageEntry { + global_first_row: rg_first_row[rg] + local_first, + num_rows, + file_offset: loc.offset, + compressed_size: loc.compressed_page_size, + null_count, + min_long, + max_long, + rg_idx: rg, + local_first_row: local_first, + }); + } + } + None => { + // Fallback: treat the whole row group as a single "page". + let cc = metadata.row_group(rg).column(leaf_idx); + let null_count = cc + .statistics() + .and_then(|s| s.null_count_opt()) + .map(|n| n as i64) + .unwrap_or(-1); + let compressed = cc.compressed_size().min(i32::MAX as i64) as i32; + pages.push(PageEntry { + global_first_row: rg_first_row[rg], + num_rows: rg_num_rows[rg], + file_offset: cc.data_page_offset(), + compressed_size: compressed, + null_count, + min_long: MINMAX_UNKNOWN.0, + max_long: MINMAX_UNKNOWN.1, + rg_idx: rg, + local_first_row: 0, + }); + } + } + } + + pages +} + +/// Sentinel pair meaning "min/max unknown": the widest possible range, so a consumer making +/// skip decisions (the DocValuesSkipper) can never wrongly exclude the page. Distinguishable +/// from real data only in that real data spanning the full i64 range behaves identically — +/// which is exactly the safe behavior. +const MINMAX_UNKNOWN: (i64, i64) = (i64::MIN, i64::MAX); + +/// Extracts the per-page min/max as raw i64 bits from a typed ColumnIndex. +/// Returns [`MINMAX_UNKNOWN`] for byte-array/unsupported columns (binary min/max is not +/// exchanged as i64) and for pages whose stats are absent. +fn page_min_max(ci: &ColumnIndexMetaData, idx: usize, _phys: PhysicalType) -> (i64, i64) { + // A stat may be absent per page (stats disabled, or an all-null page). Report the unknown + // sentinel rather than 0 — 0 is indistinguishable from a real value and would let a + // skipper wrongly exclude pages. + match ci { + ColumnIndexMetaData::INT32(p) => match (p.min_value(idx), p.max_value(idx)) { + (Some(min), Some(max)) => (*min as i64, *max as i64), + _ => MINMAX_UNKNOWN, + }, + ColumnIndexMetaData::INT64(p) => match (p.min_value(idx), p.max_value(idx)) { + (Some(min), Some(max)) => (*min, *max), + _ => MINMAX_UNKNOWN, + }, + ColumnIndexMetaData::FLOAT(p) => match (p.min_value(idx), p.max_value(idx)) { + (Some(min), Some(max)) => (min.to_bits() as i64, max.to_bits() as i64), + _ => MINMAX_UNKNOWN, + }, + ColumnIndexMetaData::DOUBLE(p) => match (p.min_value(idx), p.max_value(idx)) { + (Some(min), Some(max)) => (min.to_bits() as i64, max.to_bits() as i64), + _ => MINMAX_UNKNOWN, + }, + ColumnIndexMetaData::BOOLEAN(p) => match (p.min_value(idx), p.max_value(idx)) { + (Some(min), Some(max)) => (if *min { 1 } else { 0 }, if *max { 1 } else { 0 }), + _ => MINMAX_UNKNOWN, + }, + _ => MINMAX_UNKNOWN, + } +} + +/// Maps a Parquet physical type to the Java-facing `expected_type` discriminant. +/// Returns `-1` for unsupported physical types (e.g. INT96), which can never +/// match a valid expectation and therefore surfaces as a clear mismatch error. +fn physical_type_code(t: PhysicalType) -> i32 { + match t { + PhysicalType::BOOLEAN => TYPE_BOOL, + PhysicalType::INT32 => TYPE_INT32, + PhysicalType::INT64 => TYPE_INT64, + PhysicalType::FLOAT => TYPE_FLOAT, + PhysicalType::DOUBLE => TYPE_DOUBLE, + PhysicalType::BYTE_ARRAY => TYPE_BYTE_ARRAY, + PhysicalType::FIXED_LEN_BYTE_ARRAY => TYPE_BYTE_ARRAY, + PhysicalType::INT96 => -1, + } +} + +/// Reads exactly one record (after skipping `skip` records) from a typed column +/// reader, returning that record's non-null values. For a single-valued column +/// the result holds 0 (null) or 1 value; for a repeated column it holds all the +/// values of the record. +fn read_record_values( + r: &mut ColumnReaderImpl, + skip: usize, +) -> Result, String> { + if skip > 0 { + let skipped = r.skip_records(skip).map_err(|e| e.to_string())?; + if skipped < skip { + return Err(format!("requested skip of {} records but only {} available", skip, skipped)); + } + } + let mut def_levels: Vec = Vec::new(); + let mut rep_levels: Vec = Vec::new(); + let mut values: Vec = Vec::new(); + r.read_records(1, Some(&mut def_levels), Some(&mut rep_levels), &mut values) + .map_err(|e| e.to_string())?; + Ok(values) +} + +/// Cached per-file metadata: footer + page-index parse results shared across all column readers +/// for the same file. Parquet files are immutable (changed data = new file = new path), so entries +/// never need invalidation — they can only be evicted when a file is deleted (shard close). This +/// is the `.dvm` equivalent: parsed once at first column open, then every subsequent query reuses +/// it without FFM/IO, at node lifetime scope. +struct FileMetadataCache { + /// Schema descriptor pointer (for column lookup). + schema: std::sync::Arc, + /// Number of rows in the file. + row_count: i64, + /// Per-row-group: global first row. + rg_first_row: Vec, + /// Per-row-group: number of rows. + rg_num_rows: Vec, + /// Per-column page layouts, keyed by leaf column index. Computed once per file. + column_pages: HashMap>, +} + +lazy_static! { + /// Node-level file metadata cache. Keyed by absolute file path. Entries are never + /// invalidated (immutable files) — evicted only on explicit `parquet_evict_file_metadata`. + static ref FILE_METADATA_CACHE: Mutex>> = Mutex::new(HashMap::new()); + + /// Per-handle registry of open column readers, keyed by an opaque i64 handle. + /// Mirrors the writer-side handle pattern; serialised behind a single mutex + /// since column readers are not shared across threads. + static ref COLUMN_READERS: Mutex> = Mutex::new(HashMap::new()); +} + +/// Monotonic handle allocator. Always `>= 0`, so a returned handle is never +/// confused with the `< 0` error-pointer convention. +static NEXT_COLUMN_READER_HANDLE: AtomicI64 = AtomicI64::new(0); + +/// Locks the column-reader registry, converting a poisoned mutex into a normal +/// FFM error instead of propagating the panic. +fn lock_readers<'a>() -> Result>, String> { + COLUMN_READERS + .lock() + .map_err(|_| "column reader registry mutex poisoned".to_string()) +} + +/// Evicts a file's cached metadata (footer + page index) from the node-level cache, e.g. on +/// shard close or file deletion. No-op if the file isn't cached. Returns 0 on success. +#[ffm_safe] +#[no_mangle] +pub unsafe extern "C" fn parquet_evict_file_metadata( + file_ptr: *const u8, + file_len: i64, +) -> i64 { + let filename = str_from_raw(file_ptr, file_len) + .map_err(|e| format!("parquet_evict_file_metadata: {}", e))? + .to_string(); + FILE_METADATA_CACHE + .lock() + .map_err(|_| "file metadata cache mutex poisoned".to_string())? + .remove(&filename); + Ok(RC_OK) +} + +/// Opens a per-column reader over `file` for `col`, validating that the column +/// exists and its physical type matches `expected_type` +/// (0=INT32,1=INT64,2=FLOAT,3=DOUBLE,4=BOOL,5=BYTE_ARRAY). +/// +/// Returns `>= 0` handle id on success, `< 0` negated error pointer on failure. +#[ffm_safe] +#[no_mangle] +pub unsafe extern "C" fn parquet_open_column_reader( + file_ptr: *const u8, + file_len: i64, + col_ptr: *const u8, + col_len: i64, + expected_type: i32, +) -> i64 { + let filename = str_from_raw(file_ptr, file_len) + .map_err(|e| format!("parquet_open_column_reader file: {}", e))? + .to_string(); + let column = str_from_raw(col_ptr, col_len) + .map_err(|e| format!("parquet_open_column_reader column: {}", e))? + .to_string(); + + let state = ColumnReaderState::open(&filename, &column, expected_type)?; + + let handle = NEXT_COLUMN_READER_HANDLE.fetch_add(1, Ordering::SeqCst); + lock_readers()?.insert(handle, state); + log_debug!( + "parquet_open_column_reader: file={}, column={}, handle={}", + filename, column, handle + ); + Ok(handle) +} + +/// Closes a column reader handle and releases its file handle and buffers. +/// Returns `0` on success, a `< 0` error pointer if the handle is unknown. +#[ffm_safe] +#[no_mangle] +pub unsafe extern "C" fn parquet_close_column_reader(handle: i64) -> i64 { + match lock_readers()?.remove(&handle) { + Some(_) => { + log_debug!("parquet_close_column_reader: handle={}", handle); + Ok(RC_OK) + } + None => Err(format!("parquet_close_column_reader: unknown handle {}", handle)), + } +} + +/// Debug-only symbol: returns the number of currently open column-reader +/// handles. Used by Property 7 (native handle non-leakage). Never errors; +/// recovers from a poisoned mutex rather than panicking. +#[no_mangle] +pub unsafe extern "C" fn parquet_open_column_reader_count() -> i64 { + match COLUMN_READERS.lock() { + Ok(guard) => guard.len() as i64, + Err(poisoned) => poisoned.into_inner().len() as i64, + } +} + +/// Reads the single value at `row` for a single-valued column. +/// +/// On success writes: +/// - `out_present` = 1 if the row has a value, 0 if null/absent +/// - `out_long` = the value's raw bits for primitive columns: +/// INT32 sign-extended to i64; INT64 verbatim; +/// FLOAT = `f32::to_bits` (zero-extended); +/// DOUBLE = `f64::to_bits`; +/// BOOL = 0 or 1 +/// - for BYTE_ARRAY columns: the value bytes are copied into `out_buf` and +/// `out_len` is set to the byte length (or -1 when the value is null). +/// +/// Returns a `< 0` error pointer naming the row when `row >= row_count`, when +/// the handle is unknown, or when `out_buf` is too small for a BYTE_ARRAY value +/// (in which case `out_len` is set to the required length first). +#[ffm_safe] +#[no_mangle] +pub unsafe extern "C" fn parquet_read_value_at_row( + handle: i64, + row: i64, + out_present: *mut i64, + out_long: *mut i64, + out_buf: *mut u8, + out_buf_cap: i64, + out_len: *mut i64, +) -> i64 { + let mut guard = lock_readers()?; + let state = guard + .get_mut(&handle) + .ok_or_else(|| format!("parquet_read_value_at_row: unknown handle {}", handle))?; + + if row < 0 { + return Err(format!("parquet_read_value_at_row: negative row {}", row)); + } + if row >= state.row_count { + return Err(format!( + "parquet_read_value_at_row: row {} out of range (row count {})", + row, state.row_count + )); + } + + // Default outputs: absent value. + if !out_present.is_null() { + *out_present = 0; + } + if !out_long.is_null() { + *out_long = 0; + } + if !out_len.is_null() { + *out_len = -1; + } + + let (rg_idx, local) = state.locate(row)?; + let rg = state.reader.get_row_group(rg_idx).map_err(|e| e.to_string())?; + let col = rg.get_column_reader(state.leaf_idx).map_err(|e| e.to_string())?; + let local = local as usize; + + match col { + ColumnReader::Int32ColumnReader(mut r) => { + if let Some(v) = read_record_values(&mut r, local)?.first() { + set_present(out_present, out_long, *v as i64); + } + } + ColumnReader::Int64ColumnReader(mut r) => { + if let Some(v) = read_record_values(&mut r, local)?.first() { + set_present(out_present, out_long, *v); + } + } + ColumnReader::FloatColumnReader(mut r) => { + if let Some(v) = read_record_values(&mut r, local)?.first() { + set_present(out_present, out_long, v.to_bits() as i64); + } + } + ColumnReader::DoubleColumnReader(mut r) => { + if let Some(v) = read_record_values(&mut r, local)?.first() { + set_present(out_present, out_long, v.to_bits() as i64); + } + } + ColumnReader::BoolColumnReader(mut r) => { + if let Some(v) = read_record_values(&mut r, local)?.first() { + set_present(out_present, out_long, if *v { 1 } else { 0 }); + } + } + ColumnReader::ByteArrayColumnReader(mut r) => { + if let Some(v) = read_record_values(&mut r, local)?.first() { + return write_bytes_value(v.data(), out_present, out_buf, out_buf_cap, out_len); + } + } + ColumnReader::FixedLenByteArrayColumnReader(mut r) => { + if let Some(v) = read_record_values(&mut r, local)?.first() { + return write_bytes_value(v.data(), out_present, out_buf, out_buf_cap, out_len); + } + } + ColumnReader::Int96ColumnReader(_) => { + return Err("parquet_read_value_at_row: INT96 columns are not supported".to_string()); + } + } + + Ok(RC_OK) +} + +/// Marks a primitive value present and stores its raw bits. +unsafe fn set_present(out_present: *mut i64, out_long: *mut i64, bits: i64) { + if !out_present.is_null() { + *out_present = 1; + } + if !out_long.is_null() { + *out_long = bits; + } +} + +/// Copies a single BYTE_ARRAY value into the caller buffer. Sets `out_present=1` +/// and `out_len` to the byte length. Returns `RC_OVERFLOW` (after recording the +/// required length in `out_len`) when the value does not fit in `out_buf_cap`, +/// so the caller can grow its buffer and retry once. +unsafe fn write_bytes_value( + bytes: &[u8], + out_present: *mut i64, + out_buf: *mut u8, + out_buf_cap: i64, + out_len: *mut i64, +) -> Result { + if !out_present.is_null() { + *out_present = 1; + } + let n = bytes.len(); + if !out_len.is_null() { + *out_len = n as i64; + } + if (n as i64) > out_buf_cap || (n > 0 && out_buf.is_null()) { + return Ok(RC_OVERFLOW); + } + if n > 0 { + std::ptr::copy_nonoverlapping(bytes.as_ptr(), out_buf, n); + } + Ok(RC_OK) +} + +/// Reads all values at `row` for a repeated (multi-valued) column. +/// +/// Capacity contract: `out_long_cap` is the maximum element count for *both* +/// primitive and BYTE_ARRAY columns; `out_byte_offsets` (BYTE_ARRAY only) must +/// have capacity `out_long_cap + 1`. +/// +/// On success (`RC_OK`): +/// - `out_count` = number of values at the row +/// - primitive columns: raw bits (see `parquet_read_value_at_row`) written to +/// `out_longs` +/// - BYTE_ARRAY columns: concatenated bytes in `out_byte_buf`, CSR offsets +/// (length `count + 1`) in `out_byte_offsets` +/// +/// On `RC_OVERFLOW`: `out_count` holds the required element count. When the +/// element count fits but only the byte buffer is too small, the full CSR +/// offsets are still written so `out_byte_offsets[count]` reports the required +/// total byte size, enabling a single retry. +#[ffm_safe] +#[no_mangle] +pub unsafe extern "C" fn parquet_read_repeated_at_row( + handle: i64, + row: i64, + out_count: *mut i64, + out_longs: *mut i64, + out_long_cap: i64, + out_byte_buf: *mut u8, + out_byte_offsets: *mut i64, + out_byte_buf_cap: i64, +) -> i64 { + let mut guard = lock_readers()?; + let state = guard + .get_mut(&handle) + .ok_or_else(|| format!("parquet_read_repeated_at_row: unknown handle {}", handle))?; + + if row < 0 { + return Err(format!("parquet_read_repeated_at_row: negative row {}", row)); + } + if row >= state.row_count { + return Err(format!( + "parquet_read_repeated_at_row: row {} out of range (row count {})", + row, state.row_count + )); + } + + if !out_count.is_null() { + *out_count = 0; + } + + let (rg_idx, local) = state.locate(row)?; + let rg = state.reader.get_row_group(rg_idx).map_err(|e| e.to_string())?; + let col = rg.get_column_reader(state.leaf_idx).map_err(|e| e.to_string())?; + let local = local as usize; + + match col { + ColumnReader::Int32ColumnReader(mut r) => { + let vals = read_record_values(&mut r, local)?; + write_primitive_repeated(vals.iter().map(|v| *v as i64), vals.len(), out_count, out_longs, out_long_cap) + } + ColumnReader::Int64ColumnReader(mut r) => { + let vals = read_record_values(&mut r, local)?; + write_primitive_repeated(vals.iter().copied(), vals.len(), out_count, out_longs, out_long_cap) + } + ColumnReader::FloatColumnReader(mut r) => { + let vals = read_record_values(&mut r, local)?; + write_primitive_repeated(vals.iter().map(|v| v.to_bits() as i64), vals.len(), out_count, out_longs, out_long_cap) + } + ColumnReader::DoubleColumnReader(mut r) => { + let vals = read_record_values(&mut r, local)?; + write_primitive_repeated(vals.iter().map(|v| v.to_bits() as i64), vals.len(), out_count, out_longs, out_long_cap) + } + ColumnReader::BoolColumnReader(mut r) => { + let vals = read_record_values(&mut r, local)?; + write_primitive_repeated(vals.iter().map(|v| if *v { 1i64 } else { 0i64 }), vals.len(), out_count, out_longs, out_long_cap) + } + ColumnReader::ByteArrayColumnReader(mut r) => { + let vals = read_record_values(&mut r, local)?; + let slices: Vec<&[u8]> = vals.iter().map(|v| v.data()).collect(); + write_bytes_repeated(&slices, out_count, out_long_cap, out_byte_buf, out_byte_offsets, out_byte_buf_cap) + } + ColumnReader::FixedLenByteArrayColumnReader(mut r) => { + let vals = read_record_values(&mut r, local)?; + let slices: Vec<&[u8]> = vals.iter().map(|v| v.data()).collect(); + write_bytes_repeated(&slices, out_count, out_long_cap, out_byte_buf, out_byte_offsets, out_byte_buf_cap) + } + ColumnReader::Int96ColumnReader(_) => { + Err("parquet_read_repeated_at_row: INT96 columns are not supported".to_string()) + } + } +} + +/// Writes repeated primitive values to `out_longs`, or reports overflow. +unsafe fn write_primitive_repeated( + values: impl Iterator, + count: usize, + out_count: *mut i64, + out_longs: *mut i64, + out_long_cap: i64, +) -> Result { + if !out_count.is_null() { + *out_count = count as i64; + } + if (count as i64) > out_long_cap || out_longs.is_null() { + return Ok(RC_OVERFLOW); + } + for (i, v) in values.enumerate() { + *out_longs.add(i) = v; + } + Ok(RC_OK) +} + +/// Writes repeated BYTE_ARRAY values (CSR layout) to the caller buffers, or +/// reports overflow with required sizes. +unsafe fn write_bytes_repeated( + slices: &[&[u8]], + out_count: *mut i64, + out_long_cap: i64, + out_byte_buf: *mut u8, + out_byte_offsets: *mut i64, + out_byte_buf_cap: i64, +) -> Result { + let count = slices.len(); + let total_bytes: usize = slices.iter().map(|s| s.len()).sum(); + if !out_count.is_null() { + *out_count = count as i64; + } + + // Element-count overflow: cannot safely write offsets (capacity is count+1). + if (count as i64) > out_long_cap { + return Ok(RC_OVERFLOW); + } + + // Element count fits: write the full CSR offsets so that, even on a byte + // overflow, out_byte_offsets[count] == total_bytes reports the required size. + if !out_byte_offsets.is_null() { + let mut acc = 0i64; + for (i, s) in slices.iter().enumerate() { + *out_byte_offsets.add(i) = acc; + acc += s.len() as i64; + } + *out_byte_offsets.add(count) = acc; + } + + if (total_bytes as i64) > out_byte_buf_cap || (total_bytes > 0 && out_byte_buf.is_null()) { + return Ok(RC_OVERFLOW); + } + + let mut acc = 0usize; + for s in slices { + if !s.is_empty() { + std::ptr::copy_nonoverlapping(s.as_ptr(), out_byte_buf.add(acc), s.len()); + } + acc += s.len(); + } + Ok(RC_OK) +} + +// --------------------------------------------------------------------------- +// Page-index loader + page decoder +// --------------------------------------------------------------------------- +// +// Page-oriented reads used by the Java `ParquetColumnReader` to decode a whole +// page per call rather than a row per call: +// - `parquet_get_column_num_pages` — page count, so Java can pre-size buffers +// - `parquet_get_column_page_index` — per-page row-range jump table + page stats +// - `parquet_decode_page_at_row` — decoded values + presence bitset for a page +// +// All row indices exchanged here are global (file-relative). + +/// Returns the number of pages in the column (`>= 0`), or a `< 0` error pointer +/// for an unknown handle. Java reads this first to size the parallel arrays +/// passed to `parquet_get_column_page_index`. +#[ffm_safe] +#[no_mangle] +pub unsafe extern "C" fn parquet_get_column_num_pages(handle: i64) -> i64 { + let guard = lock_readers()?; + let state = guard + .get(&handle) + .ok_or_else(|| format!("parquet_get_column_num_pages: unknown handle {}", handle))?; + Ok(state.pages.len() as i64) +} + +/// Writes the column's per-page row-range jump table and page statistics into +/// caller-provided parallel arrays, each of capacity `out_buf_capacity` +/// (= the page count from `parquet_get_column_num_pages`). +/// +/// Arrays (length = page count): +/// - `out_first_row` global index of the page's first row +/// - `out_file_offset` byte offset of the page in the file (0 if unknown) +/// - `out_compressed_size` compressed page size in bytes (0 if unknown) +/// - `out_null_count` nulls in the page, or -1 when unknown +/// - `out_min_long` per-page min raw bits (numeric only; 0 otherwise) +/// - `out_max_long` per-page max raw bits (numeric only; 0 otherwise) +/// +/// `out_actual_pages` always receives the true page count. Returns `RC_OVERFLOW` +/// (a positive sentinel) without writing the arrays when `out_buf_capacity` is +/// smaller than the page count, so the caller can grow and retry. Returns a +/// `< 0` error pointer for an unknown handle. +#[ffm_safe] +#[no_mangle] +pub unsafe extern "C" fn parquet_get_column_page_index( + handle: i64, + out_first_row: *mut i64, + out_file_offset: *mut i64, + out_compressed_size: *mut i32, + out_null_count: *mut i64, + out_min_long: *mut i64, + out_max_long: *mut i64, + out_buf_capacity: i64, + out_actual_pages: *mut i64, +) -> i64 { + let guard = lock_readers()?; + let state = guard + .get(&handle) + .ok_or_else(|| format!("parquet_get_column_page_index: unknown handle {}", handle))?; + + let n = state.pages.len(); + if !out_actual_pages.is_null() { + *out_actual_pages = n as i64; + } + if (n as i64) > out_buf_capacity { + return Ok(RC_OVERFLOW); + } + + for (i, e) in state.pages.iter().enumerate() { + if !out_first_row.is_null() { + *out_first_row.add(i) = e.global_first_row; + } + if !out_file_offset.is_null() { + *out_file_offset.add(i) = e.file_offset; + } + if !out_compressed_size.is_null() { + *out_compressed_size.add(i) = e.compressed_size; + } + if !out_null_count.is_null() { + *out_null_count.add(i) = e.null_count; + } + if !out_min_long.is_null() { + *out_min_long.add(i) = e.min_long; + } + if !out_max_long.is_null() { + *out_max_long.add(i) = e.max_long; + } + } + Ok(RC_OK) +} + +/// Decodes one primitive page through a (possibly retained) typed column reader and writes +/// values + packed presence straight into the caller's out-buffers. `skip` is relative to the +/// reader's current position, NOT the row group start — the retained-cursor caller computes it +/// from the cursor position so a reused reader only skips forward the remaining distance. +#[allow(clippy::too_many_arguments)] +unsafe fn decode_primitive_page( + r: &mut ColumnReaderImpl, + skip: usize, + num_rows: usize, + max_def_level: i16, + effective_null_count: i64, + def_scratch: &mut Vec, + val_scratch: &mut Vec, + to_bits: impl Fn(T::T) -> i64, + out_value_buf: *mut u8, + out_presence_bitset: *mut i64, +) -> Result<(), String> +where + T::T: Copy, +{ + decode_page_records(r, skip, num_rows, def_scratch, val_scratch)?; + pack_presence_from_def_levels(def_scratch, max_def_level, num_rows, out_presence_bitset); + expand_to_outbuf(val_scratch, to_bits, effective_null_count, num_rows, out_value_buf, out_presence_bitset as *const i64); + Ok(()) +} + +/// Decodes one page's worth of single-valued records, returning a per-row +/// presence flag (`true` = value present) and the dense list of non-null values +/// in row order. `skip` records are skipped first, then `num_rows` records are +/// read. Works for required columns (`max_def_level == 0`, all present) and +/// optional non-repeated columns. +fn decode_page_records( + r: &mut ColumnReaderImpl, + skip: usize, + num_rows: usize, + scratch_def: &mut Vec, + scratch_vals: &mut Vec, +) -> Result<(), String> { + if skip > 0 { + let skipped = r.skip_records(skip).map_err(|e| e.to_string())?; + if skipped < skip { + return Err(format!( + "page decode: requested skip of {} records but only {} available", + skip, skipped + )); + } + } + + scratch_def.clear(); + scratch_vals.clear(); + scratch_def.reserve(num_rows.saturating_sub(scratch_def.capacity())); + scratch_vals.reserve(num_rows.saturating_sub(scratch_vals.capacity())); + + let (records_read, _values_read, _levels_read) = r + .read_records(num_rows, Some(scratch_def), None, scratch_vals) + .map_err(|e| e.to_string())?; + if records_read < num_rows { + return Err(format!( + "page decode: expected {} records but read {}", + num_rows, records_read + )); + } + Ok(()) +} + +/// Packs definition levels directly into a little-endian `long[]` bitset in the +/// caller's out-buffer, with bit `i` set when `def_levels[i] == max_def_level`. +/// When `max_def_level == 0` (required column), all bits are set. +/// +/// Writes `ceil(num_rows / 64)` words; the caller must ensure `out` has capacity +/// for that many. Uses a branchless comparison so the inner loop auto-vectorizes. +/// +/// The primitive path uses this because it has def_levels in hand and no per-row +/// bool vector; the byte path instead uses `write_presence_bitset`, packing from +/// the `Vec` it already builds to slice its values. +#[inline] +unsafe fn pack_presence_from_def_levels( + def_levels: &[i16], + max_def_level: i16, + num_rows: usize, + out: *mut i64, +) { + // The optional-column branch reads def_levels[0..num_rows] unchecked; assert the precondition + // so a corrupt page index or decoder state fails loudly in tests rather than reading OOB. + debug_assert!(def_levels.len() >= num_rows || max_def_level == 0); + let words = (num_rows + 63) / 64; + if max_def_level == 0 { + // Required column: every row is present → all-ones, mask tail. + for w in 0..words { + let remaining = num_rows - w * 64; + if remaining >= 64 { + *out.add(w) = -1i64; // all bits set + } else { + *out.add(w) = ((1u64 << remaining) - 1) as i64; + } + } + } else { + // Optional column: branchless pack. The inner loop auto-vectorizes because + // `(d == max_def_level) as u64` is a conditional-move / compare instruction. + for w in 0..words { + let mut bits: u64 = 0; + let base = w * 64; + let end = (base + 64).min(num_rows); + for b in base..end { + bits |= ((*def_levels.get_unchecked(b) == max_def_level) as u64) << (b - base); + } + *out.add(w) = bits as i64; + } + } +} + +/// Writes dense non-null values into the caller's out-buffer as per-row raw i64 +/// bits, using the packed presence bitset to scatter. Null slots hold 0. +/// +/// The inner scatter loop is split by nullability: when `null_count == 0` the +/// entire dense buffer can be converted with a tight widening loop that +/// auto-vectorizes (LLVM emits vpmovsxdq / sshll). The nullable path reads the +/// packed bits we just wrote and scatters accordingly. +/// +/// # Safety +/// `out` must be valid for `num_rows * 8` bytes. `presence_bits` must contain +/// the packed bitset already written by `pack_presence_from_def_levels`. +#[inline] +unsafe fn expand_to_outbuf( + dense: &[T], + to_bits: impl Fn(T) -> i64, + null_count: i64, + num_rows: usize, + out: *mut u8, + presence_bits: *const i64, +) { + let out_i64 = out as *mut i64; + if null_count == 0 { + // All rows present — tight conversion loop, no branching, SIMD-friendly. + for i in 0..num_rows { + *out_i64.add(i) = to_bits(*dense.get_unchecked(i)); + } + } else { + // Scatter using the packed presence bits. Read one word at a time and use + // trailing_zeros to jump to set bits (pop-and-scatter pattern). + // First zero-fill so null slots hold 0 without an explicit branch. + std::ptr::write_bytes(out, 0, num_rows * 8); + let mut di = 0usize; + let words = (num_rows + 63) / 64; + for w in 0..words { + let mut bits = *presence_bits.add(w) as u64; + while bits != 0 { + let b = bits.trailing_zeros() as usize; + let row = w * 64 + b; + *out_i64.add(row) = to_bits(*dense.get_unchecked(di)); + di += 1; + bits &= bits - 1; // clear lowest set bit + } + } + } +} + +/// Packs a per-row presence slice into a little-endian `long[]` bitset (bit i set +/// when row i is present), writing into `out` (capacity `out_words`). Returns the +/// number of words required; on capacity shortfall writes nothing and the caller +/// treats the positive required count as an overflow signal. +/// +/// The byte path uses this because it already builds a `Vec` to slice its +/// values; the primitive path instead uses `pack_presence_from_def_levels`, which +/// packs straight from def_levels without materializing a bool vector. +unsafe fn write_presence_bitset(presence: &[bool], out: *mut i64, out_words: i64) -> i64 { + let words_needed = ((presence.len() + 63) / 64) as i64; + if words_needed > out_words || out.is_null() { + return words_needed; + } + for w in 0..words_needed as usize { + let mut bits: u64 = 0; + let base = w * 64; + for b in 0..64 { + let idx = base + b; + if idx >= presence.len() { + break; + } + if presence[idx] { + bits |= 1u64 << b; + } + } + *out.add(w) = bits as i64; + } + words_needed +} + +/// Decode the page containing global row `row` into caller buffers (values + +/// presence bitset). +/// +/// On success (`RC_OK`): +/// - `out_first_row` / `out_last_row` = inclusive global row range of the page +/// - primitive columns: per-row raw bits written to `out_value_buf` +/// (interpreted as `long[]`, one slot per row; null rows hold 0); +/// `out_value_actual_len` = `rows * 8` +/// - BYTE_ARRAY columns: concatenated value bytes in `out_value_buf`, per-row +/// CSR offsets (length `rows + 1`) in `out_byte_offsets`; +/// `out_value_actual_len` = total bytes used +/// - `out_presence_bitset` = packed `long[]`, one bit per row +/// +/// On `RC_OVERFLOW` (a positive sentinel): `out_first_row`, `out_last_row` and +/// `out_value_actual_len` are populated so the caller can size every buffer +/// (values = `out_value_actual_len` bytes; offsets = `rows + 1`; presence = +/// `ceil(rows / 64)` words) and retry once. Returns a `< 0` error pointer for an +/// unknown handle, an out-of-range row, or a repeated (multi-valued) column. +#[ffm_safe] +#[no_mangle] +pub unsafe extern "C" fn parquet_decode_page_at_row( + handle: i64, + row: i64, + out_first_row: *mut i64, + out_last_row: *mut i64, + out_value_buf: *mut u8, + out_value_buf_cap: i64, + out_value_actual_len: *mut i64, + out_byte_offsets: *mut i32, + out_byte_offsets_cap: i64, + out_presence_bitset: *mut i64, + out_presence_bits_cap: i64, +) -> i64 { + let mut guard = lock_readers()?; + let state = guard + .get_mut(&handle) + .ok_or_else(|| format!("parquet_decode_page_at_row: unknown handle {}", handle))?; + + if row < 0 || row >= state.row_count { + return Err(format!( + "parquet_decode_page_at_row: row {} out of range (row count {})", + row, state.row_count + )); + } + if state.repeated { + return Err(format!( + "parquet_decode_page_at_row: column is repeated (multi-valued); use parquet_read_repeated_at_row (handle {})", + handle + )); + } + + let page_idx = state.page_for_row(row)?; + // Copy out the page's coordinates before borrowing the reader mutably. + let (rg_idx, local_first, num_rows, first_global) = { + let e = &state.pages[page_idx]; + (e.rg_idx, e.local_first_row, e.num_rows, e.global_first_row) + }; + let num_rows_usize = num_rows as usize; + let max_def_level = state.max_def_level; + let physical_type = state.physical_type; + + // Always report the page row range so the caller can bound its cache and + // size buffers even on the overflow path. + if !out_first_row.is_null() { + *out_first_row = first_global; + } + if !out_last_row.is_null() { + *out_last_row = first_global + num_rows - 1; + } + + // For primitive types the value byte length is known before decode (num_rows * 8), so check + // capacity up front and write straight into the out-buffers. BYTE_ARRAY defers its capacity + // check to after decode, since its total byte length is data-dependent. + let is_primitive = physical_type != PhysicalType::BYTE_ARRAY + && physical_type != PhysicalType::FIXED_LEN_BYTE_ARRAY; + + if is_primitive { + let value_bytes = (num_rows_usize * 8) as i64; + if !out_value_actual_len.is_null() { + *out_value_actual_len = value_bytes; + } + let presence_words = ((num_rows_usize + 63) / 64) as i64; + if value_bytes > out_value_buf_cap + || out_value_buf.is_null() + || presence_words > out_presence_bits_cap + || out_presence_bitset.is_null() + { + return Ok(RC_OVERFLOW); + } + } + + // Get the null_count from the page entry for the expand path (0 means all-present). + let page_null_count = state.pages[page_idx].null_count; + // If null_count is unknown (-1), treat as potentially nullable. + let effective_null_count = if page_null_count < 0 { 1 } else { page_null_count }; + + // Retained-cursor reader acquisition. A typed column reader can only move forward, so it + // is reusable exactly when the target page is in the same row group at-or-ahead of the + // cursor's position; then the skip is the remaining forward distance and — crucially — + // the dictionary page and row-group metadata are NOT re-read. Any backward jump or + // row-group change falls back to a fresh reader (dictionary re-decoded once). + // The cursor is take()n up front so a decode error leaves it invalidated; each arm + // re-installs it only after a successful decode. + let (col, skip) = match state.cursor.take() { + Some(c) if c.rg_idx == rg_idx && c.position <= first_global => { + (c.col_reader, (first_global - c.position) as usize) + } + _ => { + let rg = state.reader.get_row_group(rg_idx).map_err(|e| e.to_string())?; + let col = rg.get_column_reader(state.leaf_idx).map_err(|e| e.to_string())?; + (col, local_first as usize) + } + }; + // Position of the reader after this page is consumed; stored on the re-installed cursor. + let next_position = first_global + num_rows; + + // Decode and write directly to the out-buffers. Primitive types decode with reused scratch + // buffers (no per-call allocation), pack presence branchlessly from def_levels, and scatter + // values straight to the out-buffer. BYTE_ARRAY/FIXED_LEN_BYTE_ARRAY decode into a temporary + // value vector first, since their total byte length must be computed before the overflow check. + match col { + ColumnReader::Int32ColumnReader(mut r) => { + let scratch = &mut state.scratch; + decode_primitive_page( + &mut r, skip, num_rows_usize, max_def_level, effective_null_count, + &mut scratch.def_levels, &mut scratch.values_i32, |v| v as i64, + out_value_buf, out_presence_bitset, + )?; + state.cursor = Some(CursorState { rg_idx, col_reader: ColumnReader::Int32ColumnReader(r), position: next_position }); + return Ok(RC_OK); + } + ColumnReader::Int64ColumnReader(mut r) => { + let scratch = &mut state.scratch; + decode_primitive_page( + &mut r, skip, num_rows_usize, max_def_level, effective_null_count, + &mut scratch.def_levels, &mut scratch.values_i64, |v| v, + out_value_buf, out_presence_bitset, + )?; + state.cursor = Some(CursorState { rg_idx, col_reader: ColumnReader::Int64ColumnReader(r), position: next_position }); + return Ok(RC_OK); + } + ColumnReader::FloatColumnReader(mut r) => { + let scratch = &mut state.scratch; + decode_primitive_page( + &mut r, skip, num_rows_usize, max_def_level, effective_null_count, + &mut scratch.def_levels, &mut scratch.values_f32, |v| v.to_bits() as i64, + out_value_buf, out_presence_bitset, + )?; + state.cursor = Some(CursorState { rg_idx, col_reader: ColumnReader::FloatColumnReader(r), position: next_position }); + return Ok(RC_OK); + } + ColumnReader::DoubleColumnReader(mut r) => { + let scratch = &mut state.scratch; + decode_primitive_page( + &mut r, skip, num_rows_usize, max_def_level, effective_null_count, + &mut scratch.def_levels, &mut scratch.values_f64, |v| v.to_bits() as i64, + out_value_buf, out_presence_bitset, + )?; + state.cursor = Some(CursorState { rg_idx, col_reader: ColumnReader::DoubleColumnReader(r), position: next_position }); + return Ok(RC_OK); + } + ColumnReader::BoolColumnReader(mut r) => { + let scratch = &mut state.scratch; + decode_primitive_page( + &mut r, skip, num_rows_usize, max_def_level, effective_null_count, + &mut scratch.def_levels, &mut scratch.values_bool, |v| if v { 1i64 } else { 0i64 }, + out_value_buf, out_presence_bitset, + )?; + state.cursor = Some(CursorState { rg_idx, col_reader: ColumnReader::BoolColumnReader(r), position: next_position }); + return Ok(RC_OK); + } + ColumnReader::ByteArrayColumnReader(mut r) => { + let rc = decode_byte_page( + &mut r, skip, num_rows_usize, max_def_level, &mut state.scratch.def_levels, + out_value_buf, out_value_buf_cap, out_value_actual_len, + out_byte_offsets, out_byte_offsets_cap, out_presence_bitset, out_presence_bits_cap, + )?; + // Install the cursor only on a successful decode, matching the primitive path. On + // RC_OVERFLOW the caller retries the same row, which reopens a fresh reader anyway. + if rc == RC_OK { + state.cursor = Some(CursorState { rg_idx, col_reader: ColumnReader::ByteArrayColumnReader(r), position: next_position }); + } + return Ok(rc); + } + ColumnReader::FixedLenByteArrayColumnReader(mut r) => { + let rc = decode_byte_page( + &mut r, skip, num_rows_usize, max_def_level, &mut state.scratch.def_levels, + out_value_buf, out_value_buf_cap, out_value_actual_len, + out_byte_offsets, out_byte_offsets_cap, out_presence_bitset, out_presence_bits_cap, + )?; + // See the ByteArray arm: install the cursor only on a successful decode. + if rc == RC_OK { + state.cursor = Some(CursorState { rg_idx, col_reader: ColumnReader::FixedLenByteArrayColumnReader(r), position: next_position }); + } + return Ok(rc); + } + ColumnReader::Int96ColumnReader(_) => { + let _ = physical_type; // silence unused in this arm + Err("parquet_decode_page_at_row: INT96 columns are not supported".to_string()) + } + } +} + +/// Expands dense non-null byte values into a per-row slice vector, null rows mapping to an empty +/// slice. Works for both `ByteArray` and `FixedLenByteArray` via the `AsBytes` trait they share. +fn expand_byte_values<'a, T: parquet::data_type::AsBytes>(presence: &[bool], dense: &'a [T]) -> Vec<&'a [u8]> { + let mut out: Vec<&[u8]> = Vec::with_capacity(presence.len()); + let mut di = 0usize; + for &present in presence { + if present { + out.push(dense[di].as_bytes()); + di += 1; + } else { + out.push(&[]); + } + } + out +} + +/// Decodes one BYTE_ARRAY / FIXED_LEN_BYTE_ARRAY page and writes it to the caller buffers +/// (concatenated value bytes + per-row CSR offsets + presence bitset), or returns `RC_OVERFLOW` +/// after recording the required sizes. Shared by both byte column types via the `AsBytes` trait. +/// Unlike the primitive path, the values decode into a temporary vector first because the total +/// byte length is data-dependent and must be known before the capacity check. +#[allow(clippy::too_many_arguments)] +unsafe fn decode_byte_page( + r: &mut ColumnReaderImpl, + skip: usize, + num_rows: usize, + max_def_level: i16, + def_scratch: &mut Vec, + out_value_buf: *mut u8, + out_value_buf_cap: i64, + out_value_actual_len: *mut i64, + out_byte_offsets: *mut i32, + out_byte_offsets_cap: i64, + out_presence_bitset: *mut i64, + out_presence_bits_cap: i64, +) -> Result +where + T::T: parquet::data_type::AsBytes, +{ + let mut values: Vec = Vec::new(); + decode_page_records(r, skip, num_rows, def_scratch, &mut values)?; + let presence: Vec = if max_def_level == 0 { + vec![true; num_rows] + } else { + def_scratch.iter().take(num_rows).map(|d| *d == max_def_level).collect() + }; + let slices = expand_byte_values(&presence, &values); + write_bytes_page( + &slices, &presence, out_value_buf, out_value_buf_cap, out_value_actual_len, + out_byte_offsets, out_byte_offsets_cap, out_presence_bitset, out_presence_bits_cap, + ) +} + +/// Writes a decoded BYTE_ARRAY page (concatenated bytes + CSR offsets + presence +/// bitset) to the caller buffers, or returns `RC_OVERFLOW` after recording the +/// required total byte length. +unsafe fn write_bytes_page( + slices: &[&[u8]], + presence: &[bool], + out_value_buf: *mut u8, + out_value_buf_cap: i64, + out_value_actual_len: *mut i64, + out_byte_offsets: *mut i32, + out_byte_offsets_cap: i64, + out_presence_bitset: *mut i64, + out_presence_bits_cap: i64, +) -> Result { + let total_bytes: usize = slices.iter().map(|s| s.len()).sum(); + if !out_value_actual_len.is_null() { + *out_value_actual_len = total_bytes as i64; + } + + let offsets_needed = (slices.len() + 1) as i64; + let presence_words = ((presence.len() + 63) / 64) as i64; + if (total_bytes as i64) > out_value_buf_cap + || (total_bytes > 0 && out_value_buf.is_null()) + || offsets_needed > out_byte_offsets_cap + || out_byte_offsets.is_null() + || presence_words > out_presence_bits_cap + || out_presence_bitset.is_null() + { + return Ok(RC_OVERFLOW); + } + + // Offsets are exchanged as i32 (a page is small by design, ~20k rows). Accumulate in i64 and + // reject a page whose byte total would exceed i32::MAX rather than let the offset wrap, which + // would corrupt the offsets and drive an out-of-bounds write into out_value_buf. + let mut acc: i64 = 0; + for (i, s) in slices.iter().enumerate() { + if acc > i32::MAX as i64 { + return Err(format!("byte page offset {} exceeds i32::MAX", acc)); + } + *out_byte_offsets.add(i) = acc as i32; + if !s.is_empty() { + std::ptr::copy_nonoverlapping(s.as_ptr(), out_value_buf.add(acc as usize), s.len()); + } + acc += s.len() as i64; + } + if acc > i32::MAX as i64 { + return Err(format!("byte page total {} exceeds i32::MAX", acc)); + } + *out_byte_offsets.add(slices.len()) = acc as i32; + + write_presence_bitset(presence, out_presence_bitset, out_presence_bits_cap); + Ok(RC_OK) +} diff --git a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/codec/FieldTypeMappingTests.java b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/codec/FieldTypeMappingTests.java new file mode 100644 index 0000000000000..2792f34391c59 --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/codec/FieldTypeMappingTests.java @@ -0,0 +1,82 @@ +/* + * 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.codec; + +import org.apache.lucene.index.DocValuesType; +import org.opensearch.test.OpenSearchTestCase; + +/** + * Unit tests for {@link FieldTypeMapping} (task 4.3): the OpenSearch mapping type → Lucene DV + * type + Parquet physical type table, and the compatibility {@code validate} guard. + */ +public class FieldTypeMappingTests extends OpenSearchTestCase { + + public void testNumericMappings() { + assertMapping("boolean", DocValuesType.NUMERIC, DocValuesType.SORTED_NUMERIC, ParquetPhysicalType.BOOL); + assertMapping("byte", DocValuesType.NUMERIC, DocValuesType.SORTED_NUMERIC, ParquetPhysicalType.INT32); + assertMapping("short", DocValuesType.NUMERIC, DocValuesType.SORTED_NUMERIC, ParquetPhysicalType.INT32); + assertMapping("integer", DocValuesType.NUMERIC, DocValuesType.SORTED_NUMERIC, ParquetPhysicalType.INT32); + assertMapping("long", DocValuesType.NUMERIC, DocValuesType.SORTED_NUMERIC, ParquetPhysicalType.INT64); + assertMapping("float", DocValuesType.NUMERIC, DocValuesType.SORTED_NUMERIC, ParquetPhysicalType.FLOAT); + assertMapping("double", DocValuesType.NUMERIC, DocValuesType.SORTED_NUMERIC, ParquetPhysicalType.DOUBLE); + assertMapping("date", DocValuesType.NUMERIC, DocValuesType.SORTED_NUMERIC, ParquetPhysicalType.INT64); + assertMapping("date_nanos", DocValuesType.NUMERIC, DocValuesType.SORTED_NUMERIC, ParquetPhysicalType.INT64); + } + + public void testKeywordAndIpMappings() { + assertMapping("keyword", DocValuesType.SORTED, DocValuesType.SORTED_SET, ParquetPhysicalType.BYTE_ARRAY); + assertMapping("ip", DocValuesType.SORTED, DocValuesType.SORTED_SET, ParquetPhysicalType.BYTE_ARRAY); + } + + public void testTextAndBinaryMappings() { + assertMapping("text", DocValuesType.BINARY, DocValuesType.NONE, ParquetPhysicalType.BYTE_ARRAY); + assertMapping("binary", DocValuesType.BINARY, DocValuesType.NONE, ParquetPhysicalType.BYTE_ARRAY); + } + + public void testUnsupportedTypeThrows() { + assertFalse(FieldTypeMapping.isSupported("geo_point")); + IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> FieldTypeMapping.forType("geo_point")); + assertTrue(e.getMessage().contains("geo_point")); + } + + public void testValidateAcceptsSingleAndMultiValuedForms() { + // long supports NUMERIC (single) and SORTED_NUMERIC (multi). + FieldTypeMapping.validate("age", "long", DocValuesType.NUMERIC); + FieldTypeMapping.validate("ages", "long", DocValuesType.SORTED_NUMERIC); + // keyword supports SORTED (single) and SORTED_SET (multi). + FieldTypeMapping.validate("tag", "keyword", DocValuesType.SORTED); + FieldTypeMapping.validate("tags", "keyword", DocValuesType.SORTED_SET); + } + + public void testValidateRejectsIncompatibleDvType() { + // Requesting NUMERIC for a keyword field must fail and name the field + type. + IllegalArgumentException e = expectThrows( + IllegalArgumentException.class, + () -> FieldTypeMapping.validate("tag", "keyword", DocValuesType.NUMERIC) + ); + assertTrue(e.getMessage().contains("tag")); + assertTrue(e.getMessage().contains("keyword")); + } + + public void testValidateRejectsUnsupportedMappingType() { + IllegalArgumentException e = expectThrows( + IllegalArgumentException.class, + () -> FieldTypeMapping.validate("loc", "geo_point", DocValuesType.BINARY) + ); + assertTrue(e.getMessage().contains("loc")); + assertTrue(e.getMessage().contains("geo_point")); + } + + private static void assertMapping(String type, DocValuesType single, DocValuesType multi, ParquetPhysicalType physical) { + FieldTypeMapping.Mapping m = FieldTypeMapping.forType(type); + assertEquals("single-valued DV for " + type, single, m.singleValued()); + assertEquals("multi-valued DV for " + type, multi, m.multiValued()); + assertEquals("physical type for " + type, physical, m.physical()); + } +} diff --git a/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/codec/ParquetPhysicalTypeTests.java b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/codec/ParquetPhysicalTypeTests.java new file mode 100644 index 0000000000000..0c98011973f55 --- /dev/null +++ b/sandbox/plugins/parquet-data-format/src/test/java/org/opensearch/parquet/codec/ParquetPhysicalTypeTests.java @@ -0,0 +1,51 @@ +/* + * 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.codec; + +import org.opensearch.test.OpenSearchTestCase; + +/** + * Unit tests for {@link ParquetPhysicalType}. The {@code code()} discriminants are exchanged with + * the Rust FFM bridge ({@code parquet_open_column_reader}'s {@code expected_type}) and MUST stay in + * lock-step with the {@code TYPE_*} constants in {@code ffm.rs}; these tests pin the exact integer + * values so a future reorder of the enum cannot silently desync Java from the native side. + */ +public class ParquetPhysicalTypeTests extends OpenSearchTestCase { + + public void testCodeDiscriminantsMatchNativeContract() { + // Must match the TYPE_* constants in ffm.rs. Do not change without updating the native side. + assertEquals(0, ParquetPhysicalType.INT32.code()); + assertEquals(1, ParquetPhysicalType.INT64.code()); + assertEquals(2, ParquetPhysicalType.FLOAT.code()); + assertEquals(3, ParquetPhysicalType.DOUBLE.code()); + assertEquals(4, ParquetPhysicalType.BOOL.code()); + assertEquals(5, ParquetPhysicalType.BYTE_ARRAY.code()); + } + + public void testCodesAreUnique() { + ParquetPhysicalType[] values = ParquetPhysicalType.values(); + boolean[] seen = new boolean[values.length]; + for (ParquetPhysicalType t : values) { + int c = t.code(); + assertTrue("code out of expected range: " + c, c >= 0 && c < values.length); + assertFalse("duplicate code " + c + " for " + t, seen[c]); + seen[c] = true; + } + } + + public void testIsFixedWidth() { + // Fixed-width types exchange values as raw long bits; only BYTE_ARRAY is variable-width. + assertTrue(ParquetPhysicalType.INT32.isFixedWidth()); + assertTrue(ParquetPhysicalType.INT64.isFixedWidth()); + assertTrue(ParquetPhysicalType.FLOAT.isFixedWidth()); + assertTrue(ParquetPhysicalType.DOUBLE.isFixedWidth()); + assertTrue(ParquetPhysicalType.BOOL.isFixedWidth()); + assertFalse(ParquetPhysicalType.BYTE_ARRAY.isFixedWidth()); + } +}