risingwave_frontend/optimizer/plan_node/
logical_source.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::rc::Rc;
16
17use pretty_xmlish::{Pretty, XmlNode};
18use risingwave_common::bail;
19use risingwave_common::catalog::{
20    ColumnCatalog, ICEBERG_FILE_PATH_COLUMN_NAME, ICEBERG_FILE_POS_COLUMN_NAME,
21    ICEBERG_SEQUENCE_NUM_COLUMN_NAME, ROW_ID_COLUMN_NAME,
22};
23use risingwave_pb::plan_common::GeneratedColumnDesc;
24use risingwave_pb::plan_common::column_desc::GeneratedOrDefaultColumn;
25use risingwave_pb::plan_common::source_refresh_mode::RefreshMode;
26use risingwave_sqlparser::ast::AsOf;
27
28use super::generic::{GenericPlanNode, GenericPlanRef, SourceNodeKind};
29use super::stream_watermark_filter::StreamWatermarkFilter;
30use super::utils::{Distill, childless_record};
31use super::{
32    BatchProject, BatchSource, ColPrunable, ExprRewritable, Logical, LogicalFilter,
33    LogicalPlanRef as PlanRef, LogicalProject, PlanBase, PredicatePushdown, StreamPlanRef,
34    StreamProject, StreamRowIdGen, StreamSource, StreamSourceScan, ToBatch, ToStream, generic,
35};
36use crate::catalog::source_catalog::SourceCatalog;
37use crate::error::Result;
38use crate::expr::{ExprImpl, ExprRewriter, ExprVisitor, InputRef};
39use crate::optimizer::optimizer_context::OptimizerContextRef;
40use crate::optimizer::plan_node::expr_visitable::ExprVisitable;
41use crate::optimizer::plan_node::stream_fs_fetch::StreamFsFetch;
42use crate::optimizer::plan_node::utils::column_names_pretty;
43use crate::optimizer::plan_node::{
44    ColumnPruningContext, PredicatePushdownContext, RewriteStreamContext, StreamDedup,
45    ToStreamContext,
46};
47use crate::optimizer::property::Distribution::HashShard;
48use crate::optimizer::property::{
49    Distribution, FunctionalDependencySet, MonotonicityMap, RequiredDist, StreamKind,
50    WatermarkColumns,
51};
52use crate::utils::{ColIndexMapping, Condition, IndexRewriter};
53
54/// `LogicalSource` returns contents of a table or other equivalent object
55#[derive(Debug, Clone, PartialEq, Eq, Hash)]
56pub struct LogicalSource {
57    pub base: PlanBase<Logical>,
58    pub core: generic::Source,
59
60    /// Expressions to output. This field presents and will be turned to a `Project` when
61    /// converting to a physical plan, only if there are generated columns.
62    pub(crate) output_exprs: Option<Vec<ExprImpl>>,
63    /// When there are generated columns, the `StreamRowIdGen`'s `row_id_index` is different from
64    /// the one in `core`. So we store the one in `output_exprs` here.
65    pub(crate) output_row_id_index: Option<usize>,
66}
67
68impl LogicalSource {
69    pub fn new(
70        source_catalog: Option<Rc<SourceCatalog>>,
71        column_catalog: Vec<ColumnCatalog>,
72        row_id_index: Option<usize>,
73        kind: SourceNodeKind,
74        ctx: OptimizerContextRef,
75        as_of: Option<AsOf>,
76    ) -> Result<Self> {
77        // XXX: should we reorder the columns?
78        // The order may be strange if the schema is changed, e.g., [foo:Varchar, _rw_kafka_timestamp:Timestamptz, _row_id:Serial, bar:Int32]
79        // related: https://github.com/risingwavelabs/risingwave/issues/16486
80        // The order does not matter much. The columns field is essentially a map indexed by the column id.
81        // It will affect what users will see in `SELECT *`.
82        // But not sure if we rely on the position of hidden column like `_row_id` somewhere. For `projected_row_id` we do so...
83        let core = generic::Source {
84            catalog: source_catalog,
85            column_catalog,
86            row_id_index,
87            kind,
88            ctx,
89            as_of,
90        };
91
92        if core.as_of.is_some() && !core.support_time_travel() {
93            bail!("Time travel is not supported for the source")
94        }
95
96        let base = PlanBase::new_logical_with_core(&core);
97
98        let output_exprs = Self::derive_output_exprs_from_generated_columns(&core.column_catalog)?;
99        let (core, output_row_id_index) = core.exclude_generated_columns();
100
101        Ok(LogicalSource {
102            base,
103            core,
104            output_exprs,
105            output_row_id_index,
106        })
107    }
108
109    pub fn with_catalog(
110        source_catalog: Rc<SourceCatalog>,
111        kind: SourceNodeKind,
112        ctx: OptimizerContextRef,
113        as_of: Option<AsOf>,
114    ) -> Result<Self> {
115        source_catalog.validate_iceberg_update_source()?;
116        let column_catalogs = source_catalog.columns.clone();
117        let row_id_index = source_catalog.row_id_index;
118        if !source_catalog.append_only {
119            assert!(row_id_index.is_none());
120        }
121
122        Self::new(
123            Some(source_catalog),
124            column_catalogs,
125            row_id_index,
126            kind,
127            ctx,
128            as_of,
129        )
130    }
131
132    /// If there are no generated columns, returns `None`.
133    ///
134    /// Otherwise, the returned expressions correspond to all columns.
135    /// Non-generated columns are represented by `InputRef`.
136    pub fn derive_output_exprs_from_generated_columns(
137        columns: &[ColumnCatalog],
138    ) -> Result<Option<Vec<ExprImpl>>> {
139        if !columns.iter().any(|c| c.is_generated()) {
140            return Ok(None);
141        }
142
143        let col_mapping = {
144            let mut mapping = vec![None; columns.len()];
145            let mut cur = 0;
146            for (idx, column) in columns.iter().enumerate() {
147                if !column.is_generated() {
148                    mapping[idx] = Some(cur);
149                    cur += 1;
150                } else {
151                    mapping[idx] = None;
152                }
153            }
154            ColIndexMapping::new(mapping, columns.len())
155        };
156
157        let mut rewriter = IndexRewriter::new(col_mapping);
158        let mut exprs = Vec::with_capacity(columns.len());
159        let mut cur = 0;
160        for column in columns {
161            let column_desc = &column.column_desc;
162            let ret_data_type = column_desc.data_type.clone();
163
164            if let Some(GeneratedOrDefaultColumn::GeneratedColumn(generated_column)) =
165                &column_desc.generated_or_default_column
166            {
167                let GeneratedColumnDesc { expr } = generated_column;
168                // TODO(yuhao): avoid this `from_expr_proto`.
169                let proj_expr =
170                    rewriter.rewrite_expr(ExprImpl::from_expr_proto(expr.as_ref().unwrap())?);
171                let casted_expr = proj_expr.cast_assign(&ret_data_type)?;
172                exprs.push(casted_expr);
173            } else {
174                let input_ref = InputRef {
175                    data_type: ret_data_type,
176                    index: cur,
177                };
178                cur += 1;
179                exprs.push(ExprImpl::InputRef(Box::new(input_ref)));
180            }
181        }
182
183        Ok(Some(exprs))
184    }
185
186    fn create_non_shared_source_plan(core: generic::Source) -> Result<StreamPlanRef> {
187        if core.is_iceberg_update_source() && core.as_of.is_some() {
188            bail!("Iceberg streaming_updates does not support time travel in streaming queries");
189        }
190        let mut plan;
191        if core.is_new_fs_connector() {
192            // Streaming file sources list objects repeatedly and need a persistent
193            // file-name dedup. FULL_RELOAD sources must reprocess the complete
194            // current object set on every refresh, including names seen earlier.
195            let dedup = !Self::is_full_reload_refresh(&core);
196            plan = Self::create_list_plan(core.clone(), dedup)?;
197            plan = StreamFsFetch::new(plan, core).into();
198        } else if core.is_iceberg_connector() || core.is_batch_connector() {
199            plan = Self::create_list_plan(core.clone(), false)?;
200            plan = StreamFsFetch::new(plan, core).into();
201        } else {
202            plan = StreamSource::new(core).into()
203        }
204        Ok(plan)
205    }
206
207    fn is_full_reload_refresh(core: &generic::Source) -> bool {
208        core.catalog.as_ref().is_some_and(|catalog| {
209            catalog.refresh_mode.as_ref().is_some_and(|refresh_mode| {
210                matches!(refresh_mode.refresh_mode, Some(RefreshMode::FullReload(_)))
211            })
212        })
213    }
214
215    /// `StreamSource` (list) -> shuffle -> (optional) `StreamDedup`
216    fn create_list_plan(core: generic::Source, dedup: bool) -> Result<StreamPlanRef> {
217        let updates = core.is_iceberg_update_source();
218        let downstream_columns = core.column_catalog.clone();
219        let logical_source = generic::Source::file_list_node(core);
220        let schema = logical_source.schema();
221        // Update assignments are identified by `(route, seq)` and shuffled by `route`.
222        let key = if updates { vec![0, 1] } else { vec![0] };
223        let functional_dependency = FunctionalDependencySet::with_key(schema.len(), &key);
224        let base = PlanBase::new_stream(
225            logical_source.ctx(),
226            schema,
227            Some(key),
228            functional_dependency,
229            Distribution::Single,
230            if updates {
231                StreamKind::Retract
232            } else {
233                StreamKind::AppendOnly
234            },
235            false,
236            WatermarkColumns::new(),
237            MonotonicityMap::new(),
238        );
239        let mut list_plan: StreamPlanRef = StreamSource {
240            base,
241            core: logical_source,
242            downstream_columns: Some(downstream_columns),
243        }
244        .into();
245        list_plan = RequiredDist::shard_by_key(list_plan.schema().len(), &[0])
246            .streaming_enforce_if_not_satisfies(list_plan)?;
247        if dedup {
248            list_plan = StreamDedup::new(generic::Dedup {
249                input: list_plan,
250                dedup_cols: vec![0],
251            })
252            .into();
253        }
254
255        Ok(list_plan)
256    }
257
258    pub fn source_catalog(&self) -> Option<Rc<SourceCatalog>> {
259        self.core.catalog.clone()
260    }
261
262    pub fn clone_with_column_catalog(&self, column_catalog: Vec<ColumnCatalog>) -> Result<Self> {
263        let row_id_index = column_catalog.iter().position(|c| c.is_row_id_column());
264        let kind = self.core.kind.clone();
265        let ctx = self.core.ctx.clone();
266        let as_of = self.core.as_of.clone();
267        Self::new(
268            self.source_catalog(),
269            column_catalog,
270            row_id_index,
271            kind,
272            ctx,
273            as_of,
274        )
275    }
276
277    fn prune_col_for_iceberg_source(&self, required_cols: &[usize]) -> PlanRef {
278        assert!(self.core.is_iceberg_connector());
279        // Iceberg source supports column pruning at source level
280        // Schema invariant: [table columns] + [_iceberg_sequence_number, _iceberg_file_path, _iceberg_file_pos, _row_id]
281        // The last 4 columns are always: 3 iceberg hidden columns + _row_id
282
283        let schema_len = self.schema().len();
284        assert!(
285            schema_len >= 4,
286            "Iceberg source must have at least 4 columns (3 iceberg hidden + 1 row_id)"
287        );
288
289        assert_eq!(
290            self.core.column_catalog[schema_len - 4].name(),
291            ICEBERG_SEQUENCE_NUM_COLUMN_NAME
292        );
293        assert_eq!(
294            self.core.column_catalog[schema_len - 3].name(),
295            ICEBERG_FILE_PATH_COLUMN_NAME
296        );
297        assert_eq!(
298            self.core.column_catalog[schema_len - 2].name(),
299            ICEBERG_FILE_POS_COLUMN_NAME
300        );
301        assert_eq!(
302            self.core.column_catalog[schema_len - 1].name(),
303            ROW_ID_COLUMN_NAME
304        );
305        assert_eq!(self.output_row_id_index, Some(self.schema().len() - 1));
306
307        let iceberg_start_idx = schema_len - 4;
308        let row_id_idx = schema_len - 1;
309
310        // Build source_cols: table columns from required_cols + always keep last 4 columns
311        let mut source_cols = Vec::new();
312
313        // Collect table columns (before the last 4 columns) from required_cols
314        for &idx in required_cols {
315            if idx < iceberg_start_idx {
316                // Regular table column
317                source_cols.push(idx);
318            }
319        }
320
321        // Always append the last 4 columns: [_iceberg_sequence_number, _iceberg_file_path, _iceberg_file_pos, _row_id]
322        source_cols.extend([
323            iceberg_start_idx,
324            iceberg_start_idx + 1,
325            iceberg_start_idx + 2,
326            row_id_idx,
327        ]);
328
329        // Clone with pruned columns - source_cols is never empty (always has last 4 columns)
330        let mut core = self.core.clone();
331        core.column_catalog = source_cols
332            .iter()
333            .map(|idx| core.column_catalog[*idx].clone())
334            .collect();
335        // row_id is always at the last position in the pruned schema
336        core.row_id_index = Some(source_cols.len() - 1);
337
338        let base = PlanBase::new_logical_with_core(&core);
339        let output_exprs =
340            Self::derive_output_exprs_from_generated_columns(&core.column_catalog).unwrap();
341        let (core, _) = core.exclude_generated_columns();
342
343        let pruned_source = LogicalSource {
344            base,
345            core,
346            output_exprs,
347            output_row_id_index: Some(source_cols.len() - 1),
348        };
349
350        // Build mapping from original schema indices to pruned schema indices
351        let mut old_to_new = vec![None; self.schema().len()];
352        for (new_idx, &old_idx) in source_cols.iter().enumerate() {
353            old_to_new[old_idx] = Some(new_idx);
354        }
355
356        // Map required_cols to indices in the pruned schema
357        let new_required: Vec<_> = required_cols
358            .iter()
359            .map(|&old_idx| old_to_new[old_idx].unwrap())
360            .collect();
361
362        let mapping =
363            ColIndexMapping::with_remaining_columns(&new_required, pruned_source.schema().len());
364        LogicalProject::with_mapping(pruned_source.into(), mapping).into()
365    }
366
367    pub fn is_shared_source(&self) -> bool {
368        // Create MV on source.
369        // We only check streaming_use_shared_source is true when `CREATE SOURCE`.
370        // The value does not affect the behavior of `CREATE MATERIALIZED VIEW` here.
371        self.source_catalog().is_some_and(|c| c.info.is_shared())
372    }
373}
374
375impl_plan_tree_node_for_leaf! { Logical, LogicalSource}
376impl Distill for LogicalSource {
377    fn distill<'a>(&self) -> XmlNode<'a> {
378        let fields = if let Some(catalog) = self.source_catalog() {
379            let src = Pretty::from(catalog.name.clone());
380            let mut fields = vec![
381                ("source", src),
382                ("is_shared", Pretty::debug(&catalog.info.is_shared())),
383                ("columns", column_names_pretty(self.schema())),
384            ];
385            if let Some(as_of) = &self.core.as_of {
386                fields.push(("as_of", Pretty::debug(as_of)));
387            }
388            fields
389        } else {
390            vec![]
391        };
392        childless_record("LogicalSource", fields)
393    }
394}
395
396impl ColPrunable for LogicalSource {
397    fn prune_col(&self, required_cols: &[usize], _ctx: &mut ColumnPruningContext) -> PlanRef {
398        // For refreshable iceberg table, we do not expose iceberg hidden columns to the user.
399        if self.core.is_iceberg_connector()
400            && !self.core.is_iceberg_update_source()
401            && !Self::is_full_reload_refresh(&self.core)
402        {
403            self.prune_col_for_iceberg_source(required_cols)
404        } else {
405            // For other sources, use a LogicalProject to prune columns
406            let mapping =
407                ColIndexMapping::with_remaining_columns(required_cols, self.schema().len());
408            LogicalProject::with_mapping(self.clone().into(), mapping).into()
409        }
410    }
411}
412
413impl ExprRewritable<Logical> for LogicalSource {
414    fn has_rewritable_expr(&self) -> bool {
415        self.output_exprs.is_some()
416    }
417
418    fn rewrite_exprs(&self, r: &mut dyn ExprRewriter) -> PlanRef {
419        let mut output_exprs = self.output_exprs.clone();
420
421        for expr in output_exprs.iter_mut().flatten() {
422            *expr = r.rewrite_expr(expr.clone());
423        }
424
425        Self {
426            output_exprs,
427            ..self.clone()
428        }
429        .into()
430    }
431}
432
433impl ExprVisitable for LogicalSource {
434    fn visit_exprs(&self, v: &mut dyn ExprVisitor) {
435        self.output_exprs
436            .iter()
437            .flatten()
438            .for_each(|e| v.visit_expr(e));
439    }
440}
441
442impl PredicatePushdown for LogicalSource {
443    fn predicate_pushdown(
444        &self,
445        predicate: Condition,
446        _ctx: &mut PredicatePushdownContext,
447    ) -> PlanRef {
448        LogicalFilter::create(self.clone().into(), predicate)
449    }
450}
451
452impl ToBatch for LogicalSource {
453    fn to_batch(&self) -> Result<crate::optimizer::plan_node::BatchPlanRef> {
454        assert!(
455            !self.core.is_kafka_connector(),
456            "LogicalSource with a kafka property should be converted to LogicalKafkaScan"
457        );
458        assert!(
459            !self.core.is_iceberg_connector(),
460            "LogicalSource with a iceberg property should be converted to LogicalIcebergScan"
461        );
462        let mut plan = BatchSource::new(self.core.clone()).into();
463
464        if let Some(exprs) = &self.output_exprs {
465            let logical_project = generic::Project::new(exprs.clone(), plan);
466            plan = BatchProject::new(logical_project).into();
467        }
468
469        Ok(plan)
470    }
471}
472
473impl ToStream for LogicalSource {
474    fn to_stream(
475        &self,
476        _ctx: &mut ToStreamContext,
477    ) -> Result<crate::optimizer::plan_node::StreamPlanRef> {
478        let mut plan;
479
480        match self.core.kind {
481            SourceNodeKind::CreateTable | SourceNodeKind::CreateSharedSource => {
482                // Note: for create table, row_id and generated columns is created in plan_root.gen_table_plan.
483                // for shared source, row_id and generated columns is created after SourceBackfill node.
484                plan = Self::create_non_shared_source_plan(self.core.clone())?;
485            }
486            SourceNodeKind::CreateMViewOrBatch => {
487                if self.is_shared_source() {
488                    plan = StreamSourceScan::new(self.core.clone()).into();
489                } else {
490                    // non-shared source
491                    plan = Self::create_non_shared_source_plan(self.core.clone())?;
492                }
493
494                if let Some(exprs) = &self.output_exprs {
495                    let logical_project = generic::Project::new(exprs.clone(), plan);
496                    plan = StreamProject::new(logical_project).into();
497                }
498
499                if let Some(row_id_index) = self.output_row_id_index {
500                    plan = StreamRowIdGen::new_with_dist(
501                        plan,
502                        row_id_index,
503                        HashShard(vec![row_id_index]),
504                    )
505                    .into();
506                }
507
508                if let Some(catalog) = self.source_catalog()
509                    && !catalog.watermark_descs.is_empty()
510                {
511                    plan = StreamWatermarkFilter::new(plan, catalog.watermark_descs.clone()).into();
512                }
513            }
514        }
515        Ok(plan)
516    }
517
518    fn logical_rewrite_for_stream(
519        &self,
520        _ctx: &mut RewriteStreamContext,
521    ) -> Result<(PlanRef, ColIndexMapping)> {
522        Ok((
523            self.clone().into(),
524            ColIndexMapping::identity(self.schema().len()),
525        ))
526    }
527}
risingwave_frontend/optimizer/plan_node/logical_source.rs

risingwave_frontend/optimizer/plan_node/
logical_source.rs
