This is an automated email from the ASF dual-hosted git repository. JingsongLi pushed a commit to branch codex/roaring-filter-search in repository https://gitbox.apache.org/repos/asf/paimon-full-text.git
commit dcd3ed4f49ba5c156ceacdc36f7c4877f7267a14 Author: JingsongLi <[email protected]> AuthorDate: Sat Jul 4 22:52:22 2026 +0800 Add Roaring row-id filters to full-text search --- Cargo.lock | 18 ++ Cargo.toml | 1 + README.md | 17 +- core/Cargo.toml | 1 + core/src/index.rs | 39 ++++- core/tests/core_roundtrip.rs | 84 +++++++++ docs/designs/2026-07-04-paimon-full-text-index.md | 21 ++- docs/java-api.md | 6 + docs/paimon-integration.md | 2 + docs/python-api.md | 11 +- ffi/Cargo.toml | 1 + ffi/src/lib.rs | 192 +++++++++++++++++---- include/paimon_ftindex.h | 11 ++ .../paimon/index/fulltext/FullTextIndexReader.java | 14 ++ .../paimon/index/fulltext/FullTextNative.java | 3 + jni/src/lib.rs | 48 +++++- python/paimon_ftindex/_ffi.py | 17 ++ python/paimon_ftindex/reader.py | 24 ++- python/tests/test_roundtrip.py | 57 ++++++ 19 files changed, 517 insertions(+), 50 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 35b4f2b..ccb6369 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -76,6 +76,12 @@ version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index"; checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index"; +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" + [[package]] name = "byteorder" version = "1.5.0" @@ -577,6 +583,7 @@ name = "paimon-ftindex-core" version = "0.1.0" dependencies = [ "anyhow", + "roaring", "serde", "serde_json", "tantivy", @@ -590,6 +597,7 @@ version = "0.1.0" dependencies = [ "anyhow", "paimon-ftindex-core", + "roaring", "serde_json", ] @@ -749,6 +757,16 @@ version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index"; checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" +[[package]] +name = "roaring" +version = "0.11.4" +source = "registry+https://github.com/rust-lang/crates.io-index"; +checksum = "1dedc5658c6ecb3bdb5ef5f3295bb9253f42dcf3fd1402c03f6b1f7659c3c4a9" +dependencies = [ + "bytemuck", + "byteorder", +] + [[package]] name = "rust-stemmers" version = "1.2.0" diff --git a/Cargo.toml b/Cargo.toml index 1e0c7ec..8a2e210 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,6 +13,7 @@ anyhow = "1" jni = "0.21" serde = { version = "1", features = ["derive"] } serde_json = "1" +roaring = "0.11" tantivy = "0.22" tempfile = "3" thiserror = "1" diff --git a/README.md b/README.md index 6c0f2a8..89e97a7 100644 --- a/README.md +++ b/README.md @@ -17,7 +17,8 @@ and do not depend on Paimon manifest metadata. Implemented: - Rust writer, reader, v1 envelope, and search. -- C FFI writer/reader/search JSON. +- C FFI writer/reader/search JSON, including serialized 64-bit Roaring row-id + filters. - Java API and JNI bridge. - Python ctypes package. - Cross-boundary round-trip tests for Rust core, FFI, and Python. @@ -67,6 +68,17 @@ let mut reader = FullTextIndexReader::open(SliceReader::new(bytes))?; let result = reader.search(FullTextQuery::match_query("paimon", "text"), 10)?; ``` +To restrict search to an upstream candidate set, pass a serialized +`RoaringTreemap` of allowed row ids: + +```rust +let filtered = reader.search_with_roaring_filter( + FullTextQuery::match_query("paimon", "text"), + 10, + roaring_filter_bytes, +)?; +``` + ## Python Example ```python @@ -86,4 +98,7 @@ class Input: with FullTextIndexReader(Input(out.getvalue())) as reader: ids, scores = reader.search(MatchQuery("paimon"), limit=10) + filtered_ids, filtered_scores = reader.search( + MatchQuery("paimon"), limit=10, filter_bytes=roaring_filter_bytes + ) ``` diff --git a/core/Cargo.toml b/core/Cargo.toml index a1124c0..1f07c33 100644 --- a/core/Cargo.toml +++ b/core/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true rust-version.workspace = true [dependencies] +roaring.workspace = true serde.workspace = true serde_json.workspace = true tantivy.workspace = true diff --git a/core/src/index.rs b/core/src/index.rs index 85bc987..30e7f1f 100644 --- a/core/src/index.rs +++ b/core/src/index.rs @@ -4,10 +4,11 @@ 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::tokenizer::{TokenizerConfig, TokenizerKind}; +use roaring::RoaringTreemap; use std::fmt; use std::fs; use std::path::Path; -use tantivy::collector::TopDocs; +use tantivy::collector::{FilterCollector, TopDocs}; use tantivy::directory::{Directory, RamDirectory}; use tantivy::query::{ BooleanQuery, BoostQuery, EnableScoring, Explanation, Occur, Query, QueryParser, Scorer, Weight, @@ -131,6 +132,25 @@ impl<R: SeekRead> FullTextIndexReader<R> { } pub fn search(&mut self, query: FullTextQuery, limit: usize) -> Result<FullTextSearchResult> { + self.search_with_filter(query, limit, None) + } + + pub fn search_with_roaring_filter( + &mut self, + query: FullTextQuery, + limit: usize, + roaring_filter_bytes: &[u8], + ) -> Result<FullTextSearchResult> { + let filter = decode_roaring_filter(roaring_filter_bytes)?; + self.search_with_filter(query, limit, Some(filter)) + } + + fn search_with_filter( + &mut self, + query: FullTextQuery, + limit: usize, + filter: Option<RoaringTreemap>, + ) -> Result<FullTextSearchResult> { if limit == 0 { return Err(FtIndexError::InvalidQuery( "search limit must be positive".to_string(), @@ -139,7 +159,17 @@ impl<R: SeekRead> FullTextIndexReader<R> { let reader = self.index.reader()?; let searcher = reader.searcher(); let tantivy_query = build_query(&self.index, &self.metadata.config, &query)?; - let top_docs = searcher.search(&tantivy_query, &TopDocs::with_limit(limit))?; + let top_docs = if let Some(filter) = filter { + let row_id_field = self.metadata.config.row_id_field.clone(); + let collector = FilterCollector::new( + row_id_field, + move |row_id: u64| filter.contains(row_id), + TopDocs::with_limit(limit), + ); + searcher.search(&tantivy_query, &collector)? + } else { + searcher.search(&tantivy_query, &TopDocs::with_limit(limit))? + }; let mut row_ids = Vec::with_capacity(top_docs.len()); let mut scores = Vec::with_capacity(top_docs.len()); for (score, doc_address) in top_docs { @@ -155,6 +185,11 @@ impl<R: SeekRead> FullTextIndexReader<R> { } } +fn decode_roaring_filter(bytes: &[u8]) -> Result<RoaringTreemap> { + RoaringTreemap::deserialize_from(bytes) + .map_err(|e| FtIndexError::InvalidQuery(format!("invalid RoaringTreemap filter: {e}"))) +} + fn build_schema(config: &FullTextIndexConfig) -> Schema { let mut builder = Schema::builder(); builder.add_u64_field( diff --git a/core/tests/core_roundtrip.rs b/core/tests/core_roundtrip.rs index 3f1b009..00d66ec 100644 --- a/core/tests/core_roundtrip.rs +++ b/core/tests/core_roundtrip.rs @@ -2,6 +2,7 @@ use paimon_ftindex_core::io::{PosWriter, SliceReader}; use paimon_ftindex_core::{ FullTextIndexConfig, FullTextIndexReader, FullTextIndexWriter, FullTextQuery, MatchOperator, }; +use roaring::RoaringTreemap; fn build_index() -> anyhow::Result<Vec<u8>> { let mut writer = FullTextIndexWriter::new(FullTextIndexConfig::new())?; @@ -47,6 +48,89 @@ fn match_query_and_operator_filters_terms() -> anyhow::Result<()> { Ok(()) } +#[test] +fn search_with_roaring_filter_limits_allowed_row_ids_before_top_docs() -> anyhow::Result<()> { + let bytes = build_index()?; + let query = FullTextQuery::match_query("paimon", "text"); + + let mut reader = FullTextIndexReader::open(SliceReader::new(bytes.clone()))?; + let unfiltered_top = reader.search(query.clone(), 1)?.row_ids[0]; + let allowed_id = if unfiltered_top == 10 { 12 } else { 10 }; + + let mut allowed = RoaringTreemap::new(); + allowed.insert(allowed_id as u64); + let mut filter_bytes = Vec::new(); + allowed.serialize_into(&mut filter_bytes)?; + + let mut reader = FullTextIndexReader::open(SliceReader::new(bytes))?; + let result = reader.search_with_roaring_filter(query, 1, &filter_bytes)?; + + assert_eq!(result.row_ids, vec![allowed_id]); + assert_eq!(result.scores.len(), 1); + Ok(()) +} + +#[test] +fn search_with_empty_roaring_filter_returns_empty_results() -> anyhow::Result<()> { + let bytes = build_index()?; + let empty = RoaringTreemap::new(); + let mut filter_bytes = Vec::new(); + empty.serialize_into(&mut filter_bytes)?; + + let mut reader = FullTextIndexReader::open(SliceReader::new(bytes))?; + let result = reader.search_with_roaring_filter( + FullTextQuery::match_query("paimon", "text"), + 10, + &filter_bytes, + )?; + + assert!(result.row_ids.is_empty()); + assert!(result.scores.is_empty()); + Ok(()) +} + +#[test] +fn search_with_roaring_filter_supports_64_bit_row_ids() -> anyhow::Result<()> { + let allowed_id = (1i64 << 33) + 17; + let mut writer = FullTextIndexWriter::new(FullTextIndexConfig::new())?; + writer.add_document(1, "apache paimon")?; + writer.add_document(allowed_id, "paimon filtered row")?; + + let mut bytes = Vec::new(); + writer.write(&mut PosWriter::new(&mut bytes))?; + + let mut allowed = RoaringTreemap::new(); + allowed.insert(allowed_id as u64); + let mut filter_bytes = Vec::new(); + allowed.serialize_into(&mut filter_bytes)?; + + let mut reader = FullTextIndexReader::open(SliceReader::new(bytes))?; + let result = reader.search_with_roaring_filter( + FullTextQuery::match_query("paimon", "text"), + 10, + &filter_bytes, + )?; + + assert_eq!(result.row_ids, vec![allowed_id]); + Ok(()) +} + +#[test] +fn search_rejects_invalid_roaring_filter_bytes() -> anyhow::Result<()> { + let bytes = build_index()?; + let mut reader = FullTextIndexReader::open(SliceReader::new(bytes))?; + let err = reader + .search_with_roaring_filter( + FullTextQuery::match_query("paimon", "text"), + 10, + b"not roaring", + ) + .expect_err("invalid filter bytes should fail"); + + assert!(err.to_string().contains("invalid RoaringTreemap filter")); + Ok(()) +} + #[test] fn phrase_query_uses_positions() -> anyhow::Result<()> { let bytes = build_index()?; diff --git a/docs/designs/2026-07-04-paimon-full-text-index.md b/docs/designs/2026-07-04-paimon-full-text-index.md index 78bb0e5..02e174f 100644 --- a/docs/designs/2026-07-04-paimon-full-text-index.md +++ b/docs/designs/2026-07-04-paimon-full-text-index.md @@ -250,6 +250,17 @@ int paimon_ftindex_reader_search_json( float *scores, size_t capacity, size_t *result_len); + +int paimon_ftindex_reader_search_json_with_roaring_filter( + PaimonFtindexReaderHandle *reader, + const char *query_json, + size_t limit, + const uint8_t *roaring_filter, + size_t roaring_filter_len, + int64_t *row_ids, + float *scores, + size_t capacity, + size_t *result_len); ``` All functions return `0` on success and `-1` on error. A thread-local @@ -291,6 +302,8 @@ try (FullTextIndexWriter writer = FullTextIndexWriter.create(options)) { try (FullTextIndexReader reader = new FullTextIndexReader(input)) { FullTextSearchResult result = reader.search(FullTextQuery.match("paimon", "text"), 10); + FullTextSearchResult filtered = + reader.search(FullTextQuery.match("paimon", "text"), 10, roaringFilterBytes); } ``` @@ -320,6 +333,11 @@ writer.write(output) reader = FullTextIndexReader(input) ids, scores = reader.search(MatchQuery("paimon", column="text"), limit=10) +filtered_ids, filtered_scores = reader.search( + MatchQuery("paimon", column="text"), + limit=10, + filter_bytes=roaring_filter_bytes, +) ``` Python I/O protocol: @@ -441,6 +459,3 @@ matching the caution already present in upstream `paimon-tantivy-jni`. - Should Java native loading copy prebuilt libraries from resources, or should this repository initially require `PAIMON_FTINDEX_LIB_PATH` like the Python package? -- Should v1 expose row-id prefiltering with serialized 64-bit Roaring bitmaps, - matching vector-index, or defer it until Paimon needs hybrid scalar plus - full-text search? diff --git a/docs/java-api.md b/docs/java-api.md index 73aa5d4..3a9535a 100644 --- a/docs/java-api.md +++ b/docs/java-api.md @@ -25,9 +25,15 @@ try (FullTextIndexWriter writer = FullTextIndexWriter.create(Collections.emptyMa try (FullTextIndexReader reader = new FullTextIndexReader(input)) { FullTextSearchResult result = reader.search(FullTextQuery.match("paimon", "text"), 10); + FullTextSearchResult filtered = + reader.search(FullTextQuery.match("paimon", "text"), 10, roaringFilterBytes); } ``` +`roaringFilterBytes` must be a serialized 64-bit Roaring bitmap +(`RoaringTreemap`) containing the allowed row ids. The filter is applied during +Tantivy collection, before the top results are selected. + Native loading: - Set `PAIMON_FTINDEX_JNI_LIB_PATH` to the full path of diff --git a/docs/paimon-integration.md b/docs/paimon-integration.md index 26f82b3..e5da028 100644 --- a/docs/paimon-integration.md +++ b/docs/paimon-integration.md @@ -7,6 +7,8 @@ integration should be a thin adapter: - Serialize queries to the JSON accepted by this library. - Implement a Paimon `GlobalIndexerFactory` that delegates to Java `FullTextIndexWriter` and `FullTextIndexReader`. +- Pass serialized 64-bit Roaring row-id filters to reader search when another + index or predicate pushdown has already produced an allowed candidate set. - Store produced files as global index files. Suggested index identifier: diff --git a/docs/python-api.md b/docs/python-api.md index 5411d83..a5e0b65 100644 --- a/docs/python-api.md +++ b/docs/python-api.md @@ -17,8 +17,15 @@ with FullTextIndexWriter({"tokenizer": "ngram"}) as writer: with FullTextIndexReader(input_) as reader: ids, scores = reader.search(MatchQuery("paimon"), limit=10) + filtered_ids, filtered_scores = reader.search( + MatchQuery("paimon"), limit=10, filter_bytes=roaring_filter_bytes + ) ``` +`filter_bytes` must be a serialized 64-bit Roaring bitmap (`RoaringTreemap`) +containing the allowed row ids. The filter is applied during Tantivy +collection, before the top results are selected. + The output object must provide: - `write(bytes)` @@ -31,5 +38,5 @@ The input object must provide: Native loading: - Set `PAIMON_FTINDEX_LIB_PATH` to a library file or directory, or -- build `paimon-ftindex-ffi` so the package can discover `target/debug` or - `target/release`. +- build `paimon-ftindex-ffi` so the package can discover `target/debug`, + `target/debug/deps`, `target/release`, or `target/release/deps`. diff --git a/ffi/Cargo.toml b/ffi/Cargo.toml index f552b3e..b8f479d 100644 --- a/ffi/Cargo.toml +++ b/ffi/Cargo.toml @@ -16,3 +16,4 @@ serde_json.workspace = true [dev-dependencies] anyhow.workspace = true +roaring.workspace = true diff --git a/ffi/src/lib.rs b/ffi/src/lib.rs index 4840eff..49ce058 100644 --- a/ffi/src/lib.rs +++ b/ffi/src/lib.rs @@ -237,38 +237,37 @@ pub unsafe extern "C" fn paimon_ftindex_reader_search_json( ) -> c_int { ffi_status(|| { let reader = require_mut(reader, "reader")?; - if result_len.is_null() { - return Err("result_len is null".to_string()); - } - let query_json = cstr_to_string(query_json, "query_json")?; - let query = FullTextQuery::from_json(&query_json).map_err(|e| e.to_string())?; - let result = reader - .inner - .search(query, limit) - .map_err(|e| e.to_string())?; - unsafe { - *result_len = result.row_ids.len(); - } - if result.row_ids.len() > capacity { - return Err(format!( - "result capacity {} is smaller than result length {}", - capacity, - result.row_ids.len() - )); - } - if !result.row_ids.is_empty() { - if row_ids.is_null() { - return Err("row_ids is null".to_string()); - } - if scores.is_null() { - return Err("scores is null".to_string()); - } - unsafe { - ptr::copy_nonoverlapping(result.row_ids.as_ptr(), row_ids, result.row_ids.len()); - ptr::copy_nonoverlapping(result.scores.as_ptr(), scores, result.scores.len()); - } - } - Ok(()) + search_json_impl( + reader, query_json, limit, None, row_ids, scores, capacity, result_len, + ) + }) +} + +#[no_mangle] +pub unsafe extern "C" fn paimon_ftindex_reader_search_json_with_roaring_filter( + reader: *mut PaimonFtindexReaderHandle, + query_json: *const c_char, + limit: usize, + roaring_filter: *const u8, + roaring_filter_len: usize, + row_ids: *mut i64, + scores: *mut f32, + capacity: usize, + result_len: *mut usize, +) -> c_int { + ffi_status(|| { + let reader = require_mut(reader, "reader")?; + let roaring_filter = const_slice(roaring_filter, roaring_filter_len, "roaring_filter")?; + search_json_impl( + reader, + query_json, + limit, + Some(roaring_filter), + row_ids, + scores, + capacity, + result_len, + ) }) } @@ -308,6 +307,16 @@ unsafe fn require_mut<'a, T>(ptr: *mut T, name: &str) -> Result<&'a mut T, Strin ptr.as_mut().ok_or_else(|| format!("{name} is null")) } +unsafe fn const_slice<'a, T>(ptr: *const T, len: usize, name: &str) -> Result<&'a [T], String> { + if len == 0 { + Ok(&[]) + } else if ptr.is_null() { + Err(format!("{name} is null")) + } else { + Ok(slice::from_raw_parts(ptr, len)) + } +} + unsafe fn cstr_to_string(ptr: *const c_char, name: &str) -> Result<String, String> { if ptr.is_null() { return Err(format!("{name} is null")); @@ -318,9 +327,61 @@ unsafe fn cstr_to_string(ptr: *const c_char, name: &str) -> Result<String, Strin .map_err(|e| format!("{name} is not valid UTF-8: {e}")) } +fn search_json_impl( + reader: &mut PaimonFtindexReaderHandle, + query_json: *const c_char, + limit: usize, + roaring_filter: Option<&[u8]>, + row_ids: *mut i64, + scores: *mut f32, + capacity: usize, + result_len: *mut usize, +) -> Result<(), String> { + if result_len.is_null() { + return Err("result_len is null".to_string()); + } + let query_json = unsafe { cstr_to_string(query_json, "query_json") }?; + let query = FullTextQuery::from_json(&query_json).map_err(|e| e.to_string())?; + let result = if let Some(roaring_filter) = roaring_filter { + reader + .inner + .search_with_roaring_filter(query, limit, roaring_filter) + .map_err(|e| e.to_string())? + } else { + reader + .inner + .search(query, limit) + .map_err(|e| e.to_string())? + }; + unsafe { + *result_len = result.row_ids.len(); + } + if result.row_ids.len() > capacity { + return Err(format!( + "result capacity {} is smaller than result length {}", + capacity, + result.row_ids.len() + )); + } + if !result.row_ids.is_empty() { + if row_ids.is_null() { + return Err("row_ids is null".to_string()); + } + if scores.is_null() { + return Err("scores is null".to_string()); + } + unsafe { + ptr::copy_nonoverlapping(result.row_ids.as_ptr(), row_ids, result.row_ids.len()); + ptr::copy_nonoverlapping(result.scores.as_ptr(), scores, result.scores.len()); + } + } + Ok(()) +} + #[cfg(test)] mod tests { use super::*; + use roaring::RoaringTreemap; use std::ffi::CString; unsafe extern "C" fn write_vec(ctx: *mut c_void, buf: *const u8, len: usize) -> c_int { @@ -395,4 +456,71 @@ mod tests { paimon_ftindex_reader_free(reader); } } + + #[test] + fn ffi_round_trip_search_json_with_roaring_filter() { + unsafe { + let writer = paimon_ftindex_writer_open(ptr::null(), ptr::null(), 0); + assert!(!writer.is_null()); + + let text = CString::new("Apache Paimon full text").unwrap(); + assert_eq!( + paimon_ftindex_writer_add_document(writer, 7, text.as_ptr()), + 0 + ); + let text = CString::new("Paimon filtered row").unwrap(); + assert_eq!( + paimon_ftindex_writer_add_document(writer, 9, text.as_ptr()), + 0 + ); + + let mut bytes = Vec::new(); + let output = PaimonFtindexOutputFile { + ctx: &mut bytes as *mut Vec<u8> as *mut c_void, + write_fn: Some(write_vec), + flush_fn: None, + }; + assert_eq!(paimon_ftindex_writer_write_index(writer, output), 0); + paimon_ftindex_writer_free(writer); + + let input = PaimonFtindexInputFile { + ctx: &bytes as *const Vec<u8> as *mut c_void, + read_at_fn: Some(read_vec), + }; + let reader = paimon_ftindex_reader_open(input); + assert!(!reader.is_null()); + + let query = CString::new( + paimon_ftindex_core::FullTextQuery::match_query("paimon", "text") + .to_json() + .unwrap(), + ) + .unwrap(); + let mut allowed = RoaringTreemap::new(); + allowed.insert(9); + let mut filter_bytes = Vec::new(); + allowed.serialize_into(&mut filter_bytes).unwrap(); + let mut row_ids = [0i64; 4]; + let mut scores = [0f32; 4]; + let mut result_len = 0usize; + assert_eq!( + paimon_ftindex_reader_search_json_with_roaring_filter( + reader, + query.as_ptr(), + 4, + filter_bytes.as_ptr(), + filter_bytes.len(), + row_ids.as_mut_ptr(), + scores.as_mut_ptr(), + row_ids.len(), + &mut result_len, + ), + 0 + ); + assert_eq!(result_len, 1); + assert_eq!(row_ids[0], 9); + assert!(scores[0] > 0.0); + paimon_ftindex_reader_free(reader); + } + } } diff --git a/include/paimon_ftindex.h b/include/paimon_ftindex.h index 08c9237..b4d2cea 100644 --- a/include/paimon_ftindex.h +++ b/include/paimon_ftindex.h @@ -51,6 +51,17 @@ int paimon_ftindex_reader_search_json( size_t capacity, size_t *result_len); +int paimon_ftindex_reader_search_json_with_roaring_filter( + PaimonFtindexReaderHandle *reader, + const char *query_json, + size_t limit, + const uint8_t *roaring_filter, + size_t roaring_filter_len, + int64_t *row_ids, + float *scores, + size_t capacity, + size_t *result_len); + void paimon_ftindex_reader_free(PaimonFtindexReaderHandle *reader); #ifdef __cplusplus diff --git a/java/src/main/java/org/apache/paimon/index/fulltext/FullTextIndexReader.java b/java/src/main/java/org/apache/paimon/index/fulltext/FullTextIndexReader.java index 717af65..1520532 100644 --- a/java/src/main/java/org/apache/paimon/index/fulltext/FullTextIndexReader.java +++ b/java/src/main/java/org/apache/paimon/index/fulltext/FullTextIndexReader.java @@ -21,6 +21,20 @@ public final class FullTextIndexReader implements AutoCloseable { return FullTextNative.searchJson(requireOpen(), query.toJson(), limit); } + public FullTextSearchResult search(FullTextQuery query, int limit, byte[] roaringFilter) { + if (query == null) { + throw new NullPointerException("query"); + } + if (limit <= 0) { + throw new IllegalArgumentException("search limit must be positive"); + } + if (roaringFilter == null) { + throw new NullPointerException("roaringFilter"); + } + return FullTextNative.searchJsonWithRoaringFilter( + requireOpen(), query.toJson(), limit, roaringFilter); + } + @Override public void close() { long ptr = nativePtr; diff --git a/java/src/main/java/org/apache/paimon/index/fulltext/FullTextNative.java b/java/src/main/java/org/apache/paimon/index/fulltext/FullTextNative.java index d573941..5d2a07e 100644 --- a/java/src/main/java/org/apache/paimon/index/fulltext/FullTextNative.java +++ b/java/src/main/java/org/apache/paimon/index/fulltext/FullTextNative.java @@ -25,5 +25,8 @@ final class FullTextNative { static native FullTextSearchResult searchJson(long readerPtr, String queryJson, int limit); + static native FullTextSearchResult searchJsonWithRoaringFilter( + long readerPtr, String queryJson, int limit, byte[] roaringFilter); + static native void freeReader(long readerPtr); } diff --git a/jni/src/lib.rs b/jni/src/lib.rs index e78af3d..953beaa 100644 --- a/jni/src/lib.rs +++ b/jni/src/lib.rs @@ -175,7 +175,28 @@ pub extern "system" fn Java_org_apache_paimon_index_fulltext_FullTextNative_sear query_json: JString, limit: jint, ) -> jobject { - match search_json(&mut env, reader_ptr, query_json, limit) { + match search_json(&mut env, reader_ptr, query_json, limit, None) { + Ok(obj) => obj, + Err(e) => throw_and_return(&mut env, &e, ptr::null_mut()), + } +} + +#[no_mangle] +pub extern "system" fn Java_org_apache_paimon_index_fulltext_FullTextNative_searchJsonWithRoaringFilter( + mut env: JNIEnv, + _class: JClass, + reader_ptr: jlong, + query_json: JString, + limit: jint, + roaring_filter: JByteArray, +) -> jobject { + match search_json( + &mut env, + reader_ptr, + query_json, + limit, + Some(roaring_filter), + ) { Ok(obj) => obj, Err(e) => throw_and_return(&mut env, &e, ptr::null_mut()), } @@ -244,6 +265,7 @@ fn search_json( reader_ptr: jlong, query_json: JString, limit: jint, + roaring_filter: Option<JByteArray>, ) -> Result<jobject, String> { let reader = handle_mut::<ReaderHandle>(reader_ptr, "reader")?; let query_json: String = env @@ -252,10 +274,18 @@ fn search_json( .into(); let query = FullTextQuery::from_json(&query_json).map_err(|e| e.to_string())?; let limit = validate_search_limit(limit)?; - let result = reader - .inner - .search(query, limit) - .map_err(|e| e.to_string())?; + let result = if let Some(roaring_filter) = roaring_filter { + let roaring_filter = read_byte_array(env, roaring_filter)?; + reader + .inner + .search_with_roaring_filter(query, limit, &roaring_filter) + .map_err(|e| e.to_string())? + } else { + reader + .inner + .search(query, limit) + .map_err(|e| e.to_string())? + }; let row_ids = env .new_long_array(result.row_ids.len() as i32) @@ -287,6 +317,14 @@ fn validate_search_limit(limit: jint) -> Result<usize, String> { Ok(limit as usize) } +fn read_byte_array(env: &mut JNIEnv, array: JByteArray) -> Result<Vec<u8>, String> { + if array.as_raw().is_null() { + return Err("roaringFilter is null".to_string()); + } + env.convert_byte_array(array) + .map_err(|e| format!("failed to read roaringFilter: {e}")) +} + fn options_from_arrays( env: &mut JNIEnv, keys: JObjectArray, diff --git a/python/paimon_ftindex/_ffi.py b/python/paimon_ftindex/_ffi.py index 1bb025d..af04a6c 100644 --- a/python/paimon_ftindex/_ffi.py +++ b/python/paimon_ftindex/_ffi.py @@ -39,9 +39,13 @@ def _load_library(): candidates = [ os.path.join(pkg_dir, lib_name), os.path.join(pkg_dir, "..", "..", "target", "debug", lib_name), + os.path.join(pkg_dir, "..", "..", "target", "debug", "deps", lib_name), os.path.join(pkg_dir, "..", "..", "target", "release", lib_name), + os.path.join(pkg_dir, "..", "..", "target", "release", "deps", lib_name), os.path.join(pkg_dir, "..", "..", "..", "target", "debug", lib_name), + os.path.join(pkg_dir, "..", "..", "..", "target", "debug", "deps", lib_name), os.path.join(pkg_dir, "..", "..", "..", "target", "release", lib_name), + os.path.join(pkg_dir, "..", "..", "..", "target", "release", "deps", lib_name), ] for candidate in candidates: candidate = os.path.abspath(candidate) @@ -101,6 +105,19 @@ lib.paimon_ftindex_reader_search_json.argtypes = [ ] lib.paimon_ftindex_reader_search_json.restype = c_int +lib.paimon_ftindex_reader_search_json_with_roaring_filter.argtypes = [ + c_void_p, + c_char_p, + c_size_t, + POINTER(c_uint8), + c_size_t, + POINTER(c_int64), + POINTER(c_float), + c_size_t, + POINTER(c_size_t), +] +lib.paimon_ftindex_reader_search_json_with_roaring_filter.restype = c_int + lib.paimon_ftindex_reader_free.argtypes = [c_void_p] lib.paimon_ftindex_reader_free.restype = None diff --git a/python/paimon_ftindex/reader.py b/python/paimon_ftindex/reader.py index 5cc102b..645d583 100644 --- a/python/paimon_ftindex/reader.py +++ b/python/paimon_ftindex/reader.py @@ -1,5 +1,5 @@ import ctypes -from ctypes import c_float, c_int64, c_size_t, c_void_p +from ctypes import c_float, c_int64, c_size_t, c_uint8, c_void_p from ._ffi import ( READ_AT_FN, @@ -30,7 +30,7 @@ class FullTextIndexReader: native_input = PaimonFtindexInputFile(c_void_p(0), read_at_fn) self._ptr = check_ptr(lib.paimon_ftindex_reader_open(native_input)) - def search(self, query, limit=10): + def search(self, query, limit=10, filter_bytes=None): if self._closed: raise RuntimeError("FullTextIndexReader is closed") query_json = query.to_json() if hasattr(query, "to_json") else str(query) @@ -38,8 +38,8 @@ class FullTextIndexReader: row_ids = (c_int64 * capacity)() scores = (c_float * capacity)() result_len = c_size_t() - check_status( - lib.paimon_ftindex_reader_search_json( + if filter_bytes is None: + status = lib.paimon_ftindex_reader_search_json( self._ptr, query_json.encode("utf-8"), capacity, @@ -48,7 +48,21 @@ class FullTextIndexReader: capacity, ctypes.byref(result_len), ) - ) + else: + filter_bytes = bytes(filter_bytes) + filter_buf = (c_uint8 * len(filter_bytes)).from_buffer_copy(filter_bytes) + status = lib.paimon_ftindex_reader_search_json_with_roaring_filter( + self._ptr, + query_json.encode("utf-8"), + capacity, + filter_buf, + len(filter_bytes), + row_ids, + scores, + capacity, + ctypes.byref(result_len), + ) + check_status(status) size = result_len.value return list(row_ids[:size]), list(scores[:size]) diff --git a/python/tests/test_roundtrip.py b/python/tests/test_roundtrip.py index 2d3e667..b6a98ca 100644 --- a/python/tests/test_roundtrip.py +++ b/python/tests/test_roundtrip.py @@ -1,4 +1,5 @@ from io import BytesIO +import struct from paimon_ftindex import FullTextIndexReader, FullTextIndexWriter, MatchQuery @@ -23,3 +24,59 @@ def test_python_round_trip(): assert row_ids == [1] assert scores[0] > 0 + + +def test_python_search_with_roaring_filter(): + output = BytesIO() + allowed_id = (1 << 33) + 2 + with FullTextIndexWriter() as writer: + writer.add_document(1, "Apache Paimon full text") + writer.add_document(allowed_id, "Paimon filtered row") + writer.write(output) + + filter_bytes = _roaring_treemap_bytes([allowed_id]) + with FullTextIndexReader(BytesInput(output.getvalue())) as reader: + row_ids, scores = reader.search( + MatchQuery("paimon"), limit=10, filter_bytes=filter_bytes + ) + + assert row_ids == [allowed_id] + assert scores[0] > 0 + + +def _roaring_treemap_bytes(ids): + bitmaps = {} + for value in sorted(set(ids)): + high = value >> 32 + low = value & 0xFFFFFFFF + bitmaps.setdefault(high, []).append(low) + + out = bytearray(struct.pack("<Q", len(bitmaps))) + for high, values in bitmaps.items(): + out.extend(struct.pack("<I", high)) + out.extend(_roaring_bitmap_bytes(values)) + return bytes(out) + + +def _roaring_bitmap_bytes(values): + containers = {} + for value in values: + key = value >> 16 + low = value & 0xFFFF + containers.setdefault(key, []).append(low) + + out = bytearray(struct.pack("<II", 12346, len(containers))) + descriptions = bytearray() + offsets = bytearray() + payload = bytearray() + offset = 8 + len(containers) * 8 + for key, lows in containers.items(): + lows = sorted(set(lows)) + if len(lows) > 4096: + raise ValueError("test helper only supports array containers") + descriptions.extend(struct.pack("<HH", key, len(lows) - 1)) + offsets.extend(struct.pack("<I", offset)) + values_bytes = struct.pack("<" + "H" * len(lows), *lows) + payload.extend(values_bytes) + offset += len(values_bytes) + return bytes(out + descriptions + offsets + payload)
