risingwave_frontend/optimizer/plan_node/generic/

source.rs

// Copyright 2022 RisingWave Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::rc::Rc;

use educe::Educe;
use risingwave_common::catalog::{ColumnCatalog, ColumnDesc, Field, Schema};
use risingwave_common::types::DataType;
use risingwave_common::util::sort_util::OrderType;
use risingwave_connector::WithPropertiesExt;
use risingwave_sqlparser::ast::AsOf;

use super::super::utils::TableCatalogBuilder;
use super::GenericPlanNode;
use crate::TableCatalog;
use crate::catalog::source_catalog::SourceCatalog;
use crate::optimizer::optimizer_context::OptimizerContextRef;
use crate::optimizer::property::{FunctionalDependencySet, StreamKind};

/// In which scnario the source node is created
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
#[expect(clippy::enum_variant_names)]
pub enum SourceNodeKind {
    /// `CREATE TABLE` with a connector.
    CreateTable,
    /// `CREATE SOURCE` with a streaming job (shared source).
    CreateSharedSource,
    /// `CREATE MATERIALIZED VIEW` or batch scan from a source.
    ///
    /// Note:
    /// - For non-shared source, `CREATE SOURCE` will not create a source node, and `CREATE MATERIALIZE VIEW` will create a `StreamSource`.
    /// - For shared source, `CREATE MATERIALIZE VIEW` will create `StreamSourceScan` instead of `StreamSource`.
    CreateMViewOrBatch,
}

/// [`Source`] returns contents of a table or other equivalent object
#[derive(Debug, Clone, Educe)]
#[educe(PartialEq, Eq, Hash)]
pub struct Source {
    /// If there is an external stream source, `catalog` will be `Some`. Otherwise, it is `None`.
    pub catalog: Option<Rc<SourceCatalog>>,

    // NOTE: Here we store `column_catalog` and `row_id_index`
    // because they are needed when `catalog` is None.
    // When `catalog` is Some, they are the same as these fields in `catalog`.
    pub column_catalog: Vec<ColumnCatalog>,
    pub row_id_index: Option<usize>,

    pub kind: SourceNodeKind,

    #[educe(PartialEq(ignore))]
    #[educe(Hash(ignore))]
    pub ctx: OptimizerContextRef,

    pub as_of: Option<AsOf>,
}

impl GenericPlanNode for Source {
    fn schema(&self) -> Schema {
        let fields = self
            .column_catalog
            .iter()
            .map(|c| (&c.column_desc).into())
            .collect();
        Schema { fields }
    }

    fn stream_key(&self) -> Option<Vec<usize>> {
        // FIXME: output col idx is not set. But iceberg source can prune cols.
        // XXX: there's a RISINGWAVE_ICEBERG_ROW_ID. Should we use it?
        if let Some(idx) = self.row_id_index {
            Some(vec![idx])
        } else if let Some(catalog) = &self.catalog {
            catalog
                .pk_col_ids
                .iter()
                .map(|id| {
                    self.column_catalog
                        .iter()
                        .position(|c| c.column_id() == *id)
                })
                .collect::<Option<Vec<_>>>()
        } else {
            None
        }
    }

    fn ctx(&self) -> OptimizerContextRef {
        self.ctx.clone()
    }

    fn functional_dependency(&self) -> FunctionalDependencySet {
        let pk_indices = self.stream_key();
        match pk_indices {
            Some(pk_indices) => {
                debug_assert!(
                    pk_indices
                        .iter()
                        .all(|idx| *idx < self.column_catalog.len())
                );
                FunctionalDependencySet::with_key(self.column_catalog.len(), &pk_indices)
            }
            None => FunctionalDependencySet::new(self.column_catalog.len()),
        }
    }
}

impl Source {
    pub fn stream_kind(&self) -> StreamKind {
        if let Some(catalog) = &self.catalog {
            if catalog.is_iceberg_update_source() {
                StreamKind::Retract
            } else if catalog.append_only {
                StreamKind::AppendOnly
            } else {
                // Always treat source as upsert, as we either don't parse the old record for `Update`, or we don't
                // trust the old record from external source.
                StreamKind::Upsert
            }
        } else {
            // `Source` acts only as a barrier receiver. There's no data at all.
            StreamKind::AppendOnly
        }
    }

    /// The output is [`risingwave_connector::source::filesystem::FsPageItem`] / [`iceberg::scan::FileScanTask`]
    pub fn file_list_node(core: Self) -> Self {
        let column_catalog = if core.is_iceberg_update_source() {
            // Task assignments, in the layout of the Fetch state table.
            Self::iceberg_update_task_fields()
                .iter()
                .enumerate()
                .map(|(id, field)| ColumnCatalog {
                    column_desc: ColumnDesc::from_field_with_column_id(field, id as i32),
                    is_hidden: false,
                })
                .collect()
        } else if core.is_iceberg_connector() {
            vec![
                ColumnCatalog {
                    column_desc: ColumnDesc::from_field_with_column_id(
                        &Field {
                            name: "file_path".to_owned(),
                            data_type: DataType::Varchar,
                        },
                        0,
                    ),
                    is_hidden: false,
                },
                ColumnCatalog {
                    column_desc: ColumnDesc::from_field_with_column_id(
                        &Field {
                            name: "file_scan_task".to_owned(),
                            data_type: DataType::Jsonb,
                        },
                        1,
                    ),
                    is_hidden: false,
                },
            ]
        } else if core.is_batch_connector() {
            vec![
                ColumnCatalog {
                    column_desc: ColumnDesc::from_field_with_column_id(
                        &Field {
                            name: "batch_task_id".to_owned(),
                            data_type: DataType::Varchar,
                        },
                        0,
                    ),
                    is_hidden: false,
                },
                ColumnCatalog {
                    column_desc: ColumnDesc::from_field_with_column_id(
                        &Field {
                            name: "batch_task_info".to_owned(),
                            data_type: DataType::Jsonb,
                        },
                        1,
                    ),
                    is_hidden: false,
                },
            ]
        } else if core.is_new_fs_connector() {
            vec![
                ColumnCatalog {
                    column_desc: ColumnDesc::from_field_with_column_id(
                        &Field {
                            name: "filename".to_owned(),
                            data_type: DataType::Varchar,
                        },
                        0,
                    ),
                    is_hidden: false,
                },
                // This columns seems unused.
                ColumnCatalog {
                    column_desc: ColumnDesc::from_field_with_column_id(
                        &Field {
                            name: "last_edit_time".to_owned(),
                            data_type: DataType::Timestamptz,
                        },
                        1,
                    ),
                    is_hidden: false,
                },
                ColumnCatalog {
                    column_desc: ColumnDesc::from_field_with_column_id(
                        &Field {
                            name: "file_size".to_owned(),
                            data_type: DataType::Int64,
                        },
                        2,
                    ),
                    is_hidden: false,
                },
            ]
        } else {
            unreachable!()
        };
        Self {
            column_catalog,
            row_id_index: None,
            ..core
        }
    }

    pub fn is_new_fs_connector(&self) -> bool {
        self.catalog
            .as_ref()
            .is_some_and(|catalog| catalog.with_properties.is_new_fs_connector())
    }

    pub fn is_iceberg_connector(&self) -> bool {
        self.catalog
            .as_ref()
            .is_some_and(|catalog| catalog.with_properties.is_iceberg_connector())
    }

    pub fn is_iceberg_update_source(&self) -> bool {
        self.catalog
            .as_ref()
            .is_some_and(|catalog| catalog.is_iceberg_update_source())
    }

    pub fn is_kafka_connector(&self) -> bool {
        self.catalog
            .as_ref()
            .is_some_and(|catalog| catalog.with_properties.is_kafka_connector())
    }

    pub fn is_batch_connector(&self) -> bool {
        self.catalog
            .as_ref()
            .is_some_and(|catalog| catalog.with_properties.is_batch_connector())
    }

    pub fn requires_singleton(&self) -> bool {
        self.is_iceberg_connector() || self.is_batch_connector()
    }

    /// Currently, only iceberg source supports time travel.
    pub fn support_time_travel(&self) -> bool {
        self.is_iceberg_connector()
    }

    pub fn exclude_iceberg_hidden_columns(mut self) -> Self {
        if self.is_iceberg_update_source() {
            return self;
        }
        let Some(catalog) = &mut self.catalog else {
            return self;
        };
        if catalog.info.is_shared() {
            // for shared source, we should produce all columns
            return self;
        }
        if self.kind != SourceNodeKind::CreateMViewOrBatch {
            return self;
        }

        let prune = |col: &ColumnCatalog| col.is_hidden() && !col.is_row_id_column();

        // minus the number of hidden columns before row_id_index.
        self.row_id_index = self.row_id_index.map(|idx| {
            let mut cnt = 0;
            for col in self.column_catalog.iter().take(idx + 1) {
                if prune(col) {
                    cnt += 1;
                }
            }
            idx - cnt
        });
        self.column_catalog.retain(|c| !prune(c));
        self
    }

    /// The columns in stream/batch source node indicate the actual columns it will produce,
    /// instead of the columns defined in source catalog. The difference is generated columns.
    pub fn exclude_generated_columns(mut self) -> (Self, Option<usize>) {
        let original_row_id_index = self.row_id_index;
        // minus the number of generated columns before row_id_index.
        self.row_id_index = original_row_id_index.map(|idx| {
            let mut cnt = 0;
            for col in self.column_catalog.iter().take(idx + 1) {
                if col.is_generated() {
                    cnt += 1;
                }
            }
            idx - cnt
        });
        self.column_catalog.retain(|c| !c.is_generated());
        (self, original_row_id_index)
    }

    /// Source's state table is `partition_id -> offset_info`.
    /// Its schema is irrelevant to the data's schema.
    ///
    /// ## Notes on the distribution of the state table (`is_distributed`)
    ///
    /// Source executors are always distributed, but their state tables are special.
    ///
    /// ### `StreamSourceExecutor`: singleton (only one vnode)
    ///
    /// Its states are not sharded by consistent hash.
    ///
    /// Each actor accesses (point get) some partitions (a.k.a splits).
    /// They are assigned by `SourceManager` in meta,
    /// instead of `vnode` computed from the `partition_id`.
    ///
    /// ### `StreamFsFetch`: distributed by `partition_id`
    ///
    /// Each actor accesses (range scan) splits according to the `vnode`
    /// computed from `partition_id`.
    /// This is a normal distributed table.
    pub fn infer_internal_table_catalog(is_distributed: bool) -> TableCatalog {
        let mut builder = TableCatalogBuilder::default();

        let key = Field {
            data_type: DataType::Varchar,
            name: "partition_id".to_owned(),
        };
        let value = Field {
            data_type: DataType::Jsonb,
            name: "offset_info".to_owned(),
        };

        let ordered_col_idx = builder.add_column(&key);
        builder.add_column(&value);
        builder.add_order_column(ordered_col_idx, OrderType::ascending());

        builder.build(
            if is_distributed {
                vec![ordered_col_idx]
            } else {
                vec![]
            },
            1,
        )
    }

    fn iceberg_update_task_fields() -> [Field; 3] {
        [
            Field {
                data_type: DataType::Int64,
                name: "route".to_owned(),
            },
            Field {
                data_type: DataType::Int64,
                name: "seq".to_owned(),
            },
            Field {
                data_type: DataType::Jsonb,
                name: "task".to_owned(),
            },
        ]
    }

    /// The Fetch state table of an Iceberg `streaming_updates` source. Tasks are keyed by
    /// `(route, seq)` and distributed by `route`, so a Fetch actor scans the tasks of a route in
    /// dispatch order. The List-to-Fetch exchange must hash the same `route` column.
    pub fn infer_iceberg_update_fetch_table_catalog() -> TableCatalog {
        let mut builder = TableCatalogBuilder::default();
        let [route, seq, task] = Self::iceberg_update_task_fields();
        let route = builder.add_column(&route);
        let seq = builder.add_column(&seq);
        builder.add_column(&task);
        builder.add_order_column(route, OrderType::ascending());
        builder.add_order_column(seq, OrderType::ascending());
        builder.build(vec![route], 2)
    }

    pub fn clone_with_column_catalog(&self, column_catalog: Vec<ColumnCatalog>) -> Self {
        let row_id_index = column_catalog.iter().position(|c| c.is_row_id_column());
        Self {
            catalog: self.catalog.clone(),
            column_catalog,
            row_id_index,
            kind: self.kind.clone(),
            ctx: self.ctx.clone(),
            as_of: self.as_of.clone(),
        }
    }
}