From 419c85397c4b65dcd522e62a94f14ab6df96abae Mon Sep 17 00:00:00 2001
From: Bugen Zhao <i@bugenzhao.com>
Date: Tue, 8 Oct 2024 18:04:57 +0800
Subject: [PATCH] feat(expr): support any children in `rw_vnode` (#18405)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>
---
 e2e_test/batch/functions/internal.slt.part |  3 ++
 src/expr/impl/src/scalar/vnode.rs          | 34 ++++++++++++++++------
 2 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/e2e_test/batch/functions/internal.slt.part b/e2e_test/batch/functions/internal.slt.part
index fdb240c4c371a..52ddbcbc87a58 100644
--- a/e2e_test/batch/functions/internal.slt.part
+++ b/e2e_test/batch/functions/internal.slt.part
@@ -13,5 +13,8 @@ insert into t values ('aaa', 1, 1), ('bbb', 0, 2), ('ccc', 0, 5), ('ddd', 1, 4)
 statement ok
 select rw_vnode(_row_id) as vnode, _row_id from t;
 
+statement ok
+select rw_vnode(v2 + 114), rw_vnode(514) from t;
+
 statement ok
 drop table t
diff --git a/src/expr/impl/src/scalar/vnode.rs b/src/expr/impl/src/scalar/vnode.rs
index 7d44dfb0e03b1..1cb92cb036097 100644
--- a/src/expr/impl/src/scalar/vnode.rs
+++ b/src/expr/impl/src/scalar/vnode.rs
@@ -24,17 +24,21 @@ use risingwave_expr::{build_function, Result};
 
 #[derive(Debug)]
 struct VnodeExpression {
-    dist_key_indices: Vec<usize>,
+    /// A list of expressions to get the distribution key columns. Typically `InputRef`.
+    children: Vec<BoxedExpression>,
+
+    /// Normally, we pass the distribution key indices to `VirtualNode::compute_xx` functions.
+    /// But in this case, all children columns are used to compute vnode. So we cache a vector of
+    /// all indices here and pass it later to reduce allocation.
+    all_indices: Vec<usize>,
 }
 
 #[build_function("vnode(...) -> int2")]
 fn build(_: DataType, children: Vec<BoxedExpression>) -> Result<BoxedExpression> {
-    let dist_key_indices = children
-        .into_iter()
-        .map(|child| child.input_ref_index().unwrap())
-        .collect();
-
-    Ok(Box::new(VnodeExpression { dist_key_indices }))
+    Ok(Box::new(VnodeExpression {
+        all_indices: (0..children.len()).collect(),
+        children,
+    }))
 }
 
 #[async_trait::async_trait]
@@ -44,7 +48,13 @@ impl Expression for VnodeExpression {
     }
 
     async fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
-        let vnodes = VirtualNode::compute_chunk(input, &self.dist_key_indices, vnode_count());
+        let mut arrays = Vec::with_capacity(self.children.len());
+        for child in &self.children {
+            arrays.push(child.eval(input).await?);
+        }
+        let input = DataChunk::new(arrays, input.visibility().clone());
+
+        let vnodes = VirtualNode::compute_chunk(&input, &self.all_indices, vnode_count());
         let mut builder = I16ArrayBuilder::new(input.capacity());
         vnodes
             .into_iter()
@@ -53,8 +63,14 @@ impl Expression for VnodeExpression {
     }
 
     async fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
+        let mut datums = Vec::with_capacity(self.children.len());
+        for child in &self.children {
+            datums.push(child.eval_row(input).await?);
+        }
+        let input = OwnedRow::new(datums);
+
         Ok(Some(
-            VirtualNode::compute_row(input, &self.dist_key_indices, vnode_count())
+            VirtualNode::compute_row(input, &self.all_indices, vnode_count())
                 .to_scalar()
                 .into(),
         ))