feat: prevent from pushing down filter of non primary key or ts key

table_engine/src/provider.rs

@@ -4,6 +4,7 @@
 
 use std::{
     any::Any,
+    collections::HashSet,
     fmt,
     sync::{Arc, Mutex},
     time::{Duration, Instant},
@@ -26,6 +27,7 @@ use datafusion::{
     scalar::ScalarValue,
 };
 use datafusion_expr::{TableSource, TableType};
+use df_operator::visitor;
 use log::debug;
 
 use crate::{
@@ -100,7 +102,7 @@ impl TableProviderAdapter {
             self.read_parallelism
         };
 
-        let predicate = self.predicate_from_filters(filters);
+        let predicate = self.check_and_build_predicate_from_filters(filters);
         let mut scan_table = ScanTable {
             projected_schema: ProjectedSchema::new(self.read_schema.clone(), projection.clone())
                 .map_err(|e| {
@@ -122,12 +124,61 @@ impl TableProviderAdapter {
         Ok(Arc::new(scan_table))
     }
 
-    fn predicate_from_filters(&self, filters: &[Expr]) -> PredicateRef {
+    fn check_and_build_predicate_from_filters(&self, filters: &[Expr]) -> PredicateRef {
+        // 1.Only filters of primary key or timestamp key can be pushed down.
+        // Build key set.
+        let mut key_set = HashSet::new();
+        // Timestamp key should be always picked.
+        key_set.insert(self.read_schema.timestamp_name());
+        if self.read_schema.tsid_column().is_some() {
+            // When tsid exists, that means default primary key (tsid, timestamp) is used.
+            // So, all filters of tag columns(tsid is the hash result of all tags) can be pushed down.
+            for column in self.read_schema.columns() {
+                if column.is_tag {
+                    key_set.insert(&column.name);
+                }
+            }
+        } else {
+            // When tsid does not exist, that means user defined primary key is used.
+            // So, only filters of primary key can be pushed down.
+            for primary_idx in self.read_schema.primary_key_indexes() {
+                let primary_column = self.read_schema.column(*primary_idx);
+                key_set.insert(&primary_column.name);
+            }
+        }
+
+        let push_down_filters = filters
+            .iter()
+            .filter_map(|filter| {
+                if self.should_push_down(filter, &key_set) {
+                    Some(filter.clone())
+                } else {
+                    None
+                }
+            })
+            .collect::<Vec<_>>();
+
+        // 2.Build predicates from push down filters.
         PredicateBuilder::default()
-            .add_pushdown_exprs(filters)
-            .extract_time_range(&self.read_schema, filters)
+            .add_pushdown_exprs(&push_down_filters)
+            .extract_time_range(&self.read_schema, &push_down_filters)
             .build()
     }
+
+    fn should_push_down(&self, filter: &Expr, key_set: &HashSet<&str>) -> bool {
+        let columns = visitor::find_columns_by_expr(filter);
+        let mut should_push_down = true;
+        for column in columns {
+            // Once found a column not primary key or timestamp key in `filter`,
+            // the `filter` will be marked as unable to push down.
+            if !key_set.contains(column.as_str()) {
+                should_push_down = false;
+                break;
+            }
+        }
+
+        should_push_down
+    }
 }
 
 #[async_trait]
@@ -329,3 +380,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);
+    }
+}