risingwave_frontend/expr/
input_ref.rs

1// Copyright 2025 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::fmt;
16
17use itertools::Itertools;
18use risingwave_common::catalog::Schema;
19use risingwave_common::types::DataType;
20use risingwave_pb::expr::PbInputRef;
21
22use super::Expr;
23use crate::expr::ExprType;
24#[derive(Clone, Eq, PartialEq, Hash)]
25pub struct InputRef {
26    // TODO(rc): remove `pub`, use `new()`, `index()` and `data_type()` instead
27    pub index: usize,
28    pub data_type: DataType,
29}
30
31#[derive(Clone, Copy)]
32pub struct RawInputRefDisplay(pub usize);
33
34pub fn input_ref_to_column_indices(input_refs: &[InputRef]) -> Vec<usize> {
35    input_refs.iter().map(|x| x.index()).collect_vec()
36}
37
38impl fmt::Display for RawInputRefDisplay {
39    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
40        (self as &dyn fmt::Debug).fmt(f)
41    }
42}
43
44impl fmt::Debug for RawInputRefDisplay {
45    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
46        write!(f, "${}", self.0)
47    }
48}
49
50#[derive(Clone, Copy)]
51pub struct InputRefDisplay<'a> {
52    pub input_ref: &'a InputRef,
53    pub input_schema: &'a Schema,
54}
55
56impl fmt::Display for InputRefDisplay<'_> {
57    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
58        (self as &dyn fmt::Debug).fmt(f)
59    }
60}
61
62impl fmt::Debug for InputRefDisplay<'_> {
63    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
64        write!(
65            f,
66            "{}",
67            self.input_schema
68                .fields
69                .get(self.input_ref.index)
70                .unwrap()
71                .name
72        )
73    }
74}
75
76impl fmt::Display for InputRef {
77    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
78        write!(f, "{:?}", RawInputRefDisplay(self.index))
79    }
80}
81
82impl fmt::Debug for InputRef {
83    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
84        if f.alternate() {
85            f.debug_struct("InputRef")
86                .field("index", &self.index)
87                .field("data_type", &self.data_type)
88                .finish()
89        } else {
90            write!(f, "{}", RawInputRefDisplay(self.index))
91        }
92    }
93}
94
95impl InputRef {
96    pub fn new(index: usize, data_type: DataType) -> Self {
97        InputRef { index, data_type }
98    }
99
100    /// Get a reference to the input ref's index.
101    pub fn index(&self) -> usize {
102        self.index
103    }
104
105    /// Shift the input ref's index with offset.
106    pub fn shift_with_offset(&mut self, offset: isize) {
107        self.index = (self.index as isize + offset) as usize;
108    }
109
110    pub fn clone_with_offset(&self, offset: isize) -> Self {
111        Self {
112            index: (self.index as isize + offset) as usize,
113            data_type: self.data_type.clone(),
114        }
115    }
116
117    /// Convert to protobuf.
118    pub fn to_proto(&self) -> PbInputRef {
119        PbInputRef {
120            index: self.index as _,
121            r#type: Some(self.data_type.to_protobuf()),
122        }
123    }
124
125    pub(super) fn from_expr_proto(
126        column_index: usize,
127        ret_type: DataType,
128    ) -> crate::error::Result<Self> {
129        Ok(Self {
130            index: column_index,
131            data_type: ret_type,
132        })
133    }
134}
135
136impl Expr for InputRef {
137    fn return_type(&self) -> DataType {
138        self.data_type.clone()
139    }
140
141    fn to_expr_proto(&self) -> risingwave_pb::expr::ExprNode {
142        use risingwave_pb::expr::expr_node::*;
143        use risingwave_pb::expr::*;
144        ExprNode {
145            function_type: ExprType::Unspecified.into(),
146            return_type: Some(self.return_type().to_protobuf()),
147            rex_node: Some(RexNode::InputRef(self.index() as _)),
148        }
149    }
150}
risingwave_frontend/expr/input_ref.rs

risingwave_frontend/expr/
input_ref.rs
