query/dist_plan/
remote_dyn_filter_receiver.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::cmp::Ordering;
16use std::fmt;
17use std::hash::{Hash, Hasher};
18use std::sync::Arc;
19
20use async_trait::async_trait;
21use datafusion::catalog::Session;
22use datafusion::common::Result;
23use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
24use datafusion::physical_expr::utils::conjunction;
25use datafusion::physical_plan::ExecutionPlan;
26use datafusion::physical_plan::expressions::Column;
27use datafusion::physical_plan::filter::FilterExec;
28use datafusion::physical_planner::{ExtensionPlanner, PhysicalPlanner};
29use datafusion_common::{DFSchemaRef, DataFusionError};
30use datafusion_expr::{
31    Expr, Extension, LogicalPlan, UserDefinedLogicalNode, UserDefinedLogicalNodeCore,
32};
33use datafusion_physical_expr::PhysicalExpr;
34use datafusion_physical_expr::utils::collect_columns;
35use session::context::QueryContextRef;
36
37type InjectRemoteDynFilterReceiver =
38    dyn Fn(LogicalPlan, QueryContextRef) -> LogicalPlan + Send + Sync + 'static;
39
40/// Injects a logical remote dynamic filter receiver into a query plan.
41pub struct RemoteDynFilterReceiverInjector {
42    inject: Box<InjectRemoteDynFilterReceiver>,
43}
44
45impl RemoteDynFilterReceiverInjector {
46    pub fn new(
47        inject: impl Fn(LogicalPlan, QueryContextRef) -> LogicalPlan + Send + Sync + 'static,
48    ) -> Self {
49        Self {
50            inject: Box::new(inject),
51        }
52    }
53
54    pub fn maybe_inject(&self, plan: LogicalPlan, query_ctx: QueryContextRef) -> LogicalPlan {
55        (self.inject)(plan, query_ctx)
56    }
57}
58
59pub type RemoteDynFilterReceiverInjectorRef = Arc<RemoteDynFilterReceiverInjector>;
60
61/// A logical marker that is converted to a [`FilterExec`] carrying remote dynamic filters.
62#[derive(Clone)]
63pub struct RemoteDynFilterReceiverLogicalPlan {
64    input: Arc<LogicalPlan>,
65    dyn_filters: Vec<Arc<dyn PhysicalExpr>>,
66}
67
68impl RemoteDynFilterReceiverLogicalPlan {
69    pub fn new(input: LogicalPlan, dyn_filters: Vec<Arc<dyn PhysicalExpr>>) -> Self {
70        Self {
71            input: Arc::new(input),
72            dyn_filters,
73        }
74    }
75
76    pub fn name() -> &'static str {
77        "RemoteDynFilterReceiver"
78    }
79
80    pub fn into_logical_plan(self) -> LogicalPlan {
81        LogicalPlan::Extension(Extension {
82            node: Arc::new(self),
83        })
84    }
85
86    fn dyn_filters(&self) -> &[Arc<dyn PhysicalExpr>] {
87        &self.dyn_filters
88    }
89
90    fn ord_key(&self) -> String {
91        format!("input={:?}, dyn_filters={:?}", self.input, self.dyn_filters)
92    }
93}
94
95impl fmt::Debug for RemoteDynFilterReceiverLogicalPlan {
96    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
97        UserDefinedLogicalNodeCore::fmt_for_explain(self, f)
98    }
99}
100
101impl Hash for RemoteDynFilterReceiverLogicalPlan {
102    fn hash<H: Hasher>(&self, state: &mut H) {
103        self.input.hash(state);
104        self.dyn_filters.hash(state);
105    }
106}
107
108impl PartialEq for RemoteDynFilterReceiverLogicalPlan {
109    fn eq(&self, other: &Self) -> bool {
110        self.input == other.input && self.dyn_filters == other.dyn_filters
111    }
112}
113
114impl Eq for RemoteDynFilterReceiverLogicalPlan {}
115
116impl PartialOrd for RemoteDynFilterReceiverLogicalPlan {
117    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
118        Some(self.ord_key().cmp(&other.ord_key()))
119    }
120}
121
122impl UserDefinedLogicalNodeCore for RemoteDynFilterReceiverLogicalPlan {
123    fn name(&self) -> &str {
124        Self::name()
125    }
126
127    fn inputs(&self) -> Vec<&LogicalPlan> {
128        vec![self.input.as_ref()]
129    }
130
131    fn schema(&self) -> &DFSchemaRef {
132        self.input.schema()
133    }
134
135    fn expressions(&self) -> Vec<Expr> {
136        Vec::new()
137    }
138
139    fn necessary_children_exprs(&self, output_columns: &[usize]) -> Option<Vec<Vec<usize>>> {
140        let mut required = output_columns.to_vec();
141
142        for filter in &self.dyn_filters {
143            required.extend(
144                collect_columns(filter)
145                    .into_iter()
146                    .map(|column| column.index()),
147            );
148        }
149
150        required.sort_unstable();
151        required.dedup();
152        Some(vec![required])
153    }
154
155    fn fmt_for_explain(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
156        write!(f, "{}: filters={}", Self::name(), self.dyn_filters.len())
157    }
158
159    fn with_exprs_and_inputs(
160        &self,
161        _exprs: Vec<Expr>,
162        mut inputs: Vec<LogicalPlan>,
163    ) -> Result<Self> {
164        let input = inputs.pop().ok_or_else(|| {
165            DataFusionError::Internal(format!("Expected exactly one input with {}", Self::name()))
166        })?;
167        let dyn_filters = self
168            .dyn_filters
169            .iter()
170            .map(|filter| remap_physical_expr_columns(filter.clone(), input.schema().as_arrow()))
171            .collect::<Result<Vec<_>>>()?;
172        Ok(Self::new(input, dyn_filters))
173    }
174}
175
176fn remap_physical_expr_columns(
177    expr: Arc<dyn PhysicalExpr>,
178    input_schema: &datafusion::arrow::datatypes::Schema,
179) -> Result<Arc<dyn PhysicalExpr>> {
180    if let Some(column) = expr.downcast_ref::<Column>() {
181        return Ok(Arc::new(Column::new_with_schema(
182            column.name(),
183            input_schema,
184        )?));
185    }
186
187    let children = expr.children();
188    if children.is_empty() {
189        return Ok(expr);
190    }
191
192    let new_children = children
193        .into_iter()
194        .map(|child| remap_physical_expr_columns(child.clone(), input_schema))
195        .collect::<Result<Vec<_>>>()?;
196    expr.with_new_children(new_children)
197}
198
199pub struct RemoteDynFilterReceiverExtensionPlanner;
200
201#[async_trait]
202impl ExtensionPlanner for RemoteDynFilterReceiverExtensionPlanner {
203    async fn plan_extension(
204        &self,
205        _planner: &dyn PhysicalPlanner,
206        node: &dyn UserDefinedLogicalNode,
207        _logical_inputs: &[&LogicalPlan],
208        physical_inputs: &[Arc<dyn ExecutionPlan>],
209        _session: &dyn Session,
210        _planning_ctx: &PhysicalPlanningContext,
211    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
212        let Some(receiver) = node
213            .as_any()
214            .downcast_ref::<RemoteDynFilterReceiverLogicalPlan>()
215        else {
216            return Ok(None);
217        };
218
219        let input = physical_inputs.first().cloned().ok_or_else(|| {
220            DataFusionError::Internal(format!("Expected exactly one input with {}", Self::name()))
221        })?;
222        if receiver.dyn_filters().is_empty() {
223            return Ok(Some(input));
224        }
225
226        let predicate = conjunction(receiver.dyn_filters().to_vec());
227        Ok(Some(Arc::new(FilterExec::try_new(predicate, input)?) as _))
228    }
229}
230
231impl RemoteDynFilterReceiverExtensionPlanner {
232    fn name() -> &'static str {
233        RemoteDynFilterReceiverLogicalPlan::name()
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use std::sync::Arc;
240
241    use datafusion::arrow::datatypes::{DataType, Field, Schema};
242    use datafusion::physical_plan::expressions::{Column, DynamicFilterPhysicalExpr, lit};
243    use datafusion_common::DFSchema;
244    use datafusion_expr::{EmptyRelation, UserDefinedLogicalNodeCore};
245
246    use super::*;
247
248    fn empty_input() -> LogicalPlan {
249        let schema = Arc::new(Schema::new(vec![
250            Field::new(
251                "ts",
252                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
253                false,
254            ),
255            Field::new("value", DataType::Float64, true),
256            Field::new("instance", DataType::Utf8, true),
257            Field::new("job", DataType::Utf8, true),
258        ]));
259        LogicalPlan::EmptyRelation(EmptyRelation {
260            produce_one_row: false,
261            schema: Arc::new(DFSchema::try_from(schema).unwrap()),
262        })
263    }
264
265    fn pruned_input() -> LogicalPlan {
266        let schema = Arc::new(Schema::new(vec![
267            Field::new(
268                "ts",
269                DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None),
270                false,
271            ),
272            Field::new("instance", DataType::Utf8, true),
273        ]));
274        LogicalPlan::EmptyRelation(EmptyRelation {
275            produce_one_row: false,
276            schema: Arc::new(DFSchema::try_from(schema).unwrap()),
277        })
278    }
279
280    #[test]
281    fn necessary_children_exprs_keeps_parent_and_filter_columns() {
282        let dyn_filter = Arc::new(DynamicFilterPhysicalExpr::new(
283            vec![Arc::new(Column::new("ts", 0)) as Arc<_>],
284            lit(true) as _,
285        ));
286        let plan = RemoteDynFilterReceiverLogicalPlan::new(empty_input(), vec![dyn_filter]);
287
288        // Parent only needs `value`, but the receiver must still keep `ts` for
289        // evaluating its dynamic filter after logical projection pruning.
290        let required = UserDefinedLogicalNodeCore::necessary_children_exprs(&plan, &[1]).unwrap();
291        assert_eq!(required, vec![vec![0, 1]]);
292    }
293
294    #[test]
295    fn necessary_children_exprs_is_transparent_without_filters() {
296        let plan = RemoteDynFilterReceiverLogicalPlan::new(empty_input(), vec![]);
297
298        let required =
299            UserDefinedLogicalNodeCore::necessary_children_exprs(&plan, &[1, 3]).unwrap();
300        assert_eq!(required, vec![vec![1, 3]]);
301    }
302
303    #[test]
304    fn with_exprs_and_inputs_remaps_dyn_filter_columns_to_pruned_input() {
305        let dyn_filter = Arc::new(DynamicFilterPhysicalExpr::new(
306            vec![Arc::new(Column::new("instance", 2)) as Arc<_>],
307            lit(true) as _,
308        ));
309        let plan = RemoteDynFilterReceiverLogicalPlan::new(empty_input(), vec![dyn_filter]);
310
311        let remapped =
312            UserDefinedLogicalNodeCore::with_exprs_and_inputs(&plan, vec![], vec![pruned_input()])
313                .unwrap();
314
315        let columns = collect_columns(&remapped.dyn_filters()[0]);
316        assert_eq!(columns.len(), 1);
317        let column = columns.iter().next().unwrap();
318        assert_eq!(column.name(), "instance");
319        assert_eq!(column.index(), 1);
320    }
321}
query/dist_plan/remote_dyn_filter_receiver.rs

query/dist_plan/
remote_dyn_filter_receiver.rs
