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

risingwave_frontend/optimizer/plan_node/
logical_source.rs
