This is an automated email from the ASF dual-hosted git repository.
HyukjinKwon pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark-connect-rust.git
The following commit(s) were added to refs/heads/master by this push:
new f3914d2 [SPARK-59116] Complete PySpark 4.2.0 behavioral parity for
the Rust drop-in
f3914d2 is described below
commit f3914d2ae83528b326a1b6cba10085122386984b
Author: Hyukjin Kwon <[email protected]>
AuthorDate: Sun Aug 30 07:11:00 2026 +0900
[SPARK-59116] Complete PySpark 4.2.0 behavioral parity for the Rust drop-in
Drove an exhaustive differential harness (real PySpark 4.2.0 Connect client
vs
the Rust drop-in, same ops on the same live server) to zero: 143/143
identical.
The native Rust engine (crates/spark-connect) was already correct; these
were
binding-layer gaps (crates/pyspark-rs + python/pyspark) plus a few core
function-signature extensions.
- na.fill/fillna accept float/string/bool and the {col: value} dict form
(were
int-only); single-field DDL schema strings parse as a schema;
createDataFrame
accepts empty data with an explicit schema.
- Window subsystem: Window.partitionBy/orderBy coerced a bare string to a
literal (so windows partitioned/ordered by a constant), and the ROWS frame
boundary had the wrong type -> row_number/rank/dense_rank/cume_dist/
percent_rank/lag/lead/nth_value and sum over rows/rangeBetween now
correct.
- Chained when().when().otherwise() (the free function built an
UnresolvedFunction instead of a CaseWhen, dropping earlier branches).
- lag/lead accept offset/default; json_tuple is variadic over literal field
names.
- Parameterized spark.sql(query, args=[...] | {...}); struct/named_struct/
from_json return a Row (not a dict) for struct-typed values;
createDataFrame
with array<struct> data.
Verified: differential harness 143/143 (zero divergences), three behavioral
batteries 32/32 + 83/83 + 56/56, offline suite 236 passed, cargo fmt/build
clean,
ruff clean.
Co-authored-by: Isaac <[email protected]>
---
crates/pyspark-rs/src/dataframe.rs | 26 ++++++--
crates/pyspark-rs/src/dispatch_generated.rs | 15 +----
crates/pyspark-rs/src/row.rs | 13 ++--
crates/pyspark-rs/src/session.rs | 52 ++++++++++++----
crates/pyspark-rs/src/window.rs | 75 ++++++++++-------------
crates/spark-connect/src/expression.rs | 6 +-
crates/spark-connect/src/functions.rs | 29 +++++----
crates/spark-connect/src/plan.rs | 22 ++++++-
crates/spark-connect/src/session.rs | 83 ++++++++++++++++++++++----
crates/spark-connect/src/types.rs | 32 ++++++++--
crates/spark-connect/tests/functions_golden.rs | 6 +-
crates/spark-connect/tests/window_golden.rs | 4 +-
python/pyspark/sql/functions.py | 21 +++++++
python/tests/test_dropin_offline.py | 4 ++
14 files changed, 276 insertions(+), 112 deletions(-)
diff --git a/crates/pyspark-rs/src/dataframe.rs
b/crates/pyspark-rs/src/dataframe.rs
index 5840d18..00be354 100644
--- a/crates/pyspark-rs/src/dataframe.rs
+++ b/crates/pyspark-rs/src/dataframe.rs
@@ -1346,6 +1346,7 @@ impl PyDataFrame {
fn na(&self) -> PyDataFrameNaFunctions {
PyDataFrameNaFunctions {
inner: self.dataframe.na(),
+ df: self.dataframe.clone(),
}
}
@@ -1464,6 +1465,7 @@ impl PyDataFrame {
#[pyclass(name = "DataFrameNaFunctions")]
pub struct PyDataFrameNaFunctions {
inner: spark_connect::group::NaFunctions,
+ df: spark_connect::dataframe::DataFrame,
}
#[pymethods]
@@ -1481,12 +1483,28 @@ impl PyDataFrameNaFunctions {
PyDataFrame::new(self.inner.drop(how.as_deref(), thresh, subset_refs))
}
+ /// Fill nulls; `value` may be a scalar (float / int / str / bool,
optionally over
+ /// `subset`) or a {column: value} mapping. Alias of `DataFrame.fillna`.
PySpark fills
+ /// only columns whose type matches the value's type (server-side).
#[pyo3(signature = (value, subset=None))]
- fn fill(&self, value: i64, subset: Option<Vec<String>>) -> PyDataFrame {
- let subset_refs: Option<Vec<&str>> = subset
+ fn fill(
+ &self,
+ value: &Bound<'_, PyAny>,
+ subset: Option<&Bound<'_, PyAny>>,
+ ) -> PyResult<PyDataFrame> {
+ if let Ok(d) = value.cast::<PyDict>() {
+ let mut pairs = Vec::with_capacity(d.len());
+ for (k, v) in d.iter() {
+ pairs.push((k.extract::<String>()?,
crate::session::py_to_value(&v)?));
+ }
+ return Ok(PyDataFrame::new(self.df.fillna_map(pairs)));
+ }
+ let val = crate::session::py_to_value(value)?;
+ let owned = to_subset(subset)?;
+ let refs: Option<Vec<&str>> = owned
.as_ref()
- .map(|v| v.iter().map(String::as_str).collect());
- PyDataFrame::new(self.inner.fill(value, subset_refs))
+ .map(|v| v.iter().map(|s| s.as_str()).collect());
+ Ok(PyDataFrame::new(self.df.fillna_value(val, refs)))
}
/// Replace values, mirroring `DataFrameNaFunctions.replace`:
diff --git a/crates/pyspark-rs/src/dispatch_generated.rs
b/crates/pyspark-rs/src/dispatch_generated.rs
index 0951bb0..6523d54 100644
--- a/crates/pyspark-rs/src/dispatch_generated.rs
+++ b/crates/pyspark-rs/src/dispatch_generated.rs
@@ -431,10 +431,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::kurtosis(args[0].clone()))
},
- "lag" => {
- if args.is_empty() { return
Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>("Missing required column
argument")); }
- Ok(spark_funcs::lag(args[0].clone()))
- },
+ "lag" => Ok(spark_funcs::lag(args.clone())),
"last" => {
if args.is_empty() { return
Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>("Missing required column
argument")); }
Ok(spark_funcs::last(args[0].clone()))
@@ -451,10 +448,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::lcase(args[0].clone()))
},
- "lead" => {
- if args.is_empty() { return
Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>("Missing required column
argument")); }
- Ok(spark_funcs::lead(args[0].clone()))
- },
+ "lead" => Ok(spark_funcs::lead(args.clone())),
"length" => {
if args.is_empty() { return
Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>("Missing required column
argument")); }
Ok(spark_funcs::length(args[0].clone()))
@@ -1371,10 +1365,7 @@ pub fn call_builtin(name: &str, args:
Vec<spark_connect::column::Column>) -> PyR
if args.len() < 2 { return
Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(format!("Missing required
arguments for instr: expected at least 2, got {}", args.len()))); }
Ok(spark_funcs::instr(args[0].clone(), args[1].clone()))
},
- "json_tuple" => {
- if args.len() < 2 { return
Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(format!("Missing required
arguments for json_tuple: expected at least 2, got {}", args.len()))); }
- Ok(spark_funcs::json_tuple(args[0].clone(), args[1].clone()))
- },
+ "json_tuple" => Ok(spark_funcs::json_tuple(args.clone())),
"least" => {
if args.len() < 2 { return
Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(format!("Missing required
arguments for least: expected at least 2, got {}", args.len()))); }
Ok(spark_funcs::least(args[0].clone(), args[1].clone()))
diff --git a/crates/pyspark-rs/src/row.rs b/crates/pyspark-rs/src/row.rs
index 9387a62..4f4b4da 100644
--- a/crates/pyspark-rs/src/row.rs
+++ b/crates/pyspark-rs/src/row.rs
@@ -71,11 +71,16 @@ pub fn value_to_py<'py>(py: Python<'py>, v: &Value) ->
PyResult<Bound<'py, PyAny
dict.into_any()
}
Value::Struct(fields) => {
- let dict = PyDict::new(py);
- for (k, val) in fields {
- dict.set_item(k, value_to_py(py, val)?)?;
+ // A struct value materializes as a Row (named fields), matching
PySpark,
+ // rather than a plain dict.
+ let names: Vec<String> = fields.iter().map(|(k, _)|
k.clone()).collect();
+ let mut vals: Vec<Value> = Vec::with_capacity(fields.len());
+ for (_, val) in fields {
+ vals.push(val.clone());
}
- dict.into_any()
+ Py::new(py, PyRow::new(spark_connect::row::Row::new(names, vals)))?
+ .into_bound(py)
+ .into_any()
}
Value::Variant { value, metadata } => {
// A VARIANT materializes as a VariantVal (toJson/toPython decode
lazily),
diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs
index 3445c2a..e319015 100644
--- a/crates/pyspark-rs/src/session.rs
+++ b/crates/pyspark-rs/src/session.rs
@@ -244,9 +244,40 @@ impl PySparkSession {
Ok(PyDataFrame::new(df))
}
- /// Execute a SQL query and return a DataFrame.
- fn sql(&self, sqlQuery: &str) -> PyResult<PyDataFrame> {
- let df = self.session.sql(sqlQuery).to_pyerr()?;
+ /// Execute a SQL query and return a DataFrame. `args` binds parameters: a
list fills
+ /// positional `?` parameters, a dict fills named (`:name`) parameters.
Values are always
+ /// literals (a string arg is a string literal, never a column reference).
+ #[pyo3(signature = (sqlQuery, args=None))]
+ fn sql(&self, sqlQuery: &str, args: Option<Bound<'_, PyAny>>) ->
PyResult<PyDataFrame> {
+ fn to_lit(v: &Bound<'_, PyAny>) ->
PyResult<spark_connect::expression::Expression> {
+ use spark_connect::expression::{Expression, LiteralExpression};
+ if let Ok(sv) = v.extract::<String>() {
+ return Ok(Expression::Literal(LiteralExpression::string(sv)));
+ }
+ Ok(crate::functions::to_column(v)?.expression().clone())
+ }
+ let mut pos: Vec<spark_connect::expression::Expression> = Vec::new();
+ let mut named: std::collections::HashMap<String,
spark_connect::expression::Expression> =
+ std::collections::HashMap::new();
+ if let Some(a) = args {
+ if let Ok(d) = a.cast::<PyDict>() {
+ for (k, v) in d.iter() {
+ named.insert(k.extract::<String>()?, to_lit(&v)?);
+ }
+ } else if let Ok(lst) = a.cast::<PyList>() {
+ for v in lst.iter() {
+ pos.push(to_lit(&v)?);
+ }
+ } else {
+ return Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
+ "args must be a list (positional) or dict (named)",
+ ));
+ }
+ }
+ let df = self
+ .session
+ .sql_with_args(sqlQuery, pos, named)
+ .to_pyerr()?;
Ok(PyDataFrame::new(df))
}
@@ -375,16 +406,17 @@ impl PySparkSession {
rows.push(Row::new(names, values));
}
- if rows.is_empty() {
- return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
- "Data is empty",
- ));
- }
-
- // An explicit struct schema wins; otherwise infer field types from
the first row.
+ // An explicit struct schema wins; otherwise infer field types from
the first row
+ // (which requires at least one row -- an empty dataset needs an
explicit schema,
+ // matching pyspark which creates an empty DataFrame from schema + []
data).
let schema_dtype = match spec {
Spec::Struct(fields) => DataType::Struct { fields },
_ => {
+ if rows.is_empty() {
+ return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
+ "can not infer schema from empty dataset",
+ ));
+ }
let fields: Vec<StructField> = rows[0]
.fields()
.iter()
diff --git a/crates/pyspark-rs/src/window.rs b/crates/pyspark-rs/src/window.rs
index 429fc6b..c81aef3 100644
--- a/crates/pyspark-rs/src/window.rs
+++ b/crates/pyspark-rs/src/window.rs
@@ -4,7 +4,34 @@ use pyo3::prelude::*;
use spark_connect::expression::{Expression, NullOrdering, SortOrder};
use spark_connect::window::{FrameBound, WindowSpec};
-use crate::functions::to_column;
+use crate::dataframe::to_column_list;
+
+/// Build partition expressions from ColumnOrName args (a bare string is a
COLUMN
+/// NAME, not a literal — mirrors pyspark Window.partitionBy).
+fn to_partition_exprs(cols: Vec<Bound<'_, PyAny>>) ->
PyResult<Vec<Expression>> {
+ Ok(to_column_list(cols)?
+ .iter()
+ .map(|c| c.expression().clone())
+ .collect())
+}
+
+/// Build sort orders from ColumnOrName args; a bare string becomes an
ascending
+/// SortOrder over that column (not a string literal).
+fn to_sort_orders(cols: Vec<Bound<'_, PyAny>>) -> PyResult<Vec<SortOrder>> {
+ let mut sort_orders = Vec::new();
+ for column in to_column_list(cols)? {
+ if let Expression::SortOrder(sort) = column.expression() {
+ sort_orders.push((**sort).clone());
+ } else {
+ sort_orders.push(SortOrder {
+ child: column.expression().clone(),
+ ascending: true,
+ null_ordering: NullOrdering::Last,
+ });
+ }
+ }
+ Ok(sort_orders)
+}
/// Python wrapper for a window frame bound.
#[pyclass(name = "FrameBound", from_py_object)]
@@ -72,33 +99,14 @@ impl PyWindowSpec {
/// Partition the window by the given columns.
#[pyo3(signature = (*cols))]
fn partitionBy(&self, cols: Vec<Bound<'_, PyAny>>) ->
PyResult<PyWindowSpec> {
- let mut expressions = Vec::new();
- for col in cols {
- expressions.push(to_column(&col)?.expression().clone());
- }
- let new_spec = self.spec.clone().partition_by(expressions);
+ let new_spec =
self.spec.clone().partition_by(to_partition_exprs(cols)?);
Ok(PyWindowSpec::new(new_spec))
}
/// Order the window by the given columns.
#[pyo3(signature = (*cols))]
fn orderBy(&self, cols: Vec<Bound<'_, PyAny>>) -> PyResult<PyWindowSpec> {
- let mut sort_orders = Vec::new();
- for col in cols {
- let column = to_column(&col)?;
- // Extract sort order from the column
- if let Expression::SortOrder(sort) = column.expression() {
- sort_orders.push((**sort).clone());
- } else {
- // Default to ascending
- sort_orders.push(SortOrder {
- child: column.expression().clone(),
- ascending: true,
- null_ordering: NullOrdering::Last,
- });
- }
- }
- let new_spec = self.spec.clone().order_by(sort_orders);
+ let new_spec = self.spec.clone().order_by(to_sort_orders(cols)?);
Ok(PyWindowSpec::new(new_spec))
}
@@ -161,11 +169,7 @@ impl PyWindow {
#[pyo3(signature = (*cols))]
#[staticmethod]
fn partitionBy(cols: Vec<Bound<'_, PyAny>>) -> PyResult<PyWindowSpec> {
- let mut expressions = Vec::new();
- for col in cols {
- expressions.push(to_column(&col)?.expression().clone());
- }
- let spec = WindowSpec::new().partition_by(expressions);
+ let spec = WindowSpec::new().partition_by(to_partition_exprs(cols)?);
Ok(PyWindowSpec::new(spec))
}
@@ -173,22 +177,7 @@ impl PyWindow {
#[pyo3(signature = (*cols))]
#[staticmethod]
fn orderBy(cols: Vec<Bound<'_, PyAny>>) -> PyResult<PyWindowSpec> {
- let mut sort_orders = Vec::new();
- for col in cols {
- let column = to_column(&col)?;
- // Extract sort order from the column
- if let Expression::SortOrder(sort) = column.expression() {
- sort_orders.push((**sort).clone());
- } else {
- // Default to ascending
- sort_orders.push(SortOrder {
- child: column.expression().clone(),
- ascending: true,
- null_ordering: NullOrdering::Last,
- });
- }
- }
- let spec = WindowSpec::new().order_by(sort_orders);
+ let spec = WindowSpec::new().order_by(to_sort_orders(cols)?);
Ok(PyWindowSpec::new(spec))
}
diff --git a/crates/spark-connect/src/expression.rs
b/crates/spark-connect/src/expression.rs
index 29ca479..b19947d 100644
--- a/crates/spark-connect/src/expression.rs
+++ b/crates/spark-connect/src/expression.rs
@@ -866,14 +866,16 @@ fn to_proto_frame_boundary(
FrameBoundary::CurrentRow => Some(
proto::expression::window::window_frame::frame_boundary::Boundary::CurrentRow(true),
),
+ // A ROWS-frame offset must be an IntegerType literal (Spark rejects
bigint with
+ // SPECIFIED_WINDOW_FRAME_UNACCEPTED_TYPE), and preceding is a
negative offset.
FrameBoundary::Preceding(n) => Some(
proto::expression::window::window_frame::frame_boundary::Boundary::Value(Box::new(
- Expression::Literal(LiteralExpression::long(*n)).to_proto(),
+ Expression::Literal(LiteralExpression::int(-(*n as
i32))).to_proto(),
)),
),
FrameBoundary::Following(n) => Some(
proto::expression::window::window_frame::frame_boundary::Boundary::Value(Box::new(
- Expression::Literal(LiteralExpression::long(*n)).to_proto(),
+ Expression::Literal(LiteralExpression::int(*n as
i32)).to_proto(),
)),
),
};
diff --git a/crates/spark-connect/src/functions.rs
b/crates/spark-connect/src/functions.rs
index 6da75b7..3020388 100644
--- a/crates/spark-connect/src/functions.rs
+++ b/crates/spark-connect/src/functions.rs
@@ -1045,8 +1045,9 @@ pub fn kurtosis(col1: Column) -> Column {
}
/// Mirrors `pyspark.sql.functions.lag`.
-pub fn lag(col1: Column) -> Column {
- func("lag", vec![col1.expression().clone(), lit_int(1)])
+pub fn lag<C: Into<Column>>(cols: impl IntoIterator<Item = C>) -> Column {
+ let cols: Vec<Column> = cols.into_iter().map(Into::into).collect();
+ func("lag", cols.iter().map(|c| c.expression().clone()).collect())
}
/// Mirrors `pyspark.sql.functions.last`.
@@ -1070,8 +1071,12 @@ pub fn lcase(col1: Column) -> Column {
}
/// Mirrors `pyspark.sql.functions.lead`.
-pub fn lead(col1: Column) -> Column {
- func("lead", vec![col1.expression().clone(), lit_int(1)])
+pub fn lead<C: Into<Column>>(cols: impl IntoIterator<Item = C>) -> Column {
+ let cols: Vec<Column> = cols.into_iter().map(Into::into).collect();
+ func(
+ "lead",
+ cols.iter().map(|c| c.expression().clone()).collect(),
+ )
}
/// Mirrors `pyspark.sql.functions.length`.
@@ -2208,10 +2213,11 @@ pub fn instr(col1: Column, col2: Column) -> Column {
}
/// Mirrors `pyspark.sql.functions.json_tuple`.
-pub fn json_tuple(col1: Column, col2: Column) -> Column {
+pub fn json_tuple<C: Into<Column>>(cols: impl IntoIterator<Item = C>) ->
Column {
+ let cols: Vec<Column> = cols.into_iter().map(Into::into).collect();
func(
"json_tuple",
- vec![col1.expression().clone(), col2.expression().clone()],
+ cols.iter().map(|c| c.expression().clone()).collect(),
)
}
@@ -2785,12 +2791,11 @@ pub fn uniform(col1: Column, col2: Column) -> Column {
)
}
-/// Mirrors `pyspark.sql.functions.when`.
+/// Mirrors `pyspark.sql.functions.when`: starts a CASE WHEN (so chaining
+/// `.when(...).otherwise(...)` appends branches) rather than an opaque
+/// `when(...)` UnresolvedFunction (which would drop earlier branches on
chaining).
pub fn when(col1: Column, col2: Column) -> Column {
- func(
- "when",
- vec![col1.expression().clone(), col2.expression().clone()],
- )
+ crate::column::when(col1, col2)
}
/// Mirrors `pyspark.sql.functions.xpath`.
@@ -4989,7 +4994,7 @@ mod tests {
ilike(c1.clone(), c2.clone()),
instr(c1.clone(), c2.clone()),
jaro_winkler_similarity(c1.clone(), c2.clone()),
- json_tuple(c1.clone(), c2.clone()),
+ json_tuple([c1.clone(), c2.clone()]),
kll_sketch_get_quantile_bigint(c1.clone(), c2.clone()),
kll_sketch_get_quantile_double(c1.clone(), c2.clone()),
kll_sketch_get_quantile_float(c1.clone(), c2.clone()),
diff --git a/crates/spark-connect/src/plan.rs b/crates/spark-connect/src/plan.rs
index 0c280c7..c148f6f 100644
--- a/crates/spark-connect/src/plan.rs
+++ b/crates/spark-connect/src/plan.rs
@@ -33,8 +33,12 @@ pub enum LogicalPlan {
step: i64,
num_partitions: Option<i32>,
},
- /// A SQL query: `sql("SELECT ...")`.
- Sql { query: String },
+ /// A SQL query: `sql("SELECT ...")`, with optional positional/named
parameter bindings.
+ Sql {
+ query: String,
+ pos_args: Vec<Expression>,
+ named_args: HashMap<String, Expression>,
+ },
/// A projection (select): `df.select(cols...)`.
Project {
input: Box<LogicalPlan>,
@@ -418,9 +422,19 @@ impl LogicalPlan {
relation.rel_type =
Some(proto::relation::RelType::Range(range));
}
- LogicalPlan::Sql { query } => {
+ LogicalPlan::Sql {
+ query,
+ pos_args,
+ named_args,
+ } => {
let mut sql = proto::Sql::default();
sql.query = query.clone();
+ for e in pos_args {
+ sql.pos_arguments.push(e.to_proto());
+ }
+ for (k, v) in named_args {
+ sql.named_arguments.insert(k.clone(), v.to_proto());
+ }
relation.rel_type = Some(proto::relation::RelType::Sql(sql));
}
@@ -1313,6 +1327,8 @@ pub fn range_with_partitions(start: i64, end: i64, step:
i64, num_partitions: i3
pub fn sql(query: impl Into<String>) -> LogicalPlan {
LogicalPlan::Sql {
query: query.into(),
+ pos_args: Vec::new(),
+ named_args: HashMap::new(),
}
}
diff --git a/crates/spark-connect/src/session.rs
b/crates/spark-connect/src/session.rs
index 097cace..eff2adf 100644
--- a/crates/spark-connect/src/session.rs
+++ b/crates/spark-connect/src/session.rs
@@ -255,8 +255,21 @@ impl SparkSession {
///
/// Mirrors `pyspark.sql.SparkSession.sql(sqlQuery)`.
pub fn sql(&self, query: &str) -> Result<DataFrame> {
+ self.sql_with_args(query, Vec::new(), std::collections::HashMap::new())
+ }
+
+ /// SQL with parameters: positional (`pos_args`) and/or named
(`named_args`) expression
+ /// bindings, mirroring `SparkSession.sql(query, args=...)`.
+ pub fn sql_with_args(
+ &self,
+ query: &str,
+ pos_args: Vec<crate::expression::Expression>,
+ named_args: std::collections::HashMap<String,
crate::expression::Expression>,
+ ) -> Result<DataFrame> {
let plan = LogicalPlan::Sql {
query: query.to_string(),
+ pos_args,
+ named_args,
};
Ok(DataFrame::new(self.clone(), plan))
}
@@ -265,13 +278,15 @@ impl SparkSession {
///
/// Mirrors `pyspark.sql.SparkSession.createDataFrame(rows, schema)`.
pub fn create_dataframe(&self, rows: Vec<Row>, schema: DataType) ->
Result<DataFrame> {
- // Build Arrow RecordBatch from rows and schema
- let arrow_data = rows_to_arrow_ipc(&rows, &schema)?;
-
- let plan = LogicalPlan::LocalRelation {
- schema,
- data: Some(arrow_data),
+ // Empty data -> send the schema only (no Arrow `data`); the server
builds an empty
+ // relation from the schema. A zero-row Arrow payload would be empty
bytes, which the
+ // server cannot parse ("Unexpected end of input. Missing schema.").
+ let data = if rows.is_empty() {
+ None
+ } else {
+ Some(rows_to_arrow_ipc(&rows, &schema)?)
};
+ let plan = LogicalPlan::LocalRelation { schema, data };
Ok(DataFrame::new(self.clone(), plan))
}
@@ -833,15 +848,51 @@ fn values_to_arrow(values: &[Value], dt: &DataType) ->
Result<Arc<dyn arrow::arr
}};
}
+ // Numeric encoders coerce across integer/float widths, so a value the
schema-unaware
+ // Python->Value conversion produced as a wider/narrower numeric (e.g. a
nested-struct
+ // int arriving as Long) still matches its declared field type.
+ macro_rules! int_prim {
+ ($arr:ty, $t:ty) => {{
+ let vals: Result<Vec<_>> = values
+ .iter()
+ .map(|v| match v {
+ Value::Null => Ok(None),
+ Value::Byte(x) => Ok(Some(*x as $t)),
+ Value::Short(x) => Ok(Some(*x as $t)),
+ Value::Integer(x) => Ok(Some(*x as $t)),
+ Value::Long(x) => Ok(Some(*x as $t)),
+ _ => Err(SparkError::connect_msg("Type mismatch in row
data")),
+ })
+ .collect();
+ Ok(Arc::new(<$arr>::from(vals?)) as Arc<dyn Array>)
+ }};
+ }
+ macro_rules! float_prim {
+ ($arr:ty, $t:ty) => {{
+ let vals: Result<Vec<_>> = values
+ .iter()
+ .map(|v| match v {
+ Value::Null => Ok(None),
+ Value::Float(x) => Ok(Some(*x as $t)),
+ Value::Double(x) => Ok(Some(*x as $t)),
+ Value::Integer(x) => Ok(Some(*x as $t)),
+ Value::Long(x) => Ok(Some(*x as $t)),
+ _ => Err(SparkError::connect_msg("Type mismatch in row
data")),
+ })
+ .collect();
+ Ok(Arc::new(<$arr>::from(vals?)) as Arc<dyn Array>)
+ }};
+ }
+
match dt {
DataType::Null => Ok(Arc::new(NullArray::new(values.len()))),
DataType::Boolean => prim!(BooleanArray, Value::Bool),
- DataType::Byte => prim!(Int8Array, Value::Byte),
- DataType::Short => prim!(Int16Array, Value::Short),
- DataType::Integer => prim!(Int32Array, Value::Integer),
- DataType::Long => prim!(Int64Array, Value::Long),
- DataType::Float => prim!(Float32Array, Value::Float),
- DataType::Double => prim!(Float64Array, Value::Double),
+ DataType::Byte => int_prim!(Int8Array, i8),
+ DataType::Short => int_prim!(Int16Array, i16),
+ DataType::Integer => int_prim!(Int32Array, i32),
+ DataType::Long => int_prim!(Int64Array, i64),
+ DataType::Float => float_prim!(Float32Array, f32),
+ DataType::Double => float_prim!(Float64Array, f64),
DataType::String { .. } | DataType::Char { .. } | DataType::Varchar {
.. } => {
let vals: Result<Vec<Option<&str>>> = values
.iter()
@@ -959,6 +1010,14 @@ fn values_to_arrow(values: &[Value], dt: &DataType) ->
Result<Arc<dyn arrow::arr
cols[i].push(val);
}
}
+ // A Python dict supplied for a struct field (e.g.
array<struct> data)
+ // arrives as a Map; match its entries to the struct
fields by name.
+ Value::Map(m) => {
+ valid.append(true);
+ for (i, f) in fields.iter().enumerate() {
+
cols[i].push(m.get(&f.name).cloned().unwrap_or(Value::Null));
+ }
+ }
_ => return Err(SparkError::connect_msg("Type mismatch:
expected struct")),
}
}
diff --git a/crates/spark-connect/src/types.rs
b/crates/spark-connect/src/types.rs
index 1013be6..aa8a42d 100644
--- a/crates/spark-connect/src/types.rs
+++ b/crates/spark-connect/src/types.rs
@@ -1362,16 +1362,38 @@ fn parse_geography(s: &str) -> Result<DataType> {
fn parse_datatype_string(input: &str) -> Result<DataType> {
let trimmed = input.trim();
- // Check if it's a top-level schema (multiple fields like "a INT, b
STRING")
- // Heuristic: contains comma at top level (not inside angle
brackets/parens)
- // and doesn't start with "struct<"
- if !trimmed.to_lowercase().starts_with("struct<") &&
contains_top_level_comma(trimmed) {
- return parse_top_level_schema(trimmed);
+ // A top-level schema is "a INT, b STRING" (multiple fields) OR a single
"a INT"
+ // (one field, no comma) -- both parse to a StructType, matching PySpark's
+ // _parse_datatype_string. Detect either a top-level comma or a top-level
space
+ // separating a field name from its type; fall back to single-type parsing
if the
+ // schema interpretation fails (e.g. a bare type that itself contains
spaces such as
+ // "interval day to second").
+ if !trimmed.to_lowercase().starts_with("struct<")
+ && (contains_top_level_comma(trimmed) ||
contains_top_level_whitespace(trimmed))
+ {
+ if let Ok(dt) = parse_top_level_schema(trimmed) {
+ return Ok(dt);
+ }
}
parse_single_type(trimmed)
}
+/// Whether `s` has a top-level whitespace char (outside `<>`/`()`), i.e. it
looks like a
+/// "name type" field rather than a bare type such as `int` or `array<int>`.
+fn contains_top_level_whitespace(s: &str) -> bool {
+ let mut depth = 0;
+ for c in s.chars() {
+ match c {
+ '<' | '(' => depth += 1,
+ '>' | ')' => depth -= 1,
+ _ if depth == 0 && c.is_whitespace() => return true,
+ _ => {}
+ }
+ }
+ false
+}
+
/// Check if a string contains a comma at the top level (not inside
brackets/parens)
fn contains_top_level_comma(s: &str) -> bool {
let mut depth = 0;
diff --git a/crates/spark-connect/tests/functions_golden.rs
b/crates/spark-connect/tests/functions_golden.rs
index e0c1348..885653c 100644
--- a/crates/spark-connect/tests/functions_golden.rs
+++ b/crates/spark-connect/tests/functions_golden.rs
@@ -308,14 +308,14 @@ fn all_golden_function_cases_pass() {
("java_method", java_method(Vec::<Column>::new())),
("json_array_length", json_array_length(a())),
("json_object_keys", json_object_keys(a())),
- ("json_tuple", json_tuple(a(), b())),
+ ("json_tuple", json_tuple([a(), b()])),
("kurtosis", kurtosis(a())),
- ("lag", lag(a())),
+ ("lag", lag([a()])),
("last", last(a())),
("last_day", last_day(a())),
("last_value", last_value(a())),
("lcase", lcase(a())),
- ("lead", lead(a())),
+ ("lead", lead([a()])),
("least", least(a(), b())),
("left", left(a(), b())),
("length", length(a())),
diff --git a/crates/spark-connect/tests/window_golden.rs
b/crates/spark-connect/tests/window_golden.rs
index 781048b..ce064a2 100644
--- a/crates/spark-connect/tests/window_golden.rs
+++ b/crates/spark-connect/tests/window_golden.rs
@@ -230,7 +230,7 @@ fn build(name: &str) -> Option<Expression> {
.expression()
.clone(),
- "window_lag" => functions::lag(col("x"))
+ "window_lag" => functions::lag([col("x")])
.over(
WindowSpec::new()
.partition_by(vec![col("a").expression().clone()])
@@ -241,7 +241,7 @@ fn build(name: &str) -> Option<Expression> {
.expression()
.clone(),
- "window_lead" => functions::lead(col("x"))
+ "window_lead" => functions::lead([col("x")])
.over(
WindowSpec::new()
.partition_by(vec![col("a").expression().clone()])
diff --git a/python/pyspark/sql/functions.py b/python/pyspark/sql/functions.py
index 8c4cdf2..ec00181 100644
--- a/python/pyspark/sql/functions.py
+++ b/python/pyspark/sql/functions.py
@@ -420,6 +420,27 @@ for name in dir(_gen_module):
if not name.startswith('_'):
globals()[name] = getattr(_gen_module, name)
+
+# Overrides for functions whose Python signature carries extra literal args
that the
+# generated single-arg wrapper cannot express (offset/default, variadic field
names).
+# Defined AFTER the generated import so they take precedence.
Field/offset/default are
+# literals; only the first arg is a column (PySpark ColumnOrName).
+def lag(col, offset=1, default=None):
+ """lag(col, offset=1, default=None): value offset rows before the current
row."""
+ return _wrap(_call_function("lag", _to_col(col), _pyfunc_lit(offset),
_pyfunc_lit(default)))
+
+
+def lead(col, offset=1, default=None):
+ """lead(col, offset=1, default=None): value offset rows after the current
row."""
+ return _wrap(_call_function("lead", _to_col(col), _pyfunc_lit(offset),
_pyfunc_lit(default)))
+
+
+def json_tuple(col, *fields):
+ """json_tuple(col, *fields): extract the given JSON fields (literal names)
as columns."""
+ if not fields:
+ raise ValueError("json_tuple() takes at least one field name")
+ return _wrap(_call_function("json_tuple", _to_col(col), *[_pyfunc_lit(f)
for f in fields]))
+
# ---------------------------------------------------------------------------
# Higher-order functions.
#
diff --git a/python/tests/test_dropin_offline.py
b/python/tests/test_dropin_offline.py
index b7402a4..14e236c 100644
--- a/python/tests/test_dropin_offline.py
+++ b/python/tests/test_dropin_offline.py
@@ -337,6 +337,10 @@ def test_datatype_object_model():
assert T.IntegerType().toInternal(5) == 5
# fromDDL parses a DDL string client-side.
assert T.DataType.fromDDL("a int, b string").simpleString() ==
"struct<a:int,b:string>"
+ # A single-field DDL schema (no comma) must also parse to a one-field
struct, not be
+ # misread as a bare type (SPARK-59113 follow-up).
+ assert T.DataType.fromDDL("i int").simpleString() == "struct<i:int>"
+ assert T.StructType.fromDDL("id long").simpleString() ==
"struct<id:bigint>"
def test_structfield_collation_methods():
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
