risingwave_frontend/optimizer/plan_node/
logical_values.rs

1// Copyright 2022 RisingWave Labs
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::vec;
16
17use itertools::Itertools;
18use risingwave_common::catalog::{Field, Schema};
19use risingwave_common::types::{DataType, ScalarImpl};
20
21use super::generic::GenericPlanRef;
22use super::utils::impl_distill_by_unit;
23use super::{
24    BatchValues, ColPrunable, ExprRewritable, Logical, LogicalFilter, LogicalPlanRef as PlanRef,
25    PlanBase, PredicatePushdown, StreamValues, ToBatch, ToStream, generic,
26};
27use crate::error::Result;
28use crate::expr::{ExprImpl, ExprRewriter, ExprVisitor, Literal};
29use crate::optimizer::optimizer_context::OptimizerContextRef;
30use crate::optimizer::plan_node::expr_visitable::ExprVisitable;
31use crate::optimizer::plan_node::{
32    ColumnPruningContext, PredicatePushdownContext, RewriteStreamContext, ToStreamContext,
33};
34use crate::utils::{ColIndexMapping, Condition};
35
36/// `LogicalValues` builds rows according to a list of expressions
37#[derive(Debug, Clone, PartialEq, Eq, Hash)]
38pub struct LogicalValues {
39    pub base: PlanBase<Logical>,
40    pub(super) core: generic::Values,
41}
42
43impl LogicalValues {
44    fn with_core(core: generic::Values) -> Self {
45        let base = PlanBase::new_logical_with_core(&core);
46        Self { base, core }
47    }
48
49    /// Create a [`LogicalValues`] node. Used internally by optimizer.
50    pub fn new(rows: Vec<Vec<ExprImpl>>, schema: Schema, ctx: OptimizerContextRef) -> Self {
51        Self::with_core(generic::Values::new(rows, schema, ctx))
52    }
53
54    /// Used only by `LogicalValues.rewrite_logical_for_stream`, set the `_row_id` column as pk
55    fn new_with_pk(
56        rows: Vec<Vec<ExprImpl>>,
57        schema: Schema,
58        ctx: OptimizerContextRef,
59        pk_index: usize,
60    ) -> Self {
61        Self::with_core(generic::Values::new_with_stream_key(
62            rows,
63            schema,
64            ctx,
65            vec![pk_index],
66        ))
67    }
68
69    /// Create a [`LogicalValues`] node. Used by planner.
70    pub fn create(rows: Vec<Vec<ExprImpl>>, schema: Schema, ctx: OptimizerContextRef) -> PlanRef {
71        // No additional checks after binder.
72        Self::new(rows, schema, ctx).into()
73    }
74
75    /// Create a [`LogicalValues`] node with a single empty row, as a dummy input for `Project` or `ProjectSet`.
76    pub fn create_empty_scalar(ctx: OptimizerContextRef) -> PlanRef {
77        Self::new(vec![vec![]], Schema::new(vec![]), ctx).into()
78    }
79
80    /// Check whether this is an empty scalar, typically created by [`LogicalValues::create_empty_scalar`].
81    pub fn is_empty_scalar(&self) -> bool {
82        self.schema().is_empty() && self.rows().len() == 1 && self.rows()[0].is_empty()
83    }
84
85    /// Get a reference to the logical values' rows.
86    pub fn rows(&self) -> &[Vec<ExprImpl>] {
87        self.core.rows()
88    }
89}
90
91impl_plan_tree_node_for_leaf! { Logical, LogicalValues }
92impl_distill_by_unit!(LogicalValues, core, "LogicalValues");
93
94impl ExprRewritable<Logical> for LogicalValues {
95    fn has_rewritable_expr(&self) -> bool {
96        true
97    }
98
99    fn rewrite_exprs(&self, r: &mut dyn ExprRewriter) -> PlanRef {
100        let mut core = self.core.clone();
101        core.rewrite_exprs(r);
102        Self::with_core(core).into()
103    }
104}
105
106impl ExprVisitable for LogicalValues {
107    fn visit_exprs(&self, v: &mut dyn ExprVisitor) {
108        self.core.visit_exprs(v);
109    }
110}
111
112impl ColPrunable for LogicalValues {
113    fn prune_col(&self, required_cols: &[usize], _ctx: &mut ColumnPruningContext) -> PlanRef {
114        let rows = self
115            .rows()
116            .iter()
117            .map(|row| required_cols.iter().map(|i| row[*i].clone()).collect())
118            .collect();
119        let fields = required_cols
120            .iter()
121            .map(|i| self.schema().fields[*i].clone())
122            .collect();
123        Self::new(rows, Schema { fields }, self.base.ctx()).into()
124    }
125}
126
127impl PredicatePushdown for LogicalValues {
128    fn predicate_pushdown(
129        &self,
130        predicate: Condition,
131        _ctx: &mut PredicatePushdownContext,
132    ) -> PlanRef {
133        LogicalFilter::create(self.clone().into(), predicate)
134    }
135}
136
137impl ToBatch for LogicalValues {
138    fn to_batch(&self) -> Result<crate::optimizer::plan_node::BatchPlanRef> {
139        Ok(BatchValues::new(self.core.clone()).into())
140    }
141}
142
143impl ToStream for LogicalValues {
144    fn to_stream(
145        &self,
146        _ctx: &mut ToStreamContext,
147    ) -> Result<crate::optimizer::plan_node::StreamPlanRef> {
148        Ok(StreamValues::new(self.core.clone()).into())
149    }
150
151    fn logical_rewrite_for_stream(
152        &self,
153        _ctx: &mut RewriteStreamContext,
154    ) -> Result<(PlanRef, ColIndexMapping)> {
155        let row_id_index = self.schema().len();
156        let col_index_mapping = ColIndexMapping::identity_or_none(row_id_index, row_id_index + 1);
157        let ctx = self.ctx();
158        let mut schema = self.schema().clone();
159        schema
160            .fields
161            .push(Field::with_name(DataType::Int64, "_row_id"));
162        let rows = self.rows().to_owned();
163        let row_with_id = rows
164            .into_iter()
165            .enumerate()
166            .map(|(i, mut r)| {
167                r.push(Literal::new(Some(ScalarImpl::Int64(i as i64)), DataType::Int64).into());
168                r
169            })
170            .collect_vec();
171        let logical_values = Self::new_with_pk(row_with_id, schema, ctx, row_id_index);
172        Ok((logical_values.into(), col_index_mapping))
173    }
174}
175
176#[cfg(test)]
177mod tests {
178
179    use risingwave_common::types::Datum;
180
181    use super::*;
182    use crate::optimizer::optimizer_context::OptimizerContext;
183
184    fn literal(val: i32) -> ExprImpl {
185        Literal::new(Datum::Some(val.into()), DataType::Int32).into()
186    }
187
188    /// Pruning
189    /// ```text
190    /// Values([[0, 1, 2], [3, 4, 5])
191    /// ```
192    /// with required columns [0, 2] will result in
193    /// ```text
194    /// Values([[0, 2], [3, 5])
195    /// ```
196    #[tokio::test]
197    async fn test_prune_filter() {
198        let ctx = OptimizerContext::mock();
199        let schema = Schema::new(vec![
200            Field::with_name(DataType::Int32, "v1"),
201            Field::with_name(DataType::Int32, "v2"),
202            Field::with_name(DataType::Int32, "v3"),
203        ]);
204        // Values([[0, 1, 2], [3, 4, 5])
205        let values: PlanRef = LogicalValues::new(
206            vec![
207                vec![literal(0), literal(1), literal(2)],
208                vec![literal(3), literal(4), literal(5)],
209            ],
210            schema,
211            ctx,
212        )
213        .into();
214
215        let required_cols = vec![0, 2];
216        let pruned = values.prune_col(
217            &required_cols,
218            &mut ColumnPruningContext::new(values.clone()),
219        );
220
221        let values = pruned.as_logical_values().unwrap();
222        let rows: &[Vec<ExprImpl>] = values.rows();
223
224        // expected output: Values([[0, 2], [3, 5])
225        assert_eq!(rows.len(), 2);
226        assert_eq!(rows[0].len(), 2);
227        assert_eq!(rows[0][0], literal(0));
228        assert_eq!(rows[0][1], literal(2));
229        assert_eq!(rows[1].len(), 2);
230        assert_eq!(rows[1][0], literal(3));
231        assert_eq!(rows[1][1], literal(5));
232    }
233}
risingwave_frontend/optimizer/plan_node/logical_values.rs

risingwave_frontend/optimizer/plan_node/
logical_values.rs
