This is an automated email from the ASF dual-hosted git repository.
HyukjinKwon pushed a commit to branch branch-4.2
in repository https://gitbox.apache.org/repos/asf/spark-connect-rust.git
The following commit(s) were added to refs/heads/branch-4.2 by this push:
new b768e97 [SPARK-59110] Fix PySpark functions parity: ColumnOrName
string args, agg dict, variadic concat_ws/format_string
b768e97 is described below
commit b768e97bd9ec8cb31b07612898decff57bb7ea0c
Author: Hyukjin Kwon <[email protected]>
AuthorDate: Sat Aug 29 21:53:07 2026 +0900
[SPARK-59110] Fix PySpark functions parity: ColumnOrName string args, agg
dict, variadic concat_ws/format_string
A behavioral parity sweep against real PySpark 4.2.0 (functions called with
string column names, plus a 51-op method sweep) found systemic gaps the
transport-level suite never exercised.
1. functions ColumnOrName parity: the shared to_column coercion mapped a
Python
str to a string LITERAL, but PySpark treats a bare str as a COLUMN NAME,
so
~130 functions were wrong (sqrt("v") errored; upper("v"),
struct("id","g")
silently used literals). _create_wrapper now resolves column-position
str args
via col(), driven by a _FN_COL_ARGS map derived from real PySpark 4.2.0
signatures; lit() and true literal params (date_format format,
regexp_replace
pattern, ...) stay literals.
2. GroupedData.agg now accepts the dict form agg({"col": "func"}).
3. concat_ws(sep, *cols) and format_string(format, *cols) dropped extra
columns
(core took a single Column; dispatch forwarded only args[0]); made
variadic
like concat.
Verified on a live Spark Connect server: functions string==Column sweep 0
divergences (was 73 errors + ~50 silently wrong), 51-op method sweep 51/51,
offline suite 232 passed.
Co-authored-by: Isaac <[email protected]>
---
crates/pyspark-rs/src/dispatch_generated.rs | 10 +-
crates/pyspark-rs/src/group.rs | 26 +++-
crates/spark-connect/src/functions.rs | 24 +++-
python/pyspark/sql/functions.py | 188 +++++++++++++++++++++++++++-
python/tests/test_dropin_offline.py | 39 ++++++
5 files changed, 270 insertions(+), 17 deletions(-)
diff --git a/crates/pyspark-rs/src/dispatch_generated.rs
b/crates/pyspark-rs/src/dispatch_generated.rs
index 4780012..0951bb0 100644
--- a/crates/pyspark-rs/src/dispatch_generated.rs
+++ b/crates/pyspark-rs/src/dispatch_generated.rs
@@ -249,10 +249,7 @@ pub fn call_builtin(name: &str, args:
Vec<spark_connect::column::Column>) -> PyR
if args.is_empty() { return
Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>("Missing required column
argument")); }
Ok(spark_funcs::collect_set(args[0].clone()))
},
- "concat_ws" => {
- if args.is_empty() { return
Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>("Missing required column
argument")); }
- Ok(spark_funcs::concat_ws(args[0].clone()))
- },
+ "concat_ws" => Ok(spark_funcs::concat_ws(args.clone())),
"cos" => {
if args.is_empty() { return
Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>("Missing required column
argument")); }
Ok(spark_funcs::cos(args[0].clone()))
@@ -357,10 +354,7 @@ pub fn call_builtin(name: &str, args:
Vec<spark_connect::column::Column>) -> PyR
if args.is_empty() { return
Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>("Missing required column
argument")); }
if args.len() >= 2 {
Ok(spark_funcs::floor_scale(args[0].clone(), args[1].clone())) } else {
Ok(spark_funcs::floor(args[0].clone())) }
},
- "format_string" => {
- if args.is_empty() { return
Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>("Missing required column
argument")); }
- Ok(spark_funcs::format_string(args[0].clone()))
- },
+ "format_string" => Ok(spark_funcs::format_string(args.clone())),
"from_unixtime" => {
if args.is_empty() { return
Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>("Missing required column
argument")); }
Ok(spark_funcs::from_unixtime(args[0].clone()))
diff --git a/crates/pyspark-rs/src/group.rs b/crates/pyspark-rs/src/group.rs
index e308af4..d4270a7 100644
--- a/crates/pyspark-rs/src/group.rs
+++ b/crates/pyspark-rs/src/group.rs
@@ -21,9 +21,33 @@ impl PyGroupedData {
#[pymethods]
impl PyGroupedData {
- /// Aggregate with expressions (accepts Column objects).
+ /// Aggregate with expressions (accepts Column objects), or the dict form
+ /// `agg({"col": "aggfunc"})` (PySpark parity), e.g. `agg({"age": "max"})`.
#[pyo3(signature = (*cols))]
fn agg(&self, _py: Python<'_>, cols: Vec<Bound<'_, PyAny>>) ->
PyResult<PyDataFrame> {
+ // Dict form: {column_name: aggregate_function_name} ->
func(col(name)) per entry.
+ if cols.len() == 1 {
+ if let Ok(dict) = cols[0].cast::<pyo3::types::PyDict>() {
+ let mut exprs = vec![];
+ for (k, v) in dict.iter() {
+ let col_name: String = k.extract()?;
+ let func_name: String = v.extract()?;
+ let c = spark_connect::functions::col(&col_name);
+ // Build func(col) as an unresolved function resolved
server-side, so
+ // aggregate names (max/min/sum/avg/count/...) all work
regardless of
+ // whether they are in the generated dispatch.
+ let agg_expr =
spark_connect::expression::Expression::UnresolvedFunction(
+ spark_connect::expression::UnresolvedFunction::new(
+ func_name,
+ vec![c.expression().clone()],
+ ),
+ );
+ exprs.push(agg_expr);
+ }
+ let df = self.grouped_data.agg(exprs);
+ return Ok(PyDataFrame::new(df));
+ }
+ }
let mut exprs = vec![];
for col in cols {
let rust_col = to_column(&col)?;
diff --git a/crates/spark-connect/src/functions.rs
b/crates/spark-connect/src/functions.rs
index 1fbe9a7..cb3dae8 100644
--- a/crates/spark-connect/src/functions.rs
+++ b/crates/spark-connect/src/functions.rs
@@ -781,9 +781,15 @@ pub fn collect_set(col1: Column) -> Column {
func("collect_set", vec![col1.expression().clone()])
}
-/// Mirrors `pyspark.sql.functions.concat_ws`.
-pub fn concat_ws(col1: Column) -> Column {
- func("concat_ws", vec![col1.expression().clone()])
+/// Mirrors `pyspark.sql.functions.concat_ws(sep, *cols)`: the separator
literal
+/// followed by a variadic number of columns. Variadic like `concat` so extra
+/// columns are not dropped.
+pub fn concat_ws<C: Into<Column>>(cols: impl IntoIterator<Item = C>) -> Column
{
+ let cols: Vec<Column> = cols.into_iter().map(Into::into).collect();
+ func(
+ "concat_ws",
+ cols.iter().map(|c| c.expression().clone()).collect(),
+ )
}
/// Mirrors `pyspark.sql.functions.cos`.
@@ -929,9 +935,15 @@ pub fn floor_scale(col1: Column, scale: Column) -> Column {
)
}
-/// Mirrors `pyspark.sql.functions.format_string`.
-pub fn format_string(col1: Column) -> Column {
- func("format_string", vec![col1.expression().clone()])
+/// Mirrors `pyspark.sql.functions.format_string(format, *cols)`: the format
+/// literal followed by a variadic number of columns. Variadic like `concat` so
+/// extra columns are not dropped.
+pub fn format_string<C: Into<Column>>(cols: impl IntoIterator<Item = C>) ->
Column {
+ let cols: Vec<Column> = cols.into_iter().map(Into::into).collect();
+ func(
+ "format_string",
+ cols.iter().map(|c| c.expression().clone()).collect(),
+ )
}
/// Mirrors `pyspark.sql.functions.from_unixtime`.
diff --git a/python/pyspark/sql/functions.py b/python/pyspark/sql/functions.py
index 4b62122..8c4cdf2 100644
--- a/python/pyspark/sql/functions.py
+++ b/python/pyspark/sql/functions.py
@@ -59,10 +59,194 @@ def _wrap(obj):
"""Wrap the result from _call_function as a Column."""
return obj
+# Auto-derived from PySpark 4.2.0 signatures: for generated functions, which
positional
+# args are ColumnOrName (a bare str is a COLUMN NAME). Functions NOT listed
here have
+# ALL args as columns. Value = (column_positions, variadic_col_start_or_None).
+_FN_COL_ARGS = {
+ 'any_value': ((0,), None),
+ 'approxCountDistinct': ((0,), None),
+ 'approx_count_distinct': ((0,), None),
+ 'approx_percentile': ((0,), None),
+ 'array_append': ((0,), None),
+ 'array_contains': ((0,), None),
+ 'array_insert': ((0, 1), None),
+ 'array_join': ((0,), None),
+ 'array_position': ((0,), None),
+ 'array_prepend': ((0,), None),
+ 'array_remove': ((0,), None),
+ 'array_sort': ((0,), None),
+ 'assert_true': ((0,), None),
+ 'bround': ((0,), None),
+ 'bucket': ((1,), None),
+ 'ceil': ((0,), None),
+ 'ceiling': ((0,), None),
+ 'collate': ((0,), None),
+ 'concat_ws': ((), 1),
+ 'conv': ((0,), None),
+ 'convert_timezone': ((2,), None),
+ 'count_min_sketch': ((0,), None),
+ 'curdate': ((), None),
+ 'current_catalog': ((), None),
+ 'current_database': ((), None),
+ 'current_date': ((), None),
+ 'current_path': ((), None),
+ 'current_schema': ((), None),
+ 'current_time': ((), None),
+ 'current_timestamp': ((), None),
+ 'current_timezone': ((), None),
+ 'current_user': ((), None),
+ 'date_format': ((0,), None),
+ 'date_part': ((1,), None),
+ 'date_trunc': ((1,), None),
+ 'datepart': ((1,), None),
+ 'decode': ((0,), None),
+ 'dense_rank': ((), None),
+ 'e': ((), None),
+ 'element_at': ((0,), None),
+ 'encode': ((0,), None),
+ 'extract': ((1,), None),
+ 'first': ((0,), None),
+ 'first_value': ((0,), None),
+ 'floor': ((0,), None),
+ 'format_number': ((0,), None),
+ 'format_string': ((), 1),
+ 'from_csv': ((0,), None),
+ 'from_json': ((0,), None),
+ 'from_unixtime': ((0,), None),
+ 'from_utc_timestamp': ((0,), None),
+ 'from_xml': ((0,), None),
+ 'get_json_object': ((0,), None),
+ 'histogram_numeric': ((0,), None),
+ 'hll_sketch_agg': ((0,), None),
+ 'hll_union': ((0, 1), None),
+ 'hll_union_agg': ((0,), None),
+ 'ilike': ((0, 1), None),
+ 'input_file_block_length': ((), None),
+ 'input_file_block_start': ((), None),
+ 'input_file_name': ((), None),
+ 'instr': ((0,), None),
+ 'kll_merge_agg_bigint': ((0,), None),
+ 'kll_merge_agg_double': ((0,), None),
+ 'kll_merge_agg_float': ((0,), None),
+ 'kll_sketch_agg_bigint': ((0,), None),
+ 'kll_sketch_agg_double': ((0,), None),
+ 'kll_sketch_agg_float': ((0,), None),
+ 'lag': ((0,), None),
+ 'last': ((0,), None),
+ 'last_value': ((0,), None),
+ 'lead': ((0,), None),
+ 'levenshtein': ((0, 1), None),
+ 'like': ((0, 1), None),
+ 'listagg': ((0,), None),
+ 'listagg_distinct': ((0,), None),
+ 'localtimestamp': ((), None),
+ 'locate': ((1,), None),
+ 'lpad': ((0,), None),
+ 'map_contains_key': ((0,), None),
+ 'max_by': ((0, 1), None),
+ 'min_by': ((0, 1), None),
+ 'mode': ((0,), None),
+ 'monotonically_increasing_id': ((), None),
+ 'months_between': ((0, 1), None),
+ 'next_day': ((0,), None),
+ 'now': ((), None),
+ 'nth_value': ((0,), None),
+ 'ntile': ((), None),
+ 'percent_rank': ((), None),
+ 'percentile': ((0,), None),
+ 'percentile_approx': ((0,), None),
+ 'pi': ((), None),
+ 'raise_error': ((), None),
+ 'rand': ((), None),
+ 'randn': ((), None),
+ 'randstr': ((), None),
+ 'rank': ((), None),
+ 'regexp_extract': ((0,), None),
+ 'regexp_extract_all': ((0, 1), None),
+ 'regexp_instr': ((0, 1), None),
+ 'regexp_replace': ((0,), None),
+ 'round': ((0,), None),
+ 'row_number': ((), None),
+ 'rpad': ((0,), None),
+ 'schema_of_csv': ((), None),
+ 'schema_of_json': ((), None),
+ 'schema_of_xml': ((), None),
+ 'session_user': ((), None),
+ 'session_window': ((0,), None),
+ 'shiftLeft': ((0,), None),
+ 'shiftRight': ((0,), None),
+ 'shiftRightUnsigned': ((0,), None),
+ 'shiftleft': ((0,), None),
+ 'shiftright': ((0,), None),
+ 'shiftrightunsigned': ((0,), None),
+ 'shuffle': ((0,), None),
+ 'sort_array': ((0,), None),
+ 'spark_partition_id': ((), None),
+ 'split': ((0, 2), None),
+ 'string_agg': ((0,), None),
+ 'string_agg_distinct': ((0,), None),
+ 'substring_index': ((0,), None),
+ 'theta_sketch_agg': ((0,), None),
+ 'theta_union': ((0, 1), None),
+ 'theta_union_agg': ((0,), None),
+ 'time_bucket': ((1,), None),
+ 'timestamp_add': ((1, 2), None),
+ 'timestamp_diff': ((1, 2), None),
+ 'to_avro': ((0,), None),
+ 'to_csv': ((0,), None),
+ 'to_date': ((0,), None),
+ 'to_json': ((0,), None),
+ 'to_timestamp': ((0,), None),
+ 'to_utc_timestamp': ((0,), None),
+ 'to_xml': ((0,), None),
+ 'translate': ((0,), None),
+ 'trunc': ((0,), None),
+ 'try_to_date': ((0,), None),
+ 'try_variant_get': ((0,), None),
+ 'tuple_intersection_agg_double': ((0,), None),
+ 'tuple_intersection_agg_integer': ((0,), None),
+ 'tuple_intersection_double': ((0, 1), None),
+ 'tuple_intersection_integer': ((0, 1), None),
+ 'tuple_intersection_theta_double': ((0, 1), None),
+ 'tuple_intersection_theta_integer': ((0, 1), None),
+ 'tuple_sketch_agg_double': ((0, 1), None),
+ 'tuple_sketch_agg_integer': ((0, 1), None),
+ 'tuple_sketch_summary_double': ((0,), None),
+ 'tuple_sketch_summary_integer': ((0,), None),
+ 'tuple_union_agg_double': ((0,), None),
+ 'tuple_union_agg_integer': ((0,), None),
+ 'tuple_union_double': ((0, 1), None),
+ 'tuple_union_integer': ((0, 1), None),
+ 'tuple_union_theta_double': ((0, 1), None),
+ 'tuple_union_theta_integer': ((0, 1), None),
+ 'uniform': ((), None),
+ 'unix_timestamp': ((0,), None),
+ 'user': ((), None),
+ 'uuid': ((), None),
+ 'variant_get': ((0,), None),
+ 'version': ((), None),
+ 'when': ((), None),
+}
+
def _create_wrapper(fname):
- """Create a wrapper function that calls _call_function."""
+ """Create a wrapper for a generated Spark SQL function.
+
+ PySpark's ColumnOrName contract: a bare ``str`` argument is a COLUMN NAME,
so
+ column-position ``str`` args are resolved via col() (``_to_col``);
literal-position
+ args (a format/pattern string, an int count, ...) are passed through
unchanged.
+ Functions absent from ``_FN_COL_ARGS`` have every positional arg as a
column.
+ """
+ override = _FN_COL_ARGS.get(fname)
def wrapper(*args):
- return _wrap(_call_function(fname, *[_unwrap(a) for a in args]))
+ if override is None:
+ conv = [_to_col(a) for a in args]
+ else:
+ cols, vstart = override
+ conv = [
+ _to_col(a) if (i in cols or (vstart is not None and i >=
vstart)) else a
+ for i, a in enumerate(args)
+ ]
+ return _wrap(_call_function(fname, *conv))
wrapper.__doc__ = f"Auto-generated wrapper for {fname}"
wrapper.__name__ = fname
return wrapper
diff --git a/python/tests/test_dropin_offline.py
b/python/tests/test_dropin_offline.py
index fde19af..2283de7 100644
--- a/python/tests/test_dropin_offline.py
+++ b/python/tests/test_dropin_offline.py
@@ -2449,3 +2449,42 @@ def
test_show_has_no_required_args_and_defaults_n_to_20():
)
def test_keyword_param_names_match_pyspark(cls, method, param):
assert param in inspect.signature(getattr(cls, method)).parameters
+
+
+# --- SPARK-XXXXX: string function args are column names (ColumnOrName parity)
---
+def test_fn_col_args_spec_classification():
+ from pyspark.sql import functions as F
+ spec = F._FN_COL_ARGS
+ # literal-param functions: only the listed positions are columns
+ assert spec["date_format"] == ((0,), None)
+ assert spec["regexp_replace"] == ((0,), None)
+ assert spec["split"] == ((0, 2), None)
+ assert spec["substring_index"] == ((0,), None)
+ assert spec["concat_ws"] == ((), 1) # sep literal, *cols variadic
columns
+ assert spec["current_date"] == ((), None) # no column args
+ # pure-column functions are absent -> default: every arg is a column
+ for name in ("sqrt", "upper", "abs", "mean", "struct", "array",
"coalesce"):
+ assert name not in spec, name
+
+
+def test_to_col_coerces_str_to_column():
+ from pyspark.sql import functions as F
+ col_cls = type(F.col("x"))
+ # a bare str resolves to a Column via col()
+ assert isinstance(F._to_col("x"), col_cls)
+ # an existing Column passes through unchanged
+ c = F.col("y")
+ assert F._to_col(c) is c
+
+
+def test_generated_fn_accepts_str_column_name():
+ from pyspark.sql import functions as F
+ col_cls = type(F.col("x"))
+ # string-name and Column forms build identical column expressions (lazy,
no server)
+ for build in (lambda a: F.sqrt(a), lambda a: F.upper(a), lambda a:
F.abs(a)):
+ assert isinstance(build("v"), col_cls)
+ assert repr(build("v")) == repr(build(F.col("v")))
+ # multi-col string args reference columns
+ assert repr(F.struct("id", "g")) == repr(F.struct(F.col("id"), F.col("g")))
+ # literal-position args stay literal (date_format format, regexp pattern)
+ assert "yyyy" in repr(F.date_format(F.col("d"), "yyyy"))
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
