Skip to content

Commit

Permalink
add test.
Browse files Browse the repository at this point in the history
  • Loading branch information
Rachelint committed Feb 6, 2023
1 parent d8d234b commit 0df3d89
Showing 1 changed file with 130 additions and 0 deletions.
130 changes: 130 additions & 0 deletions table_engine/src/provider.rs
Original file line number Diff line number Diff line change
Expand Up @@ -379,3 +379,133 @@ impl fmt::Debug for ScanTable {
.finish()
}
}

#[cfg(test)]
mod test {
use std::sync::Arc;

use common_types::{
column_schema,
datum::DatumKind,
schema::{Builder, Schema},
};
use datafusion::scalar::ScalarValue;
use datafusion_expr::{col, Expr};

use super::TableProviderAdapter;
use crate::{memory::MemoryTable, table::TableId};

fn build_user_defined_primary_key_schema() -> Schema {
Builder::new()
.auto_increment_column_id(true)
.add_key_column(
column_schema::Builder::new("key1".to_string(), DatumKind::String)
.build()
.expect("should succeed build column schema"),
)
.unwrap()
.add_key_column(
column_schema::Builder::new("timestamp".to_string(), DatumKind::Timestamp)
.build()
.expect("should succeed build column schema"),
)
.unwrap()
.add_normal_column(
column_schema::Builder::new("field1".to_string(), DatumKind::String)
.is_tag(true)
.build()
.expect("should succeed build column schema"),
)
.unwrap()
.add_normal_column(
column_schema::Builder::new("field2".to_string(), DatumKind::Double)
.build()
.expect("should succeed build column schema"),
)
.unwrap()
.build()
.unwrap()
}

fn build_default_primary_key_schema() -> Schema {
Builder::new()
.auto_increment_column_id(true)
.add_key_column(
column_schema::Builder::new("timestamp".to_string(), DatumKind::Timestamp)
.build()
.expect("should succeed build column schema"),
)
.unwrap()
.add_key_column(
column_schema::Builder::new("tsid".to_string(), DatumKind::UInt64)
.build()
.expect("should succeed build column schema"),
)
.unwrap()
.add_normal_column(
column_schema::Builder::new("field1".to_string(), DatumKind::String)
.is_tag(true)
.build()
.expect("should succeed build column schema"),
)
.unwrap()
.add_normal_column(
column_schema::Builder::new("field2".to_string(), DatumKind::Double)
.build()
.expect("should succeed build column schema"),
)
.unwrap()
.build()
.unwrap()
}

fn build_filters() -> Vec<Expr> {
let filter1 = col("timestamp").lt(Expr::Literal(ScalarValue::UInt64(Some(10086))));
let filter2 = col("key1").eq(Expr::Literal(ScalarValue::Utf8(Some("10086".to_string()))));
let filter3 = col("field1").eq(Expr::Literal(ScalarValue::Utf8(Some("10087".to_string()))));
let filter4 = col("field2").eq(Expr::Literal(ScalarValue::Float64(Some(10088.0))));

vec![filter1, filter2, filter3, filter4]
}

#[test]
pub fn test_push_down_in_user_defined_primary_key_case() {
let test_filters = build_filters();
let user_defined_pk_schema = build_user_defined_primary_key_schema();

let table = MemoryTable::new(
"test_table".to_string(),
TableId::new(0),
user_defined_pk_schema,
"memory".to_string(),
);
let provider = TableProviderAdapter::new(Arc::new(table), 1);
let predicate = provider.check_and_build_predicate_from_filters(&test_filters);

let expected_filters = vec![test_filters[0].clone(), test_filters[1].clone()];
assert_eq!(predicate.exprs(), &expected_filters);
}

#[test]
pub fn test_push_down_in_default_primary_key_case() {
let test_filters = build_filters();
let test_filters = vec![
test_filters[0].clone(),
test_filters[2].clone(),
test_filters[3].clone(),
];
let default_pk_schema = build_default_primary_key_schema();

let table = MemoryTable::new(
"test_table".to_string(),
TableId::new(0),
default_pk_schema,
"memory".to_string(),
);
let provider = TableProviderAdapter::new(Arc::new(table), 1);
let predicate = provider.check_and_build_predicate_from_filters(&test_filters);

let expected_filters = vec![test_filters[0].clone(), test_filters[1].clone()];
assert_eq!(predicate.exprs(), &expected_filters);
}
}

0 comments on commit 0df3d89

Please sign in to comment.