From 67df70ecd0c03dfa67ee1bb2ba24b6c9900767a6 Mon Sep 17 00:00:00 2001
From: Bugen Zhao <i@bugenzhao.com>
Date: Tue, 8 Oct 2024 18:16:43 +0800
Subject: [PATCH] refine

Signed-off-by: Bugen Zhao <i@bugenzhao.com>
---
 src/expr/impl/src/scalar/vnode.rs            | 8 +++++++-
 src/frontend/src/expr/type_inference/func.rs | 2 +-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/expr/impl/src/scalar/vnode.rs b/src/expr/impl/src/scalar/vnode.rs
index aefedf8203ea7..ce8d992aecbfd 100644
--- a/src/expr/impl/src/scalar/vnode.rs
+++ b/src/expr/impl/src/scalar/vnode.rs
@@ -25,6 +25,8 @@ use risingwave_expr::{build_function, expr_context, Result};
 
 #[derive(Debug)]
 struct VnodeExpression {
+    /// `Some` if it's from the first argument of user-facing function `VnodeUser` (`rw_vnode`),
+    /// `None` if it's from the internal function `Vnode`.
     vnode_count: Option<usize>,
 
     /// A list of expressions to get the distribution key columns. Typically `InputRef`.
@@ -55,8 +57,12 @@ fn build_user(_: DataType, children: Vec<BoxedExpression>) -> Result<BoxedExpres
         .eval_const() // required to be constant
         .ok()
         .flatten() // required to be non-null
-        .context("rw_vnode expects the first argument to be a non-null constant")?
+        .context("the first argument (vnode count) must be a non-null constant")?
         .into_int16() as usize;
+    if vnode_count == 0 {
+        return Err(anyhow::anyhow!("the first argument (vnode count) must not be zero").into());
+    }
+
     let children = children.collect_vec();
 
     Ok(Box::new(VnodeExpression {
diff --git a/src/frontend/src/expr/type_inference/func.rs b/src/frontend/src/expr/type_inference/func.rs
index 0660eee405532..883ce96fd6af8 100644
--- a/src/frontend/src/expr/type_inference/func.rs
+++ b/src/frontend/src/expr/type_inference/func.rs
@@ -662,7 +662,7 @@ fn infer_type_for_special(
         }
         ExprType::VnodeUser => {
             ensure_arity!("vnode", 2 <= | inputs |);
-            inputs[0].cast_implicit_mut(DataType::Int16)?;
+            inputs[0].cast_explicit_mut(DataType::Int16)?;
             Ok(Some(VirtualNode::RW_TYPE))
         }
         ExprType::Greatest | ExprType::Least => {