risingwave_expr_impl/scalar/
vnode.rs

1// Copyright 2023 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::sync::Arc;
16
17use anyhow::Context;
18use itertools::Itertools;
19use risingwave_common::array::{ArrayBuilder, ArrayImpl, ArrayRef, DataChunk, I16ArrayBuilder};
20use risingwave_common::hash::VirtualNode;
21use risingwave_common::row::OwnedRow;
22use risingwave_common::types::{DataType, Datum};
23use risingwave_expr::expr::{
24    AsyncExpression, AsyncExpressionBoxExt, BoxedExpression, ExpressionInfo, SyncExpression,
25    SyncExpressionBoxExt, try_into_sync_exprs,
26};
27use risingwave_expr::{Result, build_function, expr_context};
28
29#[derive(Debug)]
30struct VnodeExpression<E> {
31    /// `Some` if it's from the first argument of user-facing function `VnodeUser` (`rw_vnode`),
32    /// `None` if it's from the internal function `Vnode`.
33    vnode_count: Option<usize>,
34
35    /// A list of expressions to get the distribution key columns. Typically `InputRef`.
36    children: Vec<E>,
37
38    /// Normally, we pass the distribution key indices to `VirtualNode::compute_xx` functions.
39    /// But in this case, all children columns are used to compute vnode. So we cache a vector of
40    /// all indices here and pass it later to reduce allocation.
41    all_indices: Vec<usize>,
42}
43
44#[build_function("vnode(...) -> int2")]
45fn build(_: DataType, children: Vec<BoxedExpression>) -> Result<BoxedExpression> {
46    let all_indices = (0..children.len()).collect();
47    match try_into_sync_exprs(children) {
48        Ok(children) => Ok(VnodeExpression {
49            vnode_count: None,
50            all_indices,
51            children,
52        }
53        .boxed()),
54        Err(children) => Ok(VnodeExpression {
55            vnode_count: None,
56            all_indices,
57            children,
58        }
59        .boxed()),
60    }
61}
62
63#[build_function("vnode_user(...) -> int2")]
64fn build_user(_: DataType, children: Vec<BoxedExpression>) -> Result<BoxedExpression> {
65    let mut children = children.into_iter();
66
67    let vnode_count = children
68        .next()
69        .unwrap() // always exist, argument number enforced in binder
70        .eval_const() // required to be constant
71        .context("the first argument (vnode count) must be a constant")?
72        .context("the first argument (vnode count) must not be NULL")?
73        .into_int32(); // always int32, casted during type inference
74
75    if !(1i32..=VirtualNode::MAX_COUNT as i32).contains(&vnode_count) {
76        return Err(anyhow::anyhow!(
77            "the first argument (vnode count) must be in range 1..={}",
78            VirtualNode::MAX_COUNT
79        )
80        .into());
81    }
82
83    let children = children.collect_vec();
84    let all_indices = (0..children.len()).collect();
85    match try_into_sync_exprs(children) {
86        Ok(children) => Ok(VnodeExpression {
87            vnode_count: Some(vnode_count.try_into().unwrap()),
88            all_indices,
89            children,
90        }
91        .boxed()),
92        Err(children) => Ok(VnodeExpression {
93            vnode_count: Some(vnode_count.try_into().unwrap()),
94            all_indices,
95            children,
96        }
97        .boxed()),
98    }
99}
100
101impl<E: ExpressionInfo> ExpressionInfo for VnodeExpression<E> {
102    fn return_type(&self) -> DataType {
103        DataType::Int16
104    }
105}
106
107macro_rules! eval_vnode {
108    ($mode:ident, $this:expr, $input:expr) => {{
109        let mut arrays = Vec::with_capacity($this.children.len());
110        for child in &$this.children {
111            arrays.push(risingwave_expr::forward!($mode, child, eval($input))?);
112        }
113        let input = DataChunk::new(arrays, $input.visibility().clone());
114
115        let vnodes = VirtualNode::compute_chunk(&input, &$this.all_indices, $this.vnode_count()?);
116        let mut builder = I16ArrayBuilder::new(input.capacity());
117        vnodes
118            .into_iter()
119            .for_each(|vnode| builder.append(Some(vnode.to_scalar())));
120        Ok(Arc::new(ArrayImpl::from(builder.finish())))
121    }};
122}
123
124macro_rules! eval_row_vnode {
125    ($mode:ident, $this:expr, $input:expr) => {{
126        let mut datums = Vec::with_capacity($this.children.len());
127        for child in &$this.children {
128            datums.push(risingwave_expr::forward!($mode, child, eval_row($input))?);
129        }
130        let input = OwnedRow::new(datums);
131
132        Ok(Some(
133            VirtualNode::compute_row(input, &$this.all_indices, $this.vnode_count()?)
134                .to_scalar()
135                .into(),
136        ))
137    }};
138}
139
140impl<E: SyncExpression> SyncExpression for VnodeExpression<E> {
141    fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
142        eval_vnode!(sync, self, input)
143    }
144
145    fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
146        eval_row_vnode!(sync, self, input)
147    }
148}
149
150impl<E: AsyncExpression> AsyncExpression for VnodeExpression<E> {
151    async fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
152        eval_vnode!(async, self, input)
153    }
154
155    async fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
156        eval_row_vnode!(async, self, input)
157    }
158}
159
160impl<E> VnodeExpression<E> {
161    fn vnode_count(&self) -> Result<usize> {
162        if let Some(vnode_count) = self.vnode_count {
163            Ok(vnode_count)
164        } else {
165            expr_context::vnode_count()
166        }
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use risingwave_common::array::{DataChunk, DataChunkTestExt};
173    use risingwave_common::row::Row;
174    use risingwave_expr::expr::build_from_pretty;
175    use risingwave_expr::expr_context::VNODE_COUNT;
176
177    #[tokio::test]
178    async fn test_vnode_expr_eval() {
179        let vnode_count = 32;
180        let expr = build_from_pretty("(vnode:int2 $0:int4 $0:int8 $0:varchar)");
181        let input = DataChunk::from_pretty(
182            "i  I  T
183             1  10 abc
184             2  32 def
185             3  88 ghi",
186        );
187
188        // test eval
189        let output = VNODE_COUNT::scope(vnode_count, expr.eval(&input))
190            .await
191            .unwrap();
192        for vnode in output.iter() {
193            let vnode = vnode.unwrap().into_int16();
194            assert!((0..vnode_count as i16).contains(&vnode));
195        }
196
197        // test eval_row
198        for row in input.rows() {
199            let result = VNODE_COUNT::scope(vnode_count, expr.eval_row(&row.to_owned_row()))
200                .await
201                .unwrap();
202            let vnode = result.unwrap().into_int16();
203            assert!((0..vnode_count as i16).contains(&vnode));
204        }
205    }
206}
risingwave_expr_impl/scalar/vnode.rs

risingwave_expr_impl/scalar/
vnode.rs
