Merge branch 'main' into perf/in-place-mutable-rewrite

Author: zhuqi-lucas

datafusion/functions-window/src/nth_value.rs

@@ -125,6 +125,14 @@ impl NthValue {
     }
 }
 
+fn validate_nth_value_n(n: i64) -> Result<i64> {
+    if n == i64::MIN {
+        return exec_err!("The second argument of nth_value must not be i64::MIN");
+    }
+
+    Ok(n)
+}
+
 static FIRST_VALUE_DOCUMENTATION: LazyLock<Documentation> = LazyLock::new(|| {
     Documentation::builder(
         DOC_SECTION_ANALYTICAL,
@@ -287,6 +295,7 @@ impl WindowUDFImpl for NthValue {
         .map(|v| get_signed_integer(&v))
         {
             Some(Ok(n)) => {
+                let n = validate_nth_value_n(n)?;
                 if partition_evaluator_args.is_reversed() {
                     -n
                 } else {
@@ -660,4 +669,24 @@ mod tests {
         )?;
         Ok(())
     }
+
+    #[test]
+    fn nth_value_i64_min_returns_error() {
+        let expr = Arc::new(Column::new("c3", 0)) as Arc<dyn PhysicalExpr>;
+        let n_value = Arc::new(Literal::new(ScalarValue::Int64(Some(i64::MIN))))
+            as Arc<dyn PhysicalExpr>;
+
+        let err = NthValue::nth()
+            .partition_evaluator(PartitionEvaluatorArgs::new(
+                &[expr, n_value],
+                &[Field::new("f", DataType::Int32, true).into()],
+                false,
+                false,
+            ))
+            .unwrap_err();
+
+        assert!(err.to_string().starts_with(
+            "Execution error: The second argument of nth_value must not be i64::MIN"
+        ));
+    }
 }

datafusion/physical-expr/src/expressions/try_cast.rs

@@ -119,6 +119,56 @@ impl PhysicalExpr for TryCastExpr {
         self.expr.fmt_sql(f)?;
         write!(f, " AS {:?})", self.cast_type)
     }
+
+    #[cfg(feature = "proto")]
+    fn try_to_proto(
+        &self,
+        ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
+    ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> {
+        use datafusion_proto_models::protobuf;
+
+        Ok(Some(protobuf::PhysicalExprNode {
+            expr_id: None,
+            expr_type: Some(protobuf::physical_expr_node::ExprType::TryCast(Box::new(
+                protobuf::PhysicalTryCastNode {
+                    expr: Some(Box::new(ctx.encode_child(&self.expr)?)),
+                    arrow_type: Some(self.cast_type().try_into()?),
+                },
+            ))),
+        }))
+    }
+}
+
+#[cfg(feature = "proto")]
+impl TryCastExpr {
+    /// Reconstruct a [`TryCastExpr`] from its protobuf representation.
+    pub fn try_from_proto(
+        node: &datafusion_proto_models::protobuf::PhysicalExprNode,
+        ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>,
+    ) -> Result<Arc<dyn PhysicalExpr>> {
+        use datafusion_physical_expr_common::expect_expr_variant;
+        use datafusion_physical_expr_common::physical_expr::proto_decode::require_proto_field;
+        use datafusion_proto_models::protobuf;
+
+        let try_cast = expect_expr_variant!(
+            node,
+            protobuf::physical_expr_node::ExprType::TryCast,
+            "TryCastExpr",
+        );
+        let expr = ctx.decode_required_expression(
+            try_cast.expr.as_deref(),
+            "TryCastExpr",
+            "expr",
+        )?;
+        let arrow_type = require_proto_field(
+            try_cast.arrow_type.as_ref(),
+            "TryCastExpr",
+            "arrow_type",
+        )?;
+        let cast_type: DataType = arrow_type.try_into()?;
+
+        Ok(Arc::new(TryCastExpr::new(expr, cast_type)))
+    }
 }
 
 /// Return a PhysicalExpression representing `expr` casted to
@@ -593,3 +643,143 @@ mod tests {
         Ok(())
     }
 }
+
+#[cfg(all(test, feature = "proto"))]
+mod proto_tests {
+    use super::*;
+    use crate::expressions::{Column, col};
+    use crate::proto_test_util::{
+        StubDecoder, StubEncoder, UnreachableDecoder, column_node,
+    };
+    use arrow::datatypes::Field;
+    use datafusion_common::DataFusionError;
+    use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx;
+    use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx;
+    use datafusion_proto_models::datafusion_common::ArrowType;
+    use datafusion_proto_models::protobuf::{
+        PhysicalExprNode, PhysicalTryCastNode, physical_expr_node,
+    };
+
+    fn try_cast_fixture() -> TryCastExpr {
+        let schema = Schema::new(vec![Field::new("a", DataType::Utf8, true)]);
+        TryCastExpr::new(col("a", &schema).unwrap(), DataType::Int32)
+    }
+
+    fn int32_arrow_type() -> ArrowType {
+        (&DataType::Int32).try_into().unwrap()
+    }
+
+    fn try_cast_node(
+        expr: Option<Box<PhysicalExprNode>>,
+        arrow_type: Option<ArrowType>,
+    ) -> PhysicalExprNode {
+        PhysicalExprNode {
+            expr_id: None,
+            expr_type: Some(physical_expr_node::ExprType::TryCast(Box::new(
+                PhysicalTryCastNode { expr, arrow_type },
+            ))),
+        }
+    }
+
+    #[test]
+    fn try_to_proto_encodes_try_cast_expr() {
+        let try_cast = try_cast_fixture();
+        let encoder = StubEncoder::ok();
+        let ctx = PhysicalExprEncodeCtx::new(&encoder);
+
+        let node = try_cast
+            .try_to_proto(&ctx)
+            .unwrap()
+            .expect("TryCastExpr should encode to Some(node)");
+
+        assert!(node.expr_id.is_none());
+        let try_cast_node = match node.expr_type {
+            Some(physical_expr_node::ExprType::TryCast(boxed)) => *boxed,
+            other => panic!("expected a TryCastExpr node, got {other:?}"),
+        };
+        assert!(try_cast_node.expr.is_some());
+
+        let arrow_type = try_cast_node
+            .arrow_type
+            .as_ref()
+            .expect("try cast type should be encoded");
+        let data_type: DataType = arrow_type.try_into().unwrap();
+        assert_eq!(data_type, DataType::Int32);
+    }
+
+    #[test]
+    fn try_to_proto_propagates_child_encode_error() {
+        let try_cast = try_cast_fixture();
+        let encoder = StubEncoder::failing_on(1);
+        let ctx = PhysicalExprEncodeCtx::new(&encoder);
+        let err = try_cast.try_to_proto(&ctx).unwrap_err();
+        assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1")));
+    }
+
+    #[test]
+    fn try_from_proto_decodes_try_cast_expr() {
+        let node =
+            try_cast_node(Some(Box::new(column_node("a"))), Some(int32_arrow_type()));
+        let schema = Schema::empty();
+        let decoder = StubDecoder::ok();
+        let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
+
+        let decoded = TryCastExpr::try_from_proto(&node, &ctx).unwrap();
+        let try_cast = decoded
+            .downcast_ref::<TryCastExpr>()
+            .expect("decoded expr should be a TryCastExpr");
+
+        assert_eq!(try_cast.cast_type(), &DataType::Int32);
+        assert!(try_cast.expr().downcast_ref::<Column>().is_some());
+    }
+
+    #[test]
+    fn try_from_proto_rejects_non_try_cast_node() {
+        let node = column_node("a");
+        let schema = Schema::empty();
+        let decoder = UnreachableDecoder;
+        let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
+
+        let err = TryCastExpr::try_from_proto(&node, &ctx).unwrap_err();
+        assert!(
+            matches!(err, DataFusionError::Internal(msg) if msg.contains("PhysicalExprNode is not a TryCastExpr"))
+        );
+    }
+
+    #[test]
+    fn try_from_proto_rejects_missing_expr() {
+        let node = try_cast_node(None, Some(int32_arrow_type()));
+        let schema = Schema::empty();
+        let decoder = UnreachableDecoder;
+        let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
+
+        let err = TryCastExpr::try_from_proto(&node, &ctx).unwrap_err();
+        assert!(
+            matches!(err, DataFusionError::Internal(msg) if msg.contains("TryCastExpr is missing required field 'expr'"))
+        );
+    }
+
+    #[test]
+    fn try_from_proto_rejects_missing_arrow_type() {
+        let node = try_cast_node(Some(Box::new(column_node("a"))), None);
+        let schema = Schema::empty();
+        let decoder = StubDecoder::ok();
+        let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
+
+        let err = TryCastExpr::try_from_proto(&node, &ctx).unwrap_err();
+        assert!(
+            matches!(err, DataFusionError::Internal(msg) if msg.contains("TryCastExpr is missing required field 'arrow_type'"))
+        );
+    }
+
+    #[test]
+    fn try_from_proto_propagates_child_decode_error() {
+        let node =
+            try_cast_node(Some(Box::new(column_node("a"))), Some(int32_arrow_type()));
+        let schema = Schema::empty();
+        let decoder = StubDecoder::failing_on(1);
+        let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
+        let err = TryCastExpr::try_from_proto(&node, &ctx).unwrap_err();
+        assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1")));
+    }
+}