dbtsai commented on code in PR #102:
URL:
https://github.com/apache/spark-connect-rust/pull/102#discussion_r4384178558
##########
crates/spark-connect/src/expression.rs:
##########
@@ -17,12 +18,38 @@ pub fn next_lambda_var_index() -> u32 {
LAMBDA_VAR_COUNTER.fetch_add(1, Ordering::SeqCst)
}
-/// The base Expression type, mirroring
`pyspark.sql.connect.expressions.Expression`.
+/// An immutable, shared expression, mirroring PySpark's base Expression type.
///
-/// All expressions are represented as variants of this enum. Each variant
-/// carries the data needed to construct the corresponding proto expression.
+/// Cloning a handle shares its node rather than copying the expression
subtree.
+/// Construct nodes with [`Expression::new`] and inspect them with
[`Expression::node`].
+#[derive(Clone)]
+pub struct Expression(Arc<ExpressionNode>);
+
+impl std::fmt::Debug for Expression {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ std::fmt::Debug::fmt(self.node(), f)
+ }
+}
+
+impl PartialEq for Expression {
+ fn eq(&self, other: &Self) -> bool {
+ // Compare values even for the same allocation: NaN must remain
unequal to itself.
+ self.node().eq(other.node())
+ }
+}
+
+impl From<ExpressionNode> for Expression {
Review Comment:
Added `Expression::literal`, `Expression::column`, and
`Expression::function(name, args)` and updated the common production and test
call sites. Less common variants use the existing `From<ExpressionNode>`
conversion through `.into()` or `Expression::from`. Distinct-function flags,
metadata, and bound-column plan IDs are preserved.
##########
crates/spark-connect/src/expression.rs:
##########
@@ -17,12 +18,38 @@ pub fn next_lambda_var_index() -> u32 {
LAMBDA_VAR_COUNTER.fetch_add(1, Ordering::SeqCst)
}
-/// The base Expression type, mirroring
`pyspark.sql.connect.expressions.Expression`.
+/// An immutable, shared expression, mirroring PySpark's base Expression type.
///
-/// All expressions are represented as variants of this enum. Each variant
-/// carries the data needed to construct the corresponding proto expression.
+/// Cloning a handle shares its node rather than copying the expression
subtree.
+/// Construct nodes with [`Expression::new`] and inspect them with
[`Expression::node`].
+#[derive(Clone)]
+pub struct Expression(Arc<ExpressionNode>);
+
+impl std::fmt::Debug for Expression {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ std::fmt::Debug::fmt(self.node(), f)
+ }
+}
+
+impl PartialEq for Expression {
+ fn eq(&self, other: &Self) -> bool {
+ // Compare values even for the same allocation: NaN must remain
unequal to itself.
+ self.node().eq(other.node())
+ }
+}
+
+impl From<ExpressionNode> for Expression {
+ fn from(node: ExpressionNode) -> Self {
+ Self::new(node)
+ }
+}
+
+/// The contents of one expression node.
+///
+/// Child expressions are shared [`Expression`] handles. Direct enum
construction
+/// and matching use this type; the public Column API continues to use
Expression.
#[derive(Debug, Clone, PartialEq)]
-pub enum Expression {
+pub enum ExpressionNode {
Review Comment:
Removed both extra boxes: `DirectShufflePartitionId` now takes an
`Expression`, and `CaseWhen::else_expr` is `Option<Expression>`. Updated the
CASE and shuffle construction/serialization paths; the selected golden tests
pass. The migration note also calls out these representation changes.
##########
scripts/check_doc_snippets.py:
##########
@@ -81,14 +83,15 @@ def check_block(code: str, tmpdir: Path, idx: int) ->
tuple[bool, str]:
indented = "\n".join(" " + ln if ln.strip() else ln for ln in
code.splitlines())
src = tmpdir / f"snippet_{idx}.rs"
src.write_text(HARNESS.format(body=indented))
- out = tmpdir / f"snippet_{idx}.out"
+ out = tmpdir / f"snippet_{idx}.rmeta"
r = subprocess.run(
[
"rustc",
"--edition",
"2021",
"--crate-type",
"bin",
+ "--emit=metadata",
Review Comment:
Removed the docs-checker change, lint workflow change, and tooling tests
from this PR's aggregate diff. The extracted CI fix is prepared as a separate
follow-up and is being kept locally.
##########
scripts/tests/test_check_doc_snippets.py:
##########
@@ -0,0 +1,115 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Offline doc-checker tests; no Rust toolchain or Spark server required."""
+
+import contextlib
+import importlib.util
+import io
+import subprocess
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+SCRIPT = Path(__file__).resolve().parents[1] / "check_doc_snippets.py"
+SPEC = importlib.util.spec_from_file_location("check_doc_snippets", SCRIPT)
+checker = importlib.util.module_from_spec(SPEC)
+SPEC.loader.exec_module(checker)
+
+
+class CheckDocSnippetsTests(unittest.TestCase):
+ def test_type_checks_without_linking_a_native_executable(self):
Review Comment:
The tooling tests are now outside this PR. In the local CI follow-up,
replaced the argv assertion with compiler-backed checks: a valid snippet with
an unresolved foreign symbol type-checks without native linking, while
type-error and moved-value snippets fail. Removed the test for the pre-existing
continue-after-failure behavior.
##########
crates/spark-connect/tests/expression_sharing.rs:
##########
@@ -0,0 +1,278 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Server-free checks for immutable expression sharing and unchanged
serialization.
+
+use std::ptr;
+
+use spark_connect::column::{col, lit, when, Column};
+use spark_connect::expression::{
+ Alias, CallFunctionWrapper, Expression, ExpressionNode, LiteralExpression,
SortOrder,
+};
+use spark_connect::functions;
+use spark_connect::types::DataType;
+use spark_connect::udf::{eval_type, CommonInlineUserDefinedFunctionExpression,
PythonUDFPayload};
+use spark_connect::window::{FrameBound, WindowSpec};
+use spark_connect_proto as proto;
+
+fn assert_shared(left: &Expression, right: &Expression) {
+ assert!(ptr::eq(left.node(), right.node()));
+}
+
+#[test]
+fn clones_and_arithmetic_chains_share_existing_nodes() {
+ let source = col("value");
+ let cloned = source.clone();
+ assert_shared(source.expression(), cloned.expression());
+ assert_shared(source.expression(), &source.expression().clone());
+
+ let mut columns = vec![source];
+ for _ in 0..128 {
+ columns.push(columns.last().unwrap().clone() + lit(1));
+ }
+ for pair in columns.windows(2) {
+ let ExpressionNode::UnresolvedFunction(function) =
pair[1].expression().node() else {
+ panic!("expected addition");
+ };
+ assert_eq!(function.name, "+");
+ assert_eq!(function.args.len(), 2);
+ assert_shared(&function.args[0], pair[0].expression());
+ }
+}
+
+#[test]
+fn aliases_and_function_arguments_share_reused_subtrees() {
+ let source = col("value") + lit(1);
+ let before = source.to_proto();
+ let left = source.clone().alias("left");
+ let right = source.clone().alias("right");
+ for alias in [&left, &right] {
+ let ExpressionNode::Alias(alias) = alias.expression().node() else {
+ panic!("expected alias");
+ };
+ assert_shared(&alias.child, source.expression());
+ }
+ let function = functions::call_function("greatest", vec![source.clone(),
source.clone()]);
+ let ExpressionNode::CallFunction(function) = function.expression().node()
else {
+ panic!("expected function call");
+ };
+ assert_eq!(function.arguments.len(), 2);
+ for argument in &function.arguments {
+ assert_shared(argument, source.expression());
+ }
+ assert_eq!(source.to_proto(), before);
+ assert_eq!(left.expression().render(), "(value + 1) AS left");
+ assert_eq!(right.expression().render(), "(value + 1) AS right");
+}
+
+#[test]
+fn case_branches_share_children_without_mutating_the_original() {
+ let condition = col("value").gt(lit(0));
+ let value = col("value") + lit(1);
+ let original = when(condition.clone(), value.clone());
+ let before = original.to_proto();
+ let extended = original.clone().when(col("value").lt(lit(0)), lit(-1));
+ let completed = original.clone().otherwise(lit(0));
+ let ExpressionNode::CaseWhen(original_case) = original.expression().node()
else {
+ panic!("expected CASE");
+ };
+ let ExpressionNode::CaseWhen(extended_case) = extended.expression().node()
else {
+ panic!("expected extended CASE");
+ };
+ let ExpressionNode::CaseWhen(completed_case) =
completed.expression().node() else {
+ panic!("expected completed CASE");
+ };
+ assert_eq!(original_case.branches.len(), 1);
+ assert!(original_case.else_expr.is_none());
+ assert_eq!(extended_case.branches.len(), 2);
+ assert!(extended_case.else_expr.is_none());
+ assert_eq!(completed_case.branches.len(), 1);
+ assert!(completed_case.else_expr.is_some());
+ for case in [original_case, extended_case, completed_case] {
+ assert_shared(&case.branches[0].0, condition.expression());
+ assert_shared(&case.branches[0].1, value.expression());
+ }
+ assert_eq!(original.to_proto(), before);
+ assert_eq!(
+ completed.expression().render(),
+ "CASE WHEN (value > 0) THEN (value + 1) ELSE 0 END"
+ );
+}
+
+#[test]
+fn window_builders_share_expressions_without_mutating_the_specification() {
+ let source = col("value") + lit(1);
+ let original = WindowSpec::new()
+ .partition_by(vec![source.expression().clone()])
+ .order_by(vec![SortOrder::desc_nulls_last(
+ source.expression().clone(),
+ )]);
+ let extended = original
+ .clone()
+ .partition_by(vec![col("other").expression().clone()])
+ .rows_between(FrameBound::UnboundedPreceding, FrameBound::CurrentRow);
+ assert_eq!(original.partition_spec.len(), 1);
+ assert!(original.frame_spec.is_none());
+ assert_eq!(extended.partition_spec.len(), 2);
+ assert!(extended.frame_spec.is_some());
+ assert_shared(&original.partition_spec[0], &extended.partition_spec[0]);
+ assert_shared(&original.order_spec[0].child,
&extended.order_spec[0].child);
+
+ let windowed = source.clone().over(extended);
+ let ExpressionNode::WindowExpression(window) =
windowed.expression().node() else {
+ panic!("expected window expression");
+ };
+ assert_shared(&window.window_function, source.expression());
+ assert_shared(&window.partition_spec[0], source.expression());
+ assert_shared(&window.order_spec[0].child, source.expression());
+}
+
+#[test]
+fn binary_literal_payload_survives_without_being_cloned() {
Review Comment:
Dropped both redundant ownership/lifetime tests. The existing UDF
serialization test already checks command bytes and eval type, so that coverage
remains without another node-sharing test.
##########
crates/spark-connect/tests/expression_sharing.rs:
##########
@@ -0,0 +1,278 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Server-free checks for immutable expression sharing and unchanged
serialization.
+
+use std::ptr;
+
+use spark_connect::column::{col, lit, when, Column};
+use spark_connect::expression::{
+ Alias, CallFunctionWrapper, Expression, ExpressionNode, LiteralExpression,
SortOrder,
+};
+use spark_connect::functions;
+use spark_connect::types::DataType;
+use spark_connect::udf::{eval_type, CommonInlineUserDefinedFunctionExpression,
PythonUDFPayload};
+use spark_connect::window::{FrameBound, WindowSpec};
+use spark_connect_proto as proto;
+
+fn assert_shared(left: &Expression, right: &Expression) {
+ assert!(ptr::eq(left.node(), right.node()));
+}
+
+#[test]
+fn clones_and_arithmetic_chains_share_existing_nodes() {
+ let source = col("value");
+ let cloned = source.clone();
+ assert_shared(source.expression(), cloned.expression());
+ assert_shared(source.expression(), &source.expression().clone());
+
+ let mut columns = vec![source];
+ for _ in 0..128 {
+ columns.push(columns.last().unwrap().clone() + lit(1));
+ }
+ for pair in columns.windows(2) {
+ let ExpressionNode::UnresolvedFunction(function) =
pair[1].expression().node() else {
+ panic!("expected addition");
+ };
+ assert_eq!(function.name, "+");
+ assert_eq!(function.args.len(), 2);
+ assert_shared(&function.args[0], pair[0].expression());
+ }
+}
+
+#[test]
+fn aliases_and_function_arguments_share_reused_subtrees() {
+ let source = col("value") + lit(1);
+ let before = source.to_proto();
+ let left = source.clone().alias("left");
+ let right = source.clone().alias("right");
+ for alias in [&left, &right] {
+ let ExpressionNode::Alias(alias) = alias.expression().node() else {
+ panic!("expected alias");
+ };
+ assert_shared(&alias.child, source.expression());
+ }
+ let function = functions::call_function("greatest", vec![source.clone(),
source.clone()]);
+ let ExpressionNode::CallFunction(function) = function.expression().node()
else {
+ panic!("expected function call");
+ };
+ assert_eq!(function.arguments.len(), 2);
+ for argument in &function.arguments {
+ assert_shared(argument, source.expression());
+ }
+ assert_eq!(source.to_proto(), before);
+ assert_eq!(left.expression().render(), "(value + 1) AS left");
+ assert_eq!(right.expression().render(), "(value + 1) AS right");
+}
+
+#[test]
+fn case_branches_share_children_without_mutating_the_original() {
+ let condition = col("value").gt(lit(0));
+ let value = col("value") + lit(1);
+ let original = when(condition.clone(), value.clone());
+ let before = original.to_proto();
+ let extended = original.clone().when(col("value").lt(lit(0)), lit(-1));
+ let completed = original.clone().otherwise(lit(0));
+ let ExpressionNode::CaseWhen(original_case) = original.expression().node()
else {
+ panic!("expected CASE");
+ };
+ let ExpressionNode::CaseWhen(extended_case) = extended.expression().node()
else {
+ panic!("expected extended CASE");
+ };
+ let ExpressionNode::CaseWhen(completed_case) =
completed.expression().node() else {
+ panic!("expected completed CASE");
+ };
+ assert_eq!(original_case.branches.len(), 1);
+ assert!(original_case.else_expr.is_none());
+ assert_eq!(extended_case.branches.len(), 2);
+ assert!(extended_case.else_expr.is_none());
+ assert_eq!(completed_case.branches.len(), 1);
+ assert!(completed_case.else_expr.is_some());
+ for case in [original_case, extended_case, completed_case] {
+ assert_shared(&case.branches[0].0, condition.expression());
+ assert_shared(&case.branches[0].1, value.expression());
+ }
+ assert_eq!(original.to_proto(), before);
+ assert_eq!(
+ completed.expression().render(),
+ "CASE WHEN (value > 0) THEN (value + 1) ELSE 0 END"
+ );
+}
+
+#[test]
+fn window_builders_share_expressions_without_mutating_the_specification() {
+ let source = col("value") + lit(1);
+ let original = WindowSpec::new()
+ .partition_by(vec![source.expression().clone()])
+ .order_by(vec![SortOrder::desc_nulls_last(
+ source.expression().clone(),
+ )]);
+ let extended = original
+ .clone()
+ .partition_by(vec![col("other").expression().clone()])
+ .rows_between(FrameBound::UnboundedPreceding, FrameBound::CurrentRow);
+ assert_eq!(original.partition_spec.len(), 1);
+ assert!(original.frame_spec.is_none());
+ assert_eq!(extended.partition_spec.len(), 2);
+ assert!(extended.frame_spec.is_some());
+ assert_shared(&original.partition_spec[0], &extended.partition_spec[0]);
+ assert_shared(&original.order_spec[0].child,
&extended.order_spec[0].child);
+
+ let windowed = source.clone().over(extended);
+ let ExpressionNode::WindowExpression(window) =
windowed.expression().node() else {
+ panic!("expected window expression");
+ };
+ assert_shared(&window.window_function, source.expression());
+ assert_shared(&window.partition_spec[0], source.expression());
+ assert_shared(&window.order_spec[0].child, source.expression());
+}
+
+#[test]
+fn binary_literal_payload_survives_without_being_cloned() {
+ let original =
Expression::new(ExpressionNode::Literal(LiteralExpression::Binary(vec![
+ 7;
+ 1024
+ ])));
+ let ExpressionNode::Literal(LiteralExpression::Binary(bytes)) =
original.node() else {
+ panic!("expected binary literal");
+ };
+ let address = bytes.as_ptr();
+ let cloned = original.clone();
+ assert_shared(&original, &cloned);
+ drop(original);
+
+ let ExpressionNode::Literal(LiteralExpression::Binary(bytes)) =
cloned.node() else {
+ panic!("expected binary literal");
+ };
+ assert_eq!(bytes.as_ptr(), address);
+ assert_eq!(bytes.as_slice(), &[7; 1024]);
+ assert_eq!(
+ cloned.to_proto(),
+ LiteralExpression::Binary(vec![7; 1024]).to_proto()
+ );
+}
+
+#[test]
+fn udf_clones_share_command_bytes_and_argument_lifetimes() {
+ let argument = col("value") + lit(1);
+ let argument_node = argument.expression().node() as *const ExpressionNode;
+ let payload = PythonUDFPayload::new(
+ DataType::Long,
+ eval_type::SQL_BATCHED_UDF,
+ vec![9; 1024],
+ "3.11".to_string(),
+ );
+ let command_address = payload.command.as_ptr();
+ let udf = CommonInlineUserDefinedFunctionExpression::new(
+ "test_udf".to_string(),
+ true,
+ vec![argument.expression().clone()],
+ payload,
+ );
+ let original =
Expression::new(ExpressionNode::CommonInlineUserDefinedFunction(Box::new(
+ udf,
+ )));
+ let cloned = original.clone();
+ assert_shared(&original, &cloned);
+ drop(original);
+ drop(argument);
+
+ let ExpressionNode::CommonInlineUserDefinedFunction(udf) = cloned.node()
else {
+ panic!("expected UDF");
+ };
+ assert_eq!(udf.python_udf.command.as_ptr(), command_address);
+ assert_eq!(udf.python_udf.command.as_slice(), &[9; 1024]);
+ assert!(ptr::eq(udf.arguments[0].node(), argument_node));
+ assert_eq!(udf.arguments[0].render(), "(value + 1)");
+ let
Some(proto::expression::ExprType::CommonInlineUserDefinedFunction(encoded)) =
+ cloned.to_proto().expr_type
+ else {
+ panic!("expected encoded UDF");
+ };
+ let
Some(proto::common_inline_user_defined_function::Function::PythonUdf(payload)) =
+ encoded.function
+ else {
+ panic!("expected encoded Python UDF");
+ };
+ assert_eq!(payload.command.as_ref(), &[9; 1024]);
+ assert_eq!(payload.eval_type, eval_type::SQL_BATCHED_UDF);
+}
+
+#[test]
+fn equality_remains_structural_and_does_not_shortcut_nan() {
+ let left = (col("value") + lit(1)).expression().clone();
+ let right = (col("value") + lit(1)).expression().clone();
+ assert!(!ptr::eq(left.node(), right.node()));
+ assert_eq!(left, right);
+ assert_ne!(left, (col("value") + lit(2)).expression().clone());
+
+ for literal in [
+ LiteralExpression::Float(f32::NAN),
+ LiteralExpression::Double(f64::NAN),
+ ] {
+ let expression: Expression = ExpressionNode::Literal(literal).into();
+ let cloned = expression.clone();
+ assert_shared(&expression, &cloned);
+ assert_ne!(expression, cloned);
+ let nested = Column::new(expression) + lit(1);
+ assert_ne!(nested.expression(), nested.clone().expression());
+ }
+}
+
+#[test]
+fn debug_render_and_proto_keep_the_existing_representation() {
+ let argument: Expression =
ExpressionNode::Literal(LiteralExpression::Integer(7)).into();
+ let wrapper = CallFunctionWrapper::new("f", vec![argument]);
+ let expected_proto = wrapper.to_proto();
+ let call =
Expression::new(ExpressionNode::CallFunction(Box::new(wrapper)));
+ let expected_debug = concat!(
+ "CallFunction(CallFunctionWrapper { function_name: \"f\", ",
+ "arguments: [Literal(Integer(7))] })"
+ );
+ assert_eq!(format!("{call:?}"), expected_debug);
+ assert_eq!(call.render(), expected_debug);
+ assert_eq!(call.to_proto(), expected_proto);
+
+ let source = col("value") + lit(1);
+ let alias = source.clone().alias("total");
+ assert_eq!(alias.expression().render(), "(value + 1) AS total");
+ assert_eq!(
+ alias.to_proto(),
+ Alias::new(source.expression().clone(), "total").to_proto()
+ );
+}
+
+#[test]
+fn shared_occurrences_expand_into_independent_proto_trees() {
+ let source = col("value") + lit(1);
+ let combined = source.clone() + source.clone();
+ let ExpressionNode::UnresolvedFunction(function) =
combined.expression().node() else {
+ panic!("expected addition");
+ };
+ assert_shared(&function.args[0], &function.args[1]);
+ let expected = ((col("value") + lit(1)) + (col("value") +
lit(1))).to_proto();
+ let mut encoded = combined.to_proto();
+ assert_eq!(encoded, expected);
+ let Some(proto::expression::ExprType::UnresolvedFunction(function)) = &mut
encoded.expr_type
+ else {
+ panic!("expected encoded addition");
+ };
+ assert_eq!(function.arguments.len(), 2);
+ assert_eq!(function.arguments[0], source.to_proto());
+ assert_eq!(function.arguments[1], source.to_proto());
+ function.arguments[0].expr_type = None;
Review Comment:
Removed the owned-proto mutation assertions. Kept the equality checks that
verify a shared expression occurrence expands into two complete protobuf
subtrees.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
