This is an automated email from the ASF dual-hosted git repository. JingsongLi pushed a commit to branch codex/batched-ftindex-open in repository https://gitbox.apache.org/repos/asf/paimon-full-text.git
commit 813ce2cd1295f7f733461a8d9d310fc31c3cae8e Author: JingsongLi <[email protected]> AuthorDate: Sun Jul 5 08:51:56 2026 +0800 Optimize full-text index open reads --- core/src/index.rs | 13 +- core/src/io.rs | 6 + core/src/storage.rs | 162 ++++++++++++++++++++- core/tests/core_roundtrip.rs | 61 +++++++- docs/java-api.md | 6 + .../paimon/index/fulltext/FullTextIndexInput.java | 21 +++ .../fulltext/FullTextNativeRoundTripTest.java | 20 ++- jni/src/lib.rs | 81 ++++++++--- 8 files changed, 341 insertions(+), 29 deletions(-) diff --git a/core/src/index.rs b/core/src/index.rs index 482a93f..f358f34 100644 --- a/core/src/index.rs +++ b/core/src/index.rs @@ -2,7 +2,9 @@ use crate::config::{FullTextIndexConfig, FullTextIndexMetadata}; use crate::error::{FtIndexError, Result}; use crate::io::{SeekRead, SeekWrite}; use crate::query::{BooleanOccur, FullTextQuery, MatchOperator}; -use crate::storage::{read_exact_at, read_header, write_envelope, ArchiveFileEntry, IndexHeader}; +use crate::storage::{ + read_archive_files_with, read_header, write_envelope, ArchiveFileEntry, IndexHeader, +}; use crate::tokenizer::{TokenizerConfig, TokenizerKind}; use roaring::RoaringTreemap; use std::fmt; @@ -110,11 +112,10 @@ impl<R: SeekRead> FullTextIndexReader<R> { pub fn open(mut input: R) -> Result<Self> { let (header, body_start) = read_header(&mut input)?; let directory = RamDirectory::create(); - for file in &header.files { - let mut data = vec![0u8; file.length as usize]; - read_exact_at(&mut input, body_start + file.offset, &mut data)?; - directory.atomic_write(Path::new(&file.name), &data)?; - } + read_archive_files_with(&mut input, body_start, &header.files, |name, data| { + directory.atomic_write(Path::new(name), data)?; + Ok(()) + })?; let mut index = Index::open(directory)?; register_tokenizer(&mut index, &header.metadata.config.tokenizer)?; Ok(Self { diff --git a/core/src/io.rs b/core/src/io.rs index 2d16b31..0320cd9 100644 --- a/core/src/io.rs +++ b/core/src/io.rs @@ -5,6 +5,12 @@ pub struct ReadRequest<'a> { pub buf: &'a mut [u8], } +impl<'a> ReadRequest<'a> { + pub fn new(pos: u64, buf: &'a mut [u8]) -> Self { + Self { pos, buf } + } +} + pub trait SeekRead: Send { fn pread(&mut self, ranges: &mut [ReadRequest<'_>]) -> io::Result<()>; } diff --git a/core/src/storage.rs b/core/src/storage.rs index 8e4121f..25f18a6 100644 --- a/core/src/storage.rs +++ b/core/src/storage.rs @@ -5,6 +5,8 @@ use serde::{Deserialize, Serialize}; pub const FORMAT_MAGIC: &[u8; 8] = b"PFTIDX01"; pub const FORMAT_VERSION: u32 = 1; +const MAX_ARCHIVE_READ_BATCH_BYTES: usize = 64 * 1024 * 1024; +const MAX_ARCHIVE_READ_BATCH_RANGES: usize = 64; #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] pub struct ArchiveFileEntry { @@ -59,8 +61,166 @@ pub fn read_header<R: SeekRead>(input: &mut R) -> Result<(IndexHeader, u64)> { Ok((header, 16 + header_len as u64)) } +pub fn read_archive_files<R: SeekRead>( + input: &mut R, + body_start: u64, + files: &[ArchiveFileEntry], +) -> Result<Vec<(String, Vec<u8>)>> { + read_archive_files_with(input, body_start, files, |name, data| { + Ok((name.to_string(), data.to_vec())) + }) +} + +pub fn read_archive_files_with<R, F, T>( + input: &mut R, + body_start: u64, + files: &[ArchiveFileEntry], + mut consume: F, +) -> Result<Vec<T>> +where + R: SeekRead, + F: FnMut(&str, &[u8]) -> Result<T>, +{ + let mut output = Vec::with_capacity(files.len()); + let mut pending = Vec::new(); + let mut pending_bytes = 0usize; + + for file in files { + let length = usize::try_from(file.length).map_err(|_| { + FtIndexError::InvalidStorage(format!("archive file '{}' is too large", file.name)) + })?; + let pos = body_start.checked_add(file.offset).ok_or_else(|| { + FtIndexError::InvalidStorage(format!("archive file '{}' offset overflow", file.name)) + })?; + + if !pending.is_empty() + && (pending.len() >= MAX_ARCHIVE_READ_BATCH_RANGES + || pending_bytes.saturating_add(length) > MAX_ARCHIVE_READ_BATCH_BYTES) + { + read_archive_file_batch(input, &mut pending, &mut consume, &mut output)?; + pending_bytes = 0; + } + + pending.push(PendingArchiveFile { + name: file.name.clone(), + pos, + data: vec![0u8; length], + }); + pending_bytes = pending_bytes.saturating_add(length); + } + + if !pending.is_empty() { + read_archive_file_batch(input, &mut pending, &mut consume, &mut output)?; + } + + Ok(output) +} + +struct PendingArchiveFile { + name: String, + pos: u64, + data: Vec<u8>, +} + +fn read_archive_file_batch<R, F, T>( + input: &mut R, + pending: &mut Vec<PendingArchiveFile>, + consume: &mut F, + output: &mut Vec<T>, +) -> Result<()> +where + R: SeekRead, + F: FnMut(&str, &[u8]) -> Result<T>, +{ + { + let mut requests: Vec<_> = pending + .iter_mut() + .map(|file| ReadRequest::new(file.pos, file.data.as_mut_slice())) + .collect(); + input.pread(&mut requests)?; + } + + for file in pending.drain(..) { + output.push(consume(&file.name, &file.data)?); + } + Ok(()) +} + pub fn read_exact_at<R: SeekRead>(input: &mut R, pos: u64, buf: &mut [u8]) -> Result<()> { - let mut request = [ReadRequest { pos, buf }]; + let mut request = [ReadRequest::new(pos, buf)]; input.pread(&mut request)?; Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + struct CountingReader { + data: Vec<u8>, + pread_batches: usize, + max_ranges_per_batch: usize, + } + + impl SeekRead for CountingReader { + fn pread(&mut self, ranges: &mut [ReadRequest<'_>]) -> std::io::Result<()> { + self.pread_batches += 1; + self.max_ranges_per_batch = self.max_ranges_per_batch.max(ranges.len()); + for range in ranges { + let start = usize::try_from(range.pos).map_err(|_| { + std::io::Error::new(std::io::ErrorKind::InvalidInput, "offset overflow") + })?; + let end = start.checked_add(range.buf.len()).ok_or_else(|| { + std::io::Error::new(std::io::ErrorKind::InvalidInput, "range overflow") + })?; + if end > self.data.len() { + return Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "read past end of slice", + )); + } + range.buf.copy_from_slice(&self.data[start..end]); + } + Ok(()) + } + } + + #[test] + fn read_archive_files_uses_one_batched_pread() { + let mut reader = CountingReader { + data: b"headerabcdef".to_vec(), + pread_batches: 0, + max_ranges_per_batch: 0, + }; + let files = vec![ + ArchiveFileEntry { + name: "a".to_string(), + offset: 0, + length: 2, + }, + ArchiveFileEntry { + name: "b".to_string(), + offset: 2, + length: 3, + }, + ArchiveFileEntry { + name: "c".to_string(), + offset: 5, + length: 1, + }, + ]; + + let files = read_archive_files(&mut reader, 6, &files).expect("archive files"); + + assert_eq!(reader.pread_batches, 1); + assert_eq!(reader.max_ranges_per_batch, 3); + assert_eq!( + files, + vec![ + ("a".to_string(), b"ab".to_vec()), + ("b".to_string(), b"cde".to_vec()), + ("c".to_string(), b"f".to_vec()), + ] + ); + } +} diff --git a/core/tests/core_roundtrip.rs b/core/tests/core_roundtrip.rs index aa38ac3..02a6087 100644 --- a/core/tests/core_roundtrip.rs +++ b/core/tests/core_roundtrip.rs @@ -1,10 +1,12 @@ -use paimon_ftindex_core::io::{PosWriter, SliceReader}; +use paimon_ftindex_core::io::{PosWriter, ReadRequest, SeekRead, SliceReader}; use paimon_ftindex_core::{ FullTextIndexConfig, FullTextIndexReader, FullTextIndexWriter, FullTextQuery, MatchOperator, TokenizerConfig, TokenizerKind, }; use roaring::RoaringTreemap; use std::collections::HashMap; +use std::io; +use std::sync::{Arc, Mutex}; fn build_index() -> anyhow::Result<Vec<u8>> { let mut writer = FullTextIndexWriter::new(FullTextIndexConfig::new())?; @@ -20,6 +22,63 @@ fn build_index() -> anyhow::Result<Vec<u8>> { Ok(bytes) } +#[derive(Default)] +struct ReadStats { + max_ranges_per_batch: usize, +} + +struct CountingSliceReader { + data: Vec<u8>, + stats: Arc<Mutex<ReadStats>>, +} + +impl CountingSliceReader { + fn new(data: Vec<u8>, stats: Arc<Mutex<ReadStats>>) -> Self { + Self { data, stats } + } +} + +impl SeekRead for CountingSliceReader { + fn pread(&mut self, ranges: &mut [ReadRequest<'_>]) -> io::Result<()> { + { + let mut stats = self.stats.lock().unwrap(); + stats.max_ranges_per_batch = stats.max_ranges_per_batch.max(ranges.len()); + } + for range in ranges { + let start = usize::try_from(range.pos) + .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "offset overflow"))?; + let end = start + .checked_add(range.buf.len()) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "range overflow"))?; + if end > self.data.len() { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "read past end of slice", + )); + } + range.buf.copy_from_slice(&self.data[start..end]); + } + Ok(()) + } +} + +#[test] +fn reader_open_batches_archive_file_reads() -> anyhow::Result<()> { + let bytes = build_index()?; + let stats = Arc::new(Mutex::new(ReadStats::default())); + let input = CountingSliceReader::new(bytes, Arc::clone(&stats)); + + let mut reader = FullTextIndexReader::open(input)?; + let result = reader.search(FullTextQuery::match_query("paimon", "text"), 10)?; + + assert_eq!(result.row_ids.len(), 2); + assert!( + stats.lock().unwrap().max_ranges_per_batch > 1, + "reader open should batch archive file reads into one pread call" + ); + Ok(()) +} + #[test] fn match_query_round_trip() -> anyhow::Result<()> { let bytes = build_index()?; diff --git a/docs/java-api.md b/docs/java-api.md index 3a9535a..de859ef 100644 --- a/docs/java-api.md +++ b/docs/java-api.md @@ -34,6 +34,12 @@ try (FullTextIndexReader reader = new FullTextIndexReader(input)) { (`RoaringTreemap`) containing the allowed row ids. The filter is applied during Tantivy collection, before the top results are selected. +Input reads: + +- Implement `FullTextIndexInput.readAt(...)` for compatibility. +- Override `FullTextIndexInput.pread(long[] positions, byte[][] buffers)` when + the storage client can batch, coalesce, or parallelize positional reads. + Native loading: - Set `PAIMON_FTINDEX_JNI_LIB_PATH` to the full path of diff --git a/java/src/main/java/org/apache/paimon/index/fulltext/FullTextIndexInput.java b/java/src/main/java/org/apache/paimon/index/fulltext/FullTextIndexInput.java index ccd5452..5e37369 100644 --- a/java/src/main/java/org/apache/paimon/index/fulltext/FullTextIndexInput.java +++ b/java/src/main/java/org/apache/paimon/index/fulltext/FullTextIndexInput.java @@ -5,4 +5,25 @@ import java.io.IOException; public interface FullTextIndexInput { void readAt(long position, byte[] buffer, int offset, int length) throws IOException; + + default void pread(long[] positions, byte[][] buffers) throws IOException { + if (positions == null) { + throw new NullPointerException("positions"); + } + if (buffers == null) { + throw new NullPointerException("buffers"); + } + if (positions.length != buffers.length) { + throw new IllegalArgumentException( + "positions length " + positions.length + " does not match buffers length " + + buffers.length); + } + for (int i = 0; i < positions.length; i++) { + byte[] buffer = buffers[i]; + if (buffer == null) { + throw new NullPointerException("buffers[" + i + "]"); + } + readAt(positions[i], buffer, 0, buffer.length); + } + } } diff --git a/java/src/test/java/org/apache/paimon/index/fulltext/FullTextNativeRoundTripTest.java b/java/src/test/java/org/apache/paimon/index/fulltext/FullTextNativeRoundTripTest.java index c5c657a..77a441e 100644 --- a/java/src/test/java/org/apache/paimon/index/fulltext/FullTextNativeRoundTripTest.java +++ b/java/src/test/java/org/apache/paimon/index/fulltext/FullTextNativeRoundTripTest.java @@ -6,6 +6,7 @@ import java.io.ByteArrayOutputStream; import java.io.File; import java.io.IOException; import java.util.Collections; +import java.util.concurrent.atomic.AtomicInteger; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; @@ -28,13 +29,26 @@ public class FullTextNativeRoundTripTest { } byte[] indexBytes = output.toByteArray(); + AtomicInteger maxBatchSize = new AtomicInteger(); FullTextIndexInput input = - (position, buffer, offset, length) -> - readAt(indexBytes, position, buffer, offset, length); + new FullTextIndexInput() { + @Override + public void readAt(long position, byte[] buffer, int offset, int length) + throws IOException { + readAtBytes(indexBytes, position, buffer, offset, length); + } + + @Override + public void pread(long[] positions, byte[][] buffers) throws IOException { + maxBatchSize.updateAndGet(current -> Math.max(current, positions.length)); + FullTextIndexInput.super.pread(positions, buffers); + } + }; try (FullTextIndexReader reader = new FullTextIndexReader(input)) { FullTextSearchResult result = reader.search(FullTextQuery.match("paimon", "text"), 10); + assertTrue(maxBatchSize.get() > 1); assertEquals(1, result.size()); assertEquals(10L, result.rowIds()[0]); assertTrue(result.scores()[0] > 0.0f); @@ -53,7 +67,7 @@ public class FullTextNativeRoundTripTest { return path != null && !path.isEmpty() && new File(path).isFile(); } - private static void readAt( + private static void readAtBytes( byte[] source, long position, byte[] buffer, int offset, int length) throws IOException { long end = position + length; if (position < 0 || end > source.length || end < position) { diff --git a/jni/src/lib.rs b/jni/src/lib.rs index 953beaa..ffbb0f5 100644 --- a/jni/src/lib.rs +++ b/jni/src/lib.rs @@ -63,37 +63,82 @@ unsafe impl Send for JavaInput {} impl SeekRead for JavaInput { fn pread(&mut self, ranges: &mut [ReadRequest<'_>]) -> io::Result<()> { + if ranges.is_empty() { + return Ok(()); + } + let mut env = self .jvm .attach_current_thread() .map_err(|e| io::Error::other(format!("JNI attach failed: {e}")))?; - for range in ranges { - let array = env + + let positions = env + .new_long_array(ranges.len() as i32) + .map_err(|e| io::Error::other(format!("new_long_array failed: {e}")))?; + let position_values: Vec<i64> = ranges.iter().map(|range| range.pos as i64).collect(); + env.set_long_array_region(&positions, 0, &position_values) + .map_err(|e| io::Error::other(format!("set_long_array_region failed: {e}")))?; + + let byte_array_class = env + .find_class("[B") + .map_err(|e| io::Error::other(format!("find byte[] class failed: {e}")))?; + let buffers = env + .new_object_array(ranges.len() as i32, byte_array_class, JObject::null()) + .map_err(|e| io::Error::other(format!("new_object_array failed: {e}")))?; + for (idx, range) in ranges.iter().enumerate() { + let buffer = env .new_byte_array(range.buf.len() as i32) .map_err(|e| io::Error::other(format!("new_byte_array failed: {e}")))?; - let array_obj = JObject::from(array); - env.call_method( - self.input.as_obj(), - "readAt", - "(J[BII)V", - &[ - JValue::Long(range.pos as i64), - JValue::Object(&array_obj), - JValue::Int(0), - JValue::Int(range.buf.len() as i32), - ], - ) - .map_err(|e| io::Error::other(format!("Java input readAt failed: {e}")))?; - let array = JByteArray::from(array_obj); + env.set_object_array_element(&buffers, idx as i32, buffer) + .map_err(|e| io::Error::other(format!("set_object_array_element failed: {e}")))?; + } + + env.call_method( + self.input.as_obj(), + "pread", + "([J[[B)V", + &[JValue::Object(&positions), JValue::Object(&buffers)], + ) + .map_err(|e| io::Error::other(format!("Java input pread failed: {e}")))?; + + copy_java_buffers(&mut env, &buffers, ranges) + } +} + +fn copy_java_buffers( + env: &mut JNIEnv<'_>, + buffers: &JObjectArray<'_>, + ranges: &mut [ReadRequest<'_>], +) -> io::Result<()> { + for (idx, range) in ranges.iter_mut().enumerate() { + let object = env + .get_object_array_element(buffers, idx as i32) + .map_err(|e| io::Error::other(format!("get_object_array_element failed: {e}")))?; + let buffer = JByteArray::from(object); + let length = env + .get_array_length(&buffer) + .map_err(|e| io::Error::other(format!("get_array_length failed: {e}")))? + as usize; + if length != range.buf.len() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "Java pread returned buffer length {} != {}", + length, + range.buf.len() + ), + )); + } + if length > 0 { let mut signed = vec![0i8; range.buf.len()]; - env.get_byte_array_region(&array, 0, &mut signed) + env.get_byte_array_region(&buffer, 0, &mut signed) .map_err(|e| io::Error::other(format!("get_byte_array_region failed: {e}")))?; for (dst, src) in range.buf.iter_mut().zip(signed) { *dst = src as u8; } } - Ok(()) } + Ok(()) } struct WriterHandle {
