Skip to content

Commit

Permalink
test: indexer test
Browse files Browse the repository at this point in the history 
Signed-off-by: Zhenchi <zhongzc_arch@outlook.com>
  • Loading branch information
@zhongzc
zhongzc committed Jul 6, 2024
1 parent 21b3831 commit 3207811
Show file tree
Hide file tree
Showing 4 changed files with 130 additions and 69 deletions.
181 changes: 126 additions & 55 deletions src/mito2/src/sst/index.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -296,7 +296,7 @@ mod tests {

use api::v1::SemanticType;
use datatypes::data_type::ConcreteDataType;
use datatypes::schema::ColumnSchema;
use datatypes::schema::{ColumnSchema, FulltextOptions, FULLTEXT_KEY};
use object_store::services::Memory;
use object_store::ObjectStore;
use puffin_manager::PuffinManagerFactory;
Expand All @@ -305,12 +305,26 @@ mod tests {
use super::*;
use crate::config::{FulltextIndexConfig, Mode};

fn mock_region_metadata() -> RegionMetadataRef {
struct MetaConfig {
with_tag: bool,
with_fulltext: bool,
}

fn mock_region_metadata(
MetaConfig {
with_tag,
with_fulltext,
}: MetaConfig,
) -> RegionMetadataRef {
let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 2));
builder
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new("a", ConcreteDataType::int64_datatype(), false),
semantic_type: SemanticType::Tag,
semantic_type: if with_tag {
SemanticType::Tag
} else {
SemanticType::Field
},
column_id: 1,
})
.push_column_metadata(ColumnMetadata {
Expand All @@ -326,34 +340,33 @@ mod tests {
),
semantic_type: SemanticType::Timestamp,
column_id: 3,
})
.primary_key(vec![1]);
});

Arc::new(builder.build().unwrap())
}
if with_tag {
builder.primary_key(vec![1]);
}

fn no_tag_region_metadata() -> RegionMetadataRef {
let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 2));
builder
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new("a", ConcreteDataType::int64_datatype(), false),
semantic_type: SemanticType::Field,
column_id: 1,
if with_fulltext {
let opts = serde_json::to_string(&FulltextOptions {
enable: true,
..Default::default()
})
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new("b", ConcreteDataType::float64_datatype(), false),
.unwrap();

let mut column_schema =
ColumnSchema::new("text", ConcreteDataType::string_datatype(), true);
column_schema
.mut_metadata()
.insert(FULLTEXT_KEY.to_string(), opts);

let column = ColumnMetadata {
column_schema,
semantic_type: SemanticType::Field,
column_id: 2,
})
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new(
"c",
ConcreteDataType::timestamp_millisecond_datatype(),
false,
),
semantic_type: SemanticType::Timestamp,
column_id: 3,
});
column_id: 4,
};

builder.push_column_metadata(column);
}

Arc::new(builder.build().unwrap())
}
Expand All @@ -362,25 +375,28 @@ mod tests {
ObjectStore::new(Memory::default()).unwrap().finish()
}

fn mock_intm_mgr() -> IntermediateManager {
IntermediateManager::new(mock_object_store())
async fn mock_intm_mgr(path: impl AsRef<str>) -> IntermediateManager {
IntermediateManager::init_fs(path).await.unwrap()
}

#[tokio::test]
async fn test_build_indexer_basic() {
let (_d, factory) =
let (dir, factory) =
PuffinManagerFactory::new_for_test_async("test_build_indexer_basic_").await;
let store = mock_object_store();
let puffin_manager = factory.build(store);
let metadata = mock_region_metadata();
let intm_maanger = mock_intm_mgr(dir.path().to_string_lossy()).await;

let metadata = mock_region_metadata(MetaConfig {
with_tag: true,
with_fulltext: true,
});
let indexer = IndexerBuilder {
op_type: OperationType::Flush,
file_id: FileId::random(),
file_path: "test".to_string(),
metadata: &metadata,
row_group_size: 1024,
puffin_manager,
intermediate_manager: mock_intm_mgr(),
puffin_manager: factory.build(mock_object_store()),
intermediate_manager: intm_maanger,
index_options: IndexOptions::default(),
inverted_index_config: InvertedIndexConfig::default(),
fulltext_index_config: FulltextIndexConfig::default(),
Expand All @@ -389,23 +405,27 @@ mod tests {
.await;

assert!(indexer.inverted_indexer.is_some());
assert!(indexer.fulltext_indexer.is_some());
}

#[tokio::test]
async fn test_build_indexer_disable_create() {
let (_d, factory) =
let (dir, factory) =
PuffinManagerFactory::new_for_test_async("test_build_indexer_disable_create_").await;
let store = mock_object_store();
let puffin_manager = factory.build(store);
let metadata = mock_region_metadata();
let intm_maanger = mock_intm_mgr(dir.path().to_string_lossy()).await;

let metadata = mock_region_metadata(MetaConfig {
with_tag: true,
with_fulltext: true,
});
let indexer = IndexerBuilder {
op_type: OperationType::Flush,
file_id: FileId::random(),
file_path: "test".to_string(),
metadata: &metadata,
row_group_size: 1024,
puffin_manager,
intermediate_manager: mock_intm_mgr(),
puffin_manager: factory.build(mock_object_store()),
intermediate_manager: intm_maanger.clone(),
index_options: IndexOptions::default(),
inverted_index_config: InvertedIndexConfig {
create_on_flush: Mode::Disable,
Expand All @@ -417,23 +437,48 @@ mod tests {
.await;

assert!(indexer.inverted_indexer.is_none());
assert!(indexer.fulltext_indexer.is_some());

let indexer = IndexerBuilder {
op_type: OperationType::Compact,
file_id: FileId::random(),
file_path: "test".to_string(),
metadata: &metadata,
row_group_size: 1024,
puffin_manager: factory.build(mock_object_store()),
intermediate_manager: intm_maanger,
index_options: IndexOptions::default(),
inverted_index_config: InvertedIndexConfig::default(),
fulltext_index_config: FulltextIndexConfig {
create_on_compaction: Mode::Disable,
..Default::default()
},
}
.build()
.await;

assert!(indexer.inverted_indexer.is_some());
assert!(indexer.fulltext_indexer.is_none());
}

#[tokio::test]
async fn test_build_indexer_no_tag() {
let (_d, factory) =
PuffinManagerFactory::new_for_test_async("test_build_indexer_no_tag_").await;
let store = mock_object_store();
let puffin_manager = factory.build(store);
let metadata = no_tag_region_metadata();
async fn test_build_indexer_no_required() {
let (dir, factory) =
PuffinManagerFactory::new_for_test_async("test_build_indexer_no_required_").await;
let intm_maanger = mock_intm_mgr(dir.path().to_string_lossy()).await;

let metadata = mock_region_metadata(MetaConfig {
with_tag: false,
with_fulltext: true,
});
let indexer = IndexerBuilder {
op_type: OperationType::Flush,
file_id: FileId::random(),
file_path: "test".to_string(),
metadata: &metadata,
row_group_size: 1024,
puffin_manager,
intermediate_manager: mock_intm_mgr(),
puffin_manager: factory.build(mock_object_store()),
intermediate_manager: intm_maanger.clone(),
index_options: IndexOptions::default(),
inverted_index_config: InvertedIndexConfig::default(),
fulltext_index_config: FulltextIndexConfig::default(),
Expand All @@ -442,23 +487,49 @@ mod tests {
.await;

assert!(indexer.inverted_indexer.is_none());
assert!(indexer.fulltext_indexer.is_some());

let metadata = mock_region_metadata(MetaConfig {
with_tag: true,
with_fulltext: false,
});
let indexer = IndexerBuilder {
op_type: OperationType::Flush,
file_id: FileId::random(),
file_path: "test".to_string(),
metadata: &metadata,
row_group_size: 1024,
puffin_manager: factory.build(mock_object_store()),
intermediate_manager: intm_maanger,
index_options: IndexOptions::default(),
inverted_index_config: InvertedIndexConfig::default(),
fulltext_index_config: FulltextIndexConfig::default(),
}
.build()
.await;

assert!(indexer.inverted_indexer.is_some());
assert!(indexer.fulltext_indexer.is_none());
}

#[tokio::test]
async fn test_build_indexer_zero_row_group() {
let (_d, factory) =
let (dir, factory) =
PuffinManagerFactory::new_for_test_async("test_build_indexer_zero_row_group_").await;
let store = mock_object_store();
let puffin_manager = factory.build(store);
let metadata = mock_region_metadata();
let intm_maanger = mock_intm_mgr(dir.path().to_string_lossy()).await;

let metadata = mock_region_metadata(MetaConfig {
with_tag: true,
with_fulltext: true,
});
let indexer = IndexerBuilder {
op_type: OperationType::Flush,
file_id: FileId::random(),
file_path: "test".to_string(),
metadata: &metadata,
row_group_size: 0,
puffin_manager,
intermediate_manager: mock_intm_mgr(),
puffin_manager: factory.build(mock_object_store()),
intermediate_manager: intm_maanger,
index_options: IndexOptions::default(),
inverted_index_config: InvertedIndexConfig::default(),
fulltext_index_config: FulltextIndexConfig::default(),
Expand Down
4 changes: 1 addition & 3 deletions src/mito2/src/sst/index/indexer/finish.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,9 +51,7 @@ impl Indexer {
}

async fn build_puffin_writer(&mut self) -> Option<SstPuffinWriter> {
let Some(puffin_manager) = self.puffin_manager.take() else {
return None;
};
let puffin_manager = self.puffin_manager.take()?;

let err = match puffin_manager.writer(&self.file_path).await {
Ok(writer) => return Some(writer),
Expand Down
8 changes: 0 additions & 8 deletions src/mito2/src/sst/index/intermediate.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -77,14 +77,6 @@ impl IntermediateManager {
.join(sst_file_id.to_string())
.join(format!("fulltext-{column_id}-{uuid}"))
}

#[cfg(test)]
pub(crate) fn new(store: object_store::ObjectStore) -> Self {
Self {
base_dir: PathBuf::new(),
store: InstrumentedStore::new(store),
}
}
}

/// `IntermediateLocation` produces paths for intermediate files
Expand Down
6 changes: 3 additions & 3 deletions src/mito2/src/sst/index/inverted_index/creator.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -313,8 +313,8 @@ mod tests {
ObjectStore::new(Memory::default()).unwrap().finish()
}

fn mock_intm_mgr() -> IntermediateManager {
IntermediateManager::new(mock_object_store())
async fn new_intm_mgr(path: impl AsRef<str>) -> IntermediateManager {
IntermediateManager::init_fs(path).await.unwrap()
}

fn mock_region_metadata() -> RegionMetadataRef {
Expand Down Expand Up @@ -387,7 +387,7 @@ mod tests {
let file_path = location::index_file_path(&region_dir, sst_file_id);
let object_store = mock_object_store();
let region_metadata = mock_region_metadata();
let intm_mgr = mock_intm_mgr();
let intm_mgr = new_intm_mgr(d.path().to_string_lossy()).await;
let memory_threshold = None;
let segment_row_count = 2;

Expand Down

0 comments on commit 3207811

Please sign in to comment.