add to async ReadParquetDataSource

databendlabs · Nov 30, 2023 · 271e9fb · 271e9fb
1 parent ebcc1ea
commit 271e9fb
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 10 deletions.
diff --git a/src/query/storages/fuse/src/operations/read/parquet_data_source_reader.rs b/src/query/storages/fuse/src/operations/read/parquet_data_source_reader.rs
@@ -109,6 +109,14 @@ impl SyncSource for ReadParquetDataSource<true> {
         match self.partitions.steal_one(self.id) {
             None => Ok(None),
             Some(part) => {
+                if runtime_filter_pruner(
+                    &part,
+                    &self.runtime_filters,
+                    &self.partitions.ctx.get_function_context()?,
+                )? {
+                    return Ok(None);
+                }
+
                 if let Some(index_reader) = self.index_reader.as_ref() {
                     let fuse_part = FusePartInfo::from_part(&part)?;
                     let loc =
@@ -148,14 +156,6 @@ impl SyncSource for ReadParquetDataSource<true> {
                     &None
                 };
 
-                if runtime_filter_pruner(
-                    &part,
-                    &self.runtime_filters,
-                    &self.partitions.ctx.get_function_context()?,
-                )? {
-                    return Ok(None);
-                }
-
                 let source = self.block_reader.sync_read_columns_data_by_merge_io(
                     &ReadSettings::from_ctx(&self.partitions.ctx)?,
                     &part,
@@ -223,12 +223,18 @@ impl Processor for ReadParquetDataSource<false> {
 
     #[async_backtrace::framed]
     async fn async_process(&mut self) -> Result<()> {
-        dbg!(&self.runtime_filters);
         let parts = self.partitions.steal(self.id, self.batch_size);
 
         if !parts.is_empty() {
             let mut chunks = Vec::with_capacity(parts.len());
             for part in &parts {
+                if runtime_filter_pruner(
+                    part,
+                    &self.runtime_filters,
+                    &self.partitions.ctx.get_function_context()?,
+                )? {
+                    continue;
+                }
                 let part = part.clone();
                 let block_reader = self.block_reader.clone();
                 let settings = ReadSettings::from_ctx(&self.partitions.ctx)?;

diff --git a/src/query/storages/fuse/src/operations/read/runtime_filter_prunner.rs b/src/query/storages/fuse/src/operations/read/runtime_filter_prunner.rs
@@ -38,6 +38,8 @@ pub fn runtime_filter_pruner(
     let part = FusePartInfo::from_part(part)?;
     Ok(filters.iter().any(|(id, filter)| {
         let column_refs = filter.column_refs();
+        // Currently only support filter with one column(probe key).
+        assert_debug!(column_refs.len() == 1);
         let ty = column_refs.values().last().unwrap();
         let name = column_refs.keys().last().unwrap();
         if let Some(stats) = &part.columns_stat {
@@ -52,7 +54,6 @@ pub fn runtime_filter_pruner(
                     func_ctx,
                     &BUILTIN_FUNCTIONS,
                 );
-                dbg!(&new_expr);
                 matches!(new_expr, Expr::Constant {
                     scalar: Scalar::Boolean(false),
                     ..