leaves12138 commented on code in PR #1:
URL: https://github.com/apache/paimon-full-text/pull/1#discussion_r3523306632
##########
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![
+ (Occur::Must, build_query(index, config, positive)?),
+ (Occur::Should, boosted_negative),
Review Comment:
The fix now makes the positive side required, which solves the negative-only
result issue, but the negative clause still contributes a positive score
through `Occur::Should + BoostQuery`. With `negative_boost = 0.5`, a document
matching both positive and negative gets `positive_score + 0.5 *
negative_score`, so it can rank above a positive-only document instead of being
demoted. I reproduced this with docs `"paimon good"` and `"paimon bad"`: `Boost
{ positive: match("paimon"), negative: match("bad"), negative_boost: 0.5 }`
returns row ids `[2, 1]` with scores `[0.5288952, 0.1823216]`. If the intended
API is negative-score demotion, this needs a different implementation (for
example a custom query/scorer or result-level score adjustment) plus a
regression test for ranking. If the intended API is just optional score
composition/promotion, then the `negative`/`negative_boost` naming and design
doc should be changed to avoid promising demotion.
--
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]
