feat: precise filter for mito parquet reader (#3178)

* impl SimpleFilterEvaluator Signed-off-by: Ruihang Xia <[email protected]> * time index and field filter Signed-off-by: Ruihang Xia <[email protected]> * finish parquet filter Signed-off-by: Ruihang Xia <[email protected]> * remove empty Batch Signed-off-by: Ruihang Xia <[email protected]> * fix clippy Signed-off-by: Ruihang Xia <[email protected]> * fix fmt Signed-off-by: Ruihang Xia <[email protected]> * fix typo Signed-off-by: Ruihang Xia <[email protected]> * update metric Signed-off-by: Ruihang Xia <[email protected]> * use projected schema from batch Signed-off-by: Ruihang Xia <[email protected]> * correct naming Signed-off-by: Ruihang Xia <[email protected]> * remove unnecessary error Signed-off-by: Ruihang Xia <[email protected]> * fix clippy Signed-off-by: Ruihang Xia <[email protected]> --------- Signed-off-by: Ruihang Xia <[email protected]>
GreptimeTeam · Jan 18, 2024 · cde5a36 · cde5a36
1 parent 6320590
commit cde5a36
Show file tree

Hide file tree

Showing 17 changed files with 652 additions and 32 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/common/recordbatch/Cargo.toml b/src/common/recordbatch/Cargo.toml
@@ -6,6 +6,7 @@ license.workspace = true
 
 [dependencies]
 arc-swap = "1.6"
+common-base.workspace = true
 common-error.workspace = true
 common-macro.workspace = true
 datafusion-common.workspace = true

diff --git a/src/common/recordbatch/src/error.rs b/src/common/recordbatch/src/error.rs
@@ -107,6 +107,16 @@ pub enum Error {
         location: Location,
         source: datatypes::error::Error,
     },
+
+    #[snafu(display("Error occurs when performing arrow computation"))]
+    ArrowCompute {
+        #[snafu(source)]
+        error: datatypes::arrow::error::ArrowError,
+        location: Location,
+    },
+
+    #[snafu(display("Unsupported operation: {}", reason))]
+    UnsupportedOperation { reason: String, location: Location },
 }
 
 impl ErrorExt for Error {
@@ -120,10 +130,13 @@ impl ErrorExt for Error {
             | Error::Format { .. }
             | Error::InitRecordbatchStream { .. }
             | Error::ColumnNotExists { .. }
-            | Error::ProjectArrowRecordBatch { .. } => StatusCode::Internal,
+            | Error::ProjectArrowRecordBatch { .. }
+            | Error::ArrowCompute { .. } => StatusCode::Internal,
 
             Error::External { source, .. } => source.status_code(),
 
+            Error::UnsupportedOperation { .. } => StatusCode::Unsupported,
+
             Error::SchemaConversion { source, .. } | Error::CastVector { source, .. } => {
                 source.status_code()
             }

diff --git a/src/common/recordbatch/src/filter.rs b/src/common/recordbatch/src/filter.rs
@@ -0,0 +1,258 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Util record batch stream wrapper that can perform precise filter.
+
+use datafusion::logical_expr::{Expr, Operator};
+use datafusion_common::arrow::array::{ArrayRef, Datum, Scalar};
+use datafusion_common::arrow::buffer::BooleanBuffer;
+use datafusion_common::arrow::compute::kernels::cmp;
+use datafusion_common::ScalarValue;
+use datatypes::vectors::VectorRef;
+use snafu::ResultExt;
+
+use crate::error::{ArrowComputeSnafu, Result, UnsupportedOperationSnafu};
+
+/// An inplace expr evaluator for simple filter. Only support
+/// - `col` `op` `literal`
+/// - `literal` `op` `col`
+///
+/// And the `op` is one of `=`, `!=`, `>`, `>=`, `<`, `<=`.
+///
+/// This struct contains normalized predicate expr. In the form of
+/// `col` `op` `literal` where the `col` is provided from input.
+#[derive(Debug)]
+pub struct SimpleFilterEvaluator {
+    /// Name of the referenced column.
+    column_name: String,
+    /// The literal value.
+    literal: Scalar<ArrayRef>,
+    /// The operator.
+    op: Operator,
+}
+
+impl SimpleFilterEvaluator {
+    pub fn try_new(predicate: &Expr) -> Option<Self> {
+        match predicate {
+            Expr::BinaryExpr(binary) => {
+                // check if the expr is in the supported form
+                match binary.op {
+                    Operator::Eq
+                    | Operator::NotEq
+                    | Operator::Lt
+                    | Operator::LtEq
+                    | Operator::Gt
+                    | Operator::GtEq => {}
+                    _ => return None,
+                }
+
+                // swap the expr if it is in the form of `literal` `op` `col`
+                let mut op = binary.op;
+                let (lhs, rhs) = match (&*binary.left, &*binary.right) {
+                    (Expr::Column(ref col), Expr::Literal(ref lit)) => (col, lit),
+                    (Expr::Literal(ref lit), Expr::Column(ref col)) => {
+                        // safety: The previous check ensures the operator is able to swap.
+                        op = op.swap().unwrap();
+                        (col, lit)
+                    }
+                    _ => return None,
+                };
+
+                Some(Self {
+                    column_name: lhs.name.clone(),
+                    literal: rhs.clone().to_scalar(),
+                    op,
+                })
+            }
+            _ => None,
+        }
+    }
+
+    /// Get the name of the referenced column.
+    pub fn column_name(&self) -> &str {
+        &self.column_name
+    }
+
+    pub fn evaluate_scalar(&self, input: &ScalarValue) -> Result<bool> {
+        let result = self.evaluate_datum(&input.to_scalar())?;
+        Ok(result.value(0))
+    }
+
+    pub fn evaluate_array(&self, input: &ArrayRef) -> Result<BooleanBuffer> {
+        self.evaluate_datum(input)
+    }
+
+    pub fn evaluate_vector(&self, input: &VectorRef) -> Result<BooleanBuffer> {
+        self.evaluate_datum(&input.to_arrow_array())
+    }
+
+    fn evaluate_datum(&self, input: &impl Datum) -> Result<BooleanBuffer> {
+        let result = match self.op {
+            Operator::Eq => cmp::eq(input, &self.literal),
+            Operator::NotEq => cmp::neq(input, &self.literal),
+            Operator::Lt => cmp::lt(input, &self.literal),
+            Operator::LtEq => cmp::lt_eq(input, &self.literal),
+            Operator::Gt => cmp::gt(input, &self.literal),
+            Operator::GtEq => cmp::gt_eq(input, &self.literal),
+            _ => {
+                return UnsupportedOperationSnafu {
+                    reason: format!("{:?}", self.op),
+                }
+                .fail()
+            }
+        };
+        result
+            .context(ArrowComputeSnafu)
+            .map(|array| array.values().clone())
+    }
+}
+
+#[cfg(test)]
+mod test {
+
+    use std::sync::Arc;
+
+    use datafusion::logical_expr::BinaryExpr;
+    use datafusion_common::Column;
+
+    use super::*;
+
+    #[test]
+    fn unsupported_filter_op() {
+        // `+` is not supported
+        let expr = Expr::BinaryExpr(BinaryExpr {
+            left: Box::new(Expr::Column(Column {
+                relation: None,
+                name: "foo".to_string(),
+            })),
+            op: Operator::Plus,
+            right: Box::new(Expr::Literal(ScalarValue::Int64(Some(1)))),
+        });
+        assert!(SimpleFilterEvaluator::try_new(&expr).is_none());
+
+        // two literal is not supported
+        let expr = Expr::BinaryExpr(BinaryExpr {
+            left: Box::new(Expr::Literal(ScalarValue::Int64(Some(1)))),
+            op: Operator::Eq,
+            right: Box::new(Expr::Literal(ScalarValue::Int64(Some(1)))),
+        });
+        assert!(SimpleFilterEvaluator::try_new(&expr).is_none());
+
+        // two column is not supported
+        let expr = Expr::BinaryExpr(BinaryExpr {
+            left: Box::new(Expr::Column(Column {
+                relation: None,
+                name: "foo".to_string(),
+            })),
+            op: Operator::Eq,
+            right: Box::new(Expr::Column(Column {
+                relation: None,
+                name: "bar".to_string(),
+            })),
+        });
+        assert!(SimpleFilterEvaluator::try_new(&expr).is_none());
+
+        // compound expr is not supported
+        let expr = Expr::BinaryExpr(BinaryExpr {
+            left: Box::new(Expr::BinaryExpr(BinaryExpr {
+                left: Box::new(Expr::Column(Column {
+                    relation: None,
+                    name: "foo".to_string(),
+                })),
+                op: Operator::Eq,
+                right: Box::new(Expr::Literal(ScalarValue::Int64(Some(1)))),
+            })),
+            op: Operator::Eq,
+            right: Box::new(Expr::Literal(ScalarValue::Int64(Some(1)))),
+        });
+        assert!(SimpleFilterEvaluator::try_new(&expr).is_none());
+    }
+
+    #[test]
+    fn supported_filter_op() {
+        // equal
+        let expr = Expr::BinaryExpr(BinaryExpr {
+            left: Box::new(Expr::Column(Column {
+                relation: None,
+                name: "foo".to_string(),
+            })),
+            op: Operator::Eq,
+            right: Box::new(Expr::Literal(ScalarValue::Int64(Some(1)))),
+        });
+        let _ = SimpleFilterEvaluator::try_new(&expr).unwrap();
+
+        // swap operands
+        let expr = Expr::BinaryExpr(BinaryExpr {
+            left: Box::new(Expr::Literal(ScalarValue::Int64(Some(1)))),
+            op: Operator::Lt,
+            right: Box::new(Expr::Column(Column {
+                relation: None,
+                name: "foo".to_string(),
+            })),
+        });
+        let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap();
+        assert_eq!(evaluator.op, Operator::Gt);
+        assert_eq!(evaluator.column_name, "foo".to_string());
+    }
+
+    #[test]
+    fn run_on_array() {
+        let expr = Expr::BinaryExpr(BinaryExpr {
+            left: Box::new(Expr::Column(Column {
+                relation: None,
+                name: "foo".to_string(),
+            })),
+            op: Operator::Eq,
+            right: Box::new(Expr::Literal(ScalarValue::Int64(Some(1)))),
+        });
+        let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap();
+
+        let input_1 = Arc::new(datatypes::arrow::array::Int64Array::from(vec![1, 2, 3])) as _;
+        let result = evaluator.evaluate_array(&input_1).unwrap();
+        assert_eq!(result, BooleanBuffer::from(vec![true, false, false]));
+
+        let input_2 = Arc::new(datatypes::arrow::array::Int64Array::from(vec![1, 1, 1])) as _;
+        let result = evaluator.evaluate_array(&input_2).unwrap();
+        assert_eq!(result, BooleanBuffer::from(vec![true, true, true]));
+
+        let input_3 = Arc::new(datatypes::arrow::array::Int64Array::new_null(0)) as _;
+        let result = evaluator.evaluate_array(&input_3).unwrap();
+        assert_eq!(result, BooleanBuffer::from(vec![]));
+    }
+
+    #[test]
+    fn run_on_scalar() {
+        let expr = Expr::BinaryExpr(BinaryExpr {
+            left: Box::new(Expr::Column(Column {
+                relation: None,
+                name: "foo".to_string(),
+            })),
+            op: Operator::Lt,
+            right: Box::new(Expr::Literal(ScalarValue::Int64(Some(1)))),
+        });
+        let evaluator = SimpleFilterEvaluator::try_new(&expr).unwrap();
+
+        let input_1 = ScalarValue::Int64(Some(1));
+        let result = evaluator.evaluate_scalar(&input_1).unwrap();
+        assert!(!result);
+
+        let input_2 = ScalarValue::Int64(Some(0));
+        let result = evaluator.evaluate_scalar(&input_2).unwrap();
+        assert!(result);
+
+        let input_3 = ScalarValue::Int64(None);
+        let result = evaluator.evaluate_scalar(&input_3).unwrap();
+        assert!(!result);
+    }
+}
diff --git a/src/common/recordbatch/src/lib.rs b/src/common/recordbatch/src/lib.rs
@@ -14,6 +14,7 @@
 
 pub mod adapter;
 pub mod error;
+pub mod filter;
 mod recordbatch;
 pub mod util;
 

diff --git a/src/mito2/src/engine/prune_test.rs b/src/mito2/src/engine/prune_test.rs
@@ -76,16 +76,11 @@ async fn test_read_parquet_stats() {
 +-------+---------+---------------------+
 | tag_0 | field_0 | ts                  |
 +-------+---------+---------------------+
-| 0     | 0.0     | 1970-01-01T00:00:00 |
-| 1     | 1.0     | 1970-01-01T00:00:01 |
 | 10    | 10.0    | 1970-01-01T00:00:10 |
 | 11    | 11.0    | 1970-01-01T00:00:11 |
 | 12    | 12.0    | 1970-01-01T00:00:12 |
 | 13    | 13.0    | 1970-01-01T00:00:13 |
 | 14    | 14.0    | 1970-01-01T00:00:14 |
-| 2     | 2.0     | 1970-01-01T00:00:02 |
-| 3     | 3.0     | 1970-01-01T00:00:03 |
-| 4     | 4.0     | 1970-01-01T00:00:04 |
 | 5     | 5.0     | 1970-01-01T00:00:05 |
 | 6     | 6.0     | 1970-01-01T00:00:06 |
 | 7     | 7.0     | 1970-01-01T00:00:07 |

diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs
@@ -536,6 +536,12 @@ pub enum Error {
         error: std::io::Error,
         location: Location,
     },
+
+    #[snafu(display("Failed to filter record batch"))]
+    FilterRecordBatch {
+        source: common_recordbatch::error::Error,
+        location: Location,
+    },
 }
 
 pub type Result<T, E = Error> = std::result::Result<T, E>;
@@ -632,6 +638,7 @@ impl ErrorExt for Error {
             CleanDir { .. } => StatusCode::Unexpected,
             InvalidConfig { .. } => StatusCode::InvalidArguments,
             StaleLogEntry { .. } => StatusCode::Unexpected,
+            FilterRecordBatch { source, .. } => source.status_code(),
             Upload { .. } => StatusCode::StorageUnavailable,
         }
     }

diff --git a/src/mito2/src/metrics.rs b/src/mito2/src/metrics.rs
@@ -124,6 +124,9 @@ lazy_static! {
     /// Counter of row groups read.
     pub static ref READ_ROW_GROUPS_TOTAL: IntCounterVec =
         register_int_counter_vec!("greptime_mito_read_row_groups_total", "mito read row groups total", &[TYPE_LABEL]).unwrap();
+    /// Counter of filtered rows by precise filter.
+    pub static ref PRECISE_FILTER_ROWS_TOTAL: IntCounterVec =
+        register_int_counter_vec!("greptime_mito_precise_filter_rows_total", "mito precise filter rows total", &[TYPE_LABEL]).unwrap();
     // ------- End of query metrics.
 
     // Cache related metrics.