leaves12138 commented on code in PR #1:
URL: https://github.com/apache/paimon-full-text/pull/1#discussion_r3523284829
##########
core/src/index.rs:
##########
@@ -0,0 +1,384 @@
+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::tokenizer::{TokenizerConfig, TokenizerKind};
+use std::fs;
+use std::path::Path;
+use tantivy::collector::TopDocs;
+use tantivy::directory::{Directory, RamDirectory};
+use tantivy::query::{BooleanQuery, BoostQuery, Occur, Query, QueryParser};
+use tantivy::schema::{IndexRecordOption, NumericOptions, Schema,
TextFieldIndexing, TextOptions};
+use tantivy::tokenizer::{
+ AsciiFoldingFilter, LowerCaser, NgramTokenizer, RawTokenizer,
RemoveLongFilter,
+ SimpleTokenizer, TextAnalyzer, WhitespaceTokenizer,
+};
+use tantivy::{Index, TantivyDocument};
+use tempfile::TempDir;
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct FullTextSearchResult {
+ pub row_ids: Vec<i64>,
+ pub scores: Vec<f32>,
+}
+
+pub struct FullTextIndexWriter {
+ config: FullTextIndexConfig,
+ documents: Vec<(i64, String)>,
+}
+
+impl FullTextIndexWriter {
+ pub fn new(config: FullTextIndexConfig) -> Result<Self> {
+ config.tokenizer.validate()?;
+ Ok(Self {
+ config,
+ documents: Vec::new(),
+ })
+ }
+
+ pub fn add_document(&mut self, row_id: i64, text: impl Into<String>) ->
Result<()> {
+ if row_id < 0 {
+ return Err(FtIndexError::InvalidStorage(format!(
+ "row id must be non-negative, got {row_id}"
+ )));
+ }
+ self.documents.push((row_id, text.into()));
+ Ok(())
+ }
+
+ pub fn write<W: SeekWrite>(&mut self, output: &mut W) -> Result<()> {
+ let temp_dir = TempDir::new()?;
+ let schema = build_schema(&self.config);
+ let mut index = Index::create_in_dir(temp_dir.path(), schema.clone())?;
+ register_tokenizer(&mut index, &self.config.tokenizer)?;
+ let row_id_field = schema
+ .get_field(&self.config.row_id_field)
+ .map_err(|_| FtIndexError::InvalidStorage("missing row_id
field".to_string()))?;
+ let text_field = schema
+ .get_field(&self.config.text_field)
+ .map_err(|_| FtIndexError::InvalidStorage("missing text
field".to_string()))?;
+
+ {
+ let mut index_writer = index.writer(50_000_000)?;
+ for (row_id, text) in &self.documents {
+ let mut doc = TantivyDocument::new();
+ doc.add_u64(row_id_field, *row_id as u64);
+ doc.add_text(text_field, text);
+ index_writer.add_document(doc)?;
+ }
+ index_writer.commit()?;
+ }
+
+ let files = collect_index_files(temp_dir.path())?;
+ let mut offset = 0u64;
+ let mut entries = Vec::with_capacity(files.len());
+ for (name, data) in &files {
+ entries.push(ArchiveFileEntry {
+ name: name.clone(),
+ offset,
+ length: data.len() as u64,
+ });
+ offset += data.len() as u64;
+ }
+
+ let header = IndexHeader {
+ metadata: FullTextIndexMetadata {
+ format_version: crate::storage::FORMAT_VERSION,
+ config: self.config.clone(),
+ document_count: self.documents.len() as u64,
+ tantivy_version: tantivy::version().to_string(),
+ },
+ files: entries,
+ };
+ write_envelope(output, &header, &files)
+ }
+}
+
+pub struct FullTextIndexReader<R> {
+ _input: R,
+ index: Index,
+ metadata: FullTextIndexMetadata,
+}
+
+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)?;
+ }
+ let mut index = Index::open(directory)?;
+ register_tokenizer(&mut index, &header.metadata.config.tokenizer)?;
+ Ok(Self {
+ _input: input,
+ index,
+ metadata: header.metadata,
+ })
+ }
+
+ pub fn optimize_for_search(&mut self) -> Result<()> {
+ Ok(())
+ }
+
+ pub fn metadata(&self) -> &FullTextIndexMetadata {
+ &self.metadata
+ }
+
+ pub fn search(&mut self, query: FullTextQuery, limit: usize) ->
Result<FullTextSearchResult> {
+ if limit == 0 {
+ return Err(FtIndexError::InvalidQuery(
+ "search limit must be positive".to_string(),
+ ));
+ }
+ 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 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 {
+ let segment_reader =
searcher.segment_reader(doc_address.segment_ord);
+ let row_id_column = segment_reader
+ .fast_fields()
+ .u64(&self.metadata.config.row_id_field)?
+ .first_or_default_col(0);
+ row_ids.push(row_id_column.get_val(doc_address.doc_id) as i64);
+ scores.push(score);
+ }
+ Ok(FullTextSearchResult { row_ids, scores })
+ }
+}
+
+fn build_schema(config: &FullTextIndexConfig) -> Schema {
+ let mut builder = Schema::builder();
+ builder.add_u64_field(
+ &config.row_id_field,
+ NumericOptions::default()
+ .set_fast()
+ .set_stored()
+ .set_indexed(),
+ );
+ let index_option = if config.tokenizer.with_position {
+ IndexRecordOption::WithFreqsAndPositions
+ } else {
+ IndexRecordOption::WithFreqs
+ };
+ let tokenizer_name = tokenizer_name(&config.tokenizer);
+ let indexing = TextFieldIndexing::default()
+ .set_tokenizer(tokenizer_name)
+ .set_index_option(index_option);
+ builder.add_text_field(
+ &config.text_field,
+ TextOptions::default().set_indexing_options(indexing),
+ );
+ builder.build()
+}
+
+fn tokenizer_name(config: &TokenizerConfig) -> &'static str {
+ match config.tokenizer {
+ TokenizerKind::Default if !needs_custom_default(config) => "default",
+ TokenizerKind::Default | TokenizerKind::Simple => "paimon_custom",
+ TokenizerKind::Whitespace => "paimon_custom",
+ TokenizerKind::Raw => "paimon_custom",
+ TokenizerKind::Ngram => "paimon_ngram",
+ TokenizerKind::Jieba => "paimon_jieba",
+ }
+}
+
+fn needs_custom_default(config: &TokenizerConfig) -> bool {
+ !config.lower_case
+ || config.max_token_length != 40
+ || config.ascii_folding
+ || config.stem
+ || config.remove_stop_words
+ || !config.stop_words.is_empty()
+}
+
+fn register_tokenizer(index: &mut Index, config: &TokenizerConfig) ->
Result<()> {
+ match config.tokenizer {
+ TokenizerKind::Default if !needs_custom_default(config) => Ok(()),
+ TokenizerKind::Jieba => Err(FtIndexError::InvalidOption {
+ key: "tokenizer".to_string(),
+ message: "jieba tokenizer is not enabled in this first
implementation".to_string(),
+ }),
+ _ => {
+ let analyzer = build_text_analyzer(config)?;
+ index
+ .tokenizers()
+ .register(tokenizer_name(config), analyzer);
+ Ok(())
+ }
+ }
+}
+
+fn build_text_analyzer(config: &TokenizerConfig) -> Result<TextAnalyzer> {
+ if config.stem || config.remove_stop_words ||
!config.stop_words.is_empty() {
+ return Err(FtIndexError::InvalidOption {
+ key: "tokenizer filters".to_string(),
+ message: "stemming and stop-word filters are not enabled in this
first implementation"
+ .to_string(),
+ });
+ }
+ let mut builder = match config.tokenizer {
+ TokenizerKind::Default | TokenizerKind::Simple => {
+ TextAnalyzer::builder(SimpleTokenizer::default()).dynamic()
+ }
+ TokenizerKind::Whitespace => {
+ TextAnalyzer::builder(WhitespaceTokenizer::default()).dynamic()
+ }
+ TokenizerKind::Raw =>
TextAnalyzer::builder(RawTokenizer::default()).dynamic(),
+ TokenizerKind::Ngram => {
+ let tokenizer = NgramTokenizer::new(
+ config.ngram_min_gram,
+ config.ngram_max_gram,
+ config.ngram_prefix_only,
+ )
+ .map_err(|e| FtIndexError::InvalidOption {
+ key: "ngram".to_string(),
+ message: e.to_string(),
+ })?;
+ TextAnalyzer::builder(tokenizer).dynamic()
+ }
+ TokenizerKind::Jieba => {
+ return Err(FtIndexError::InvalidOption {
+ key: "tokenizer".to_string(),
+ message: "jieba tokenizer is not enabled in this first
implementation".to_string(),
+ })
+ }
+ };
+ builder =
builder.filter_dynamic(RemoveLongFilter::limit(config.max_token_length));
+ if config.lower_case {
+ builder = builder.filter_dynamic(LowerCaser);
+ }
+ if config.ascii_folding {
+ builder = builder.filter_dynamic(AsciiFoldingFilter);
+ }
+ Ok(builder.build())
+}
+
+fn collect_index_files(path: &Path) -> Result<Vec<(String, Vec<u8>)>> {
+ let mut paths = Vec::new();
+ for entry in fs::read_dir(path)? {
+ let entry = entry?;
+ if entry.file_type()?.is_file() {
+ paths.push(entry.path());
+ }
+ }
+ paths.sort();
+ let mut files = Vec::with_capacity(paths.len());
+ for path in paths {
+ let name = path
+ .file_name()
+ .and_then(|name| name.to_str())
+ .ok_or_else(|| FtIndexError::InvalidStorage("non-utf8 file
name".to_string()))?
+ .to_string();
+ if name.ends_with(".lock") {
+ continue;
+ }
+ files.push((name, fs::read(path)?));
+ }
+ Ok(files)
+}
+
+fn build_query(
+ index: &Index,
+ config: &FullTextIndexConfig,
+ query: &FullTextQuery,
+) -> Result<Box<dyn Query>> {
+ match query {
+ FullTextQuery::Match {
+ column,
+ terms,
+ operator,
+ boost,
+ } => {
+ validate_column(config, column)?;
+ let text_field = index
+ .schema()
+ .get_field(&config.text_field)
+ .map_err(|_| FtIndexError::InvalidQuery("missing text
field".to_string()))?;
+ let mut parser = QueryParser::for_index(index, vec![text_field]);
+ if *operator == MatchOperator::And {
+ parser.set_conjunction_by_default();
+ }
+ let parsed = parser
+ .parse_query(terms)
+ .map_err(|e| FtIndexError::InvalidQuery(e.to_string()))?;
+ if (*boost - 1.0).abs() > f32::EPSILON {
+ Ok(Box::new(BoostQuery::new(parsed, *boost)))
+ } else {
+ Ok(parsed)
+ }
+ }
+ FullTextQuery::MatchPhrase {
+ column,
+ terms,
+ slop,
+ } => {
+ validate_column(config, column)?;
+ if !config.tokenizer.with_position {
+ return Err(FtIndexError::InvalidQuery(
+ "phrase query requires positions".to_string(),
+ ));
+ }
+ let text_field = index
+ .schema()
+ .get_field(&config.text_field)
+ .map_err(|_| FtIndexError::InvalidQuery("missing text
field".to_string()))?;
+ let parser = QueryParser::for_index(index, vec![text_field]);
+ let escaped = terms.replace('\\', "\\\\").replace('"', "\\\"");
+ let query_text = if *slop == 0 {
+ format!("\"{escaped}\"")
+ } else {
+ format!("\"{escaped}\"~{slop}")
+ };
+ parser
+ .parse_query(&query_text)
+ .map_err(|e| FtIndexError::InvalidQuery(e.to_string()))
+ }
+ FullTextQuery::Boolean { queries } => {
+ if queries.is_empty() {
+ return Err(FtIndexError::InvalidQuery(
+ "boolean query must contain at least one
clause".to_string(),
+ ));
+ }
+ let mut children = Vec::with_capacity(queries.len());
+ for (occur, child) in queries {
+ let occur = match occur {
+ BooleanOccur::Should => Occur::Should,
+ BooleanOccur::Must => Occur::Must,
+ BooleanOccur::MustNot => Occur::MustNot,
+ };
+ children.push((occur, build_query(index, config, child)?));
+ }
+ Ok(Box::new(BooleanQuery::new(children)))
+ }
+ FullTextQuery::Boost {
+ positive,
+ negative,
+ negative_boost,
+ } => {
+ let boosted_negative = Box::new(BoostQuery::new(
+ build_query(index, config, negative)?,
+ *negative_boost,
+ ));
+ Ok(Box::new(BooleanQuery::new(vec![
Review Comment:
This makes the `Boost` query broaden the result set to documents that match
only the negative query, because both children are `Should` clauses. That seems
contrary to the `positive`/`negative` naming and the design doc wording of
score demotion/promotion by composing child scores. I reproduced it with two
docs (`"apache paimon"`, `"tantivy only"`) and `Boost { positive:
match("paimon"), negative: match("tantivy"), negative_boost: 0.5 }`; the result
row ids are `[1, 2]`, so the negative-only document is returned. Please either
make the positive side required (and only use the negative side to adjust
scores for matching positives), or explicitly redefine the semantics and add a
regression test for negative-only matches.
##########
jni/src/lib.rs:
##########
@@ -0,0 +1,334 @@
+use jni::objects::{GlobalRef, JByteArray, JClass, JObject, JObjectArray,
JString, JValue};
+use jni::sys::{jint, jlong, jobject};
+use jni::{JNIEnv, JavaVM};
+use paimon_ftindex_core::io::{ReadRequest, SeekRead, SeekWrite};
+use paimon_ftindex_core::{
+ FullTextIndexConfig, FullTextIndexReader, FullTextIndexWriter,
FullTextQuery, TokenizerConfig,
+};
+use std::collections::HashMap;
+use std::io;
+use std::ptr;
+
+struct JavaOutput {
+ jvm: JavaVM,
+ output: GlobalRef,
+}
+
+unsafe impl Send for JavaOutput {}
+
+impl SeekWrite for JavaOutput {
+ fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
+ let mut env = self
+ .jvm
+ .attach_current_thread()
+ .map_err(|e| io::Error::other(format!("JNI attach failed: {e}")))?;
+ let array = env
+ .new_byte_array(buf.len() as i32)
+ .map_err(|e| io::Error::other(format!("new_byte_array failed:
{e}")))?;
+ let signed: Vec<i8> = buf.iter().map(|b| *b as i8).collect();
+ env.set_byte_array_region(&array, 0, &signed)
+ .map_err(|e| io::Error::other(format!("set_byte_array_region
failed: {e}")))?;
+ let array_obj = JObject::from(array);
+ env.call_method(
+ self.output.as_obj(),
+ "write",
+ "([BII)V",
+ &[
+ JValue::Object(&array_obj),
+ JValue::Int(0),
+ JValue::Int(buf.len() as i32),
+ ],
+ )
+ .map_err(|e| io::Error::other(format!("Java output write failed:
{e}")))?;
+ Ok(())
+ }
+
+ fn flush(&mut self) -> io::Result<()> {
+ let mut env = self
+ .jvm
+ .attach_current_thread()
+ .map_err(|e| io::Error::other(format!("JNI attach failed: {e}")))?;
+ env.call_method(self.output.as_obj(), "flush", "()V", &[])
+ .map_err(|e| io::Error::other(format!("Java output flush failed:
{e}")))?;
+ Ok(())
+ }
+}
+
+struct JavaInput {
+ jvm: JavaVM,
+ input: GlobalRef,
+}
+
+unsafe impl Send for JavaInput {}
+
+impl SeekRead for JavaInput {
+ fn pread(&mut self, ranges: &mut [ReadRequest<'_>]) -> io::Result<()> {
+ 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
+ .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);
+ let mut signed = vec![0i8; range.buf.len()];
+ env.get_byte_array_region(&array, 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(())
+ }
+}
+
+struct WriterHandle {
+ inner: FullTextIndexWriter,
+}
+
+struct ReaderHandle {
+ inner: FullTextIndexReader<JavaInput>,
+}
+
+#[no_mangle]
+pub extern "system" fn
Java_org_apache_paimon_index_fulltext_FullTextNative_createWriter(
+ mut env: JNIEnv,
+ _class: JClass,
+ keys: JObjectArray,
+ values: JObjectArray,
+) -> jlong {
+ match create_writer(&mut env, keys, values) {
+ Ok(ptr) => ptr,
+ Err(e) => throw_and_return(&mut env, &e, 0),
+ }
+}
+
+#[no_mangle]
+pub extern "system" fn
Java_org_apache_paimon_index_fulltext_FullTextNative_addDocument(
+ mut env: JNIEnv,
+ _class: JClass,
+ writer_ptr: jlong,
+ row_id: jlong,
+ text: JString,
+) {
+ if let Err(e) = add_document(&mut env, writer_ptr, row_id, text) {
+ throw(&mut env, &e);
+ }
+}
+
+#[no_mangle]
+pub extern "system" fn
Java_org_apache_paimon_index_fulltext_FullTextNative_writeIndex(
+ mut env: JNIEnv,
+ _class: JClass,
+ writer_ptr: jlong,
+ output: JObject,
+) {
+ if let Err(e) = write_index(&mut env, writer_ptr, output) {
+ throw(&mut env, &e);
+ }
+}
+
+#[no_mangle]
+pub extern "system" fn
Java_org_apache_paimon_index_fulltext_FullTextNative_freeWriter(
+ _env: JNIEnv,
+ _class: JClass,
+ writer_ptr: jlong,
+) {
+ if writer_ptr != 0 {
+ unsafe {
+ drop(Box::from_raw(writer_ptr as *mut WriterHandle));
+ }
+ }
+}
+
+#[no_mangle]
+pub extern "system" fn
Java_org_apache_paimon_index_fulltext_FullTextNative_openReader(
+ mut env: JNIEnv,
+ _class: JClass,
+ input: JObject,
+) -> jlong {
+ match open_reader(&mut env, input) {
+ Ok(ptr) => ptr,
+ Err(e) => throw_and_return(&mut env, &e, 0),
+ }
+}
+
+#[no_mangle]
+pub extern "system" fn
Java_org_apache_paimon_index_fulltext_FullTextNative_searchJson(
+ mut env: JNIEnv,
+ _class: JClass,
+ reader_ptr: jlong,
+ query_json: JString,
+ limit: jint,
+) -> jobject {
+ match search_json(&mut env, reader_ptr, query_json, limit) {
+ 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_freeReader(
+ _env: JNIEnv,
+ _class: JClass,
+ reader_ptr: jlong,
+) {
+ if reader_ptr != 0 {
+ unsafe {
+ drop(Box::from_raw(reader_ptr as *mut ReaderHandle));
+ }
+ }
+}
+
+fn create_writer(
+ env: &mut JNIEnv,
+ keys: JObjectArray,
+ values: JObjectArray,
+) -> Result<jlong, String> {
+ let options = options_from_arrays(env, keys, values)?;
+ let tokenizer = TokenizerConfig::from_options(&options).map_err(|e|
e.to_string())?;
+ let config = FullTextIndexConfig::new().tokenizer(tokenizer);
+ let writer = FullTextIndexWriter::new(config).map_err(|e| e.to_string())?;
+ Ok(Box::into_raw(Box::new(WriterHandle { inner: writer })) as jlong)
+}
+
+fn add_document(
+ env: &mut JNIEnv,
+ writer_ptr: jlong,
+ row_id: jlong,
+ text: JString,
+) -> Result<(), String> {
+ let writer = handle_mut::<WriterHandle>(writer_ptr, "writer")?;
+ let text: String = env
+ .get_string(&text)
+ .map_err(|e| format!("failed to read text: {e}"))?
+ .into();
+ writer
+ .inner
+ .add_document(row_id, text)
+ .map_err(|e| e.to_string())
+}
+
+fn write_index(env: &mut JNIEnv, writer_ptr: jlong, output: JObject) ->
Result<(), String> {
+ let writer = handle_mut::<WriterHandle>(writer_ptr, "writer")?;
+ let jvm = env.get_java_vm().map_err(|e| e.to_string())?;
+ let output = env.new_global_ref(output).map_err(|e| e.to_string())?;
+ let mut output = JavaOutput { jvm, output };
+ writer.inner.write(&mut output).map_err(|e| e.to_string())
+}
+
+fn open_reader(env: &mut JNIEnv, input: JObject) -> Result<jlong, String> {
+ let jvm = env.get_java_vm().map_err(|e| e.to_string())?;
+ let input = env.new_global_ref(input).map_err(|e| e.to_string())?;
+ let input = JavaInput { jvm, input };
+ let reader = FullTextIndexReader::open(input).map_err(|e| e.to_string())?;
+ Ok(Box::into_raw(Box::new(ReaderHandle { inner: reader })) as jlong)
+}
+
+fn search_json(
+ env: &mut JNIEnv,
+ reader_ptr: jlong,
+ query_json: JString,
+ limit: jint,
+) -> Result<jobject, String> {
+ let reader = handle_mut::<ReaderHandle>(reader_ptr, "reader")?;
+ let query_json: String = env
+ .get_string(&query_json)
+ .map_err(|e| format!("failed to read query json: {e}"))?
+ .into();
+ let query = FullTextQuery::from_json(&query_json).map_err(|e|
e.to_string())?;
+ let result = reader
+ .inner
+ .search(query, limit as usize)
Review Comment:
Please validate `limit > 0` before casting the Java `int` to `usize`. Today
`search(query, -1)` reaches JNI as a negative `jint`, wraps to a huge `usize`,
and is then passed to `TopDocs::with_limit`; that can turn a simple invalid API
call into excessive allocation/work instead of a clean exception. The Java-side
Paimon full-text API rejects non-positive limits, so this bridge should keep
the same boundary.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
