This is an automated email from the ASF dual-hosted git repository.
dbtsai 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 d884e34 [SPARK-60606][CONNECT] Build local Arrow input one column at
a time
d884e34 is described below
commit d884e347452c6446c7bc54af151755c62801d8cd
Author: DB Tsai <[email protected]>
AuthorDate: Wed Oct 28 14:16:31 2026 -0700
[SPARK-60606][CONNECT] Build local Arrow input one column at a time
Local input conversion builds one Arrow column at a time, avoiding a fully
coerced row matrix and per-row field-name copies. Numeric conversion now
happens only in the schema-driven Arrow encoder, which already recurses through
arrays, maps, and structs. The redundant coercion pass and unused scalar
builder are removed.
Resolving nested struct names only once also fixes premature narrowing
across duplicate names: a Float and Double field with the same name now convert
the original source value independently. Wrong-width rows return a normal error
instead of panicking, both in the Rust input conversion and in the Python
wrapper before Row construction. A shared helper in `spark-connect` produces
the same message and PySparkValueError type for inferred, names-only, and
explicit StructType Python schema [...]
IPC tests share the existing test module and field helper. They cover
in-range scalar conversion, nulls, nested arrays/maps/structs, duplicate-name
numeric independence, input/decoder round trips, and short/long row errors. A
Python regression covers inferred, names-only, and explicit StructType schemas.
The clone-test cleanup from #106 remains included; the Int64-only decoder
optimization is absent.
Validation for this revision: 392 client library tests and 95 core library
tests passed; the complete offline Rust suite passed (server-gated tests did
not execute against a live server); all integration targets compiled with and
without `wasm-udf`; formatting and core Clippy passed. The rebuilt Python
extension passed width-error regressions and retained the same 108-class API.
The preceding revision also passed all 11 independent Python fixture
comparisons (fixture-only RPC, not liv [...]
The prior CPU/wall/RSS measurements in SPARK-60606 cover the earlier
implementation that retained coercion. They do not validate this revision.
Fresh screening and confirmation are still needed, including CPU (≤+5%), wall
(≤+10%), and RSS (≤the larger of 5% or 4 MiB) regression confidence bounds; no
performance acceptance is claimed for the revised code.
The Python wrapper's per-row `spec_names` clones are addressed separately
in [PR #113](https://github.com/apache/spark-connect-rust/pull/113). Remaining
intermediate value copies are tracked in
[SPARK-60640](https://issues.apache.org/jira/browse/SPARK-60640), and
Byte/Short-to-Float/Double conversion support in
[SPARK-60641](https://issues.apache.org/jira/browse/SPARK-60641). Integer
overflow validation also remains outside this change; the tests no longer
assert wrapping as intended [...]
[SPARK-60606 and historical benchmark
details](https://issues.apache.org/jira/browse/SPARK-60606) ·
[Harness](https://github.com/dbtsai/pyspark-connect-benchmark/tree/23bdf9046815ae615c921b2e6bab827eed00aaae)
This pull request and its description were written by Isaac.
Closes #109 from dbtsai/codex/rsi-input-int64-followup.
Authored-by: DB Tsai <[email protected]>
Signed-off-by: DB Tsai <[email protected]>
Co-authored-by: Isaac <[email protected]>
---
crates/pyspark-rs/src/session.rs | 5 +-
crates/spark-connect/src/row.rs | 15 +-
crates/spark-connect/src/session.rs | 649 ++++++++++++++---------------
crates/spark-connect/tests/e2e_coverage.rs | 2 +-
python/tests/test_dropin_offline.py | 26 ++
5 files changed, 339 insertions(+), 358 deletions(-)
diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs
index 5650e60..5f6302e 100644
--- a/crates/pyspark-rs/src/session.rs
+++ b/crates/pyspark-rs/src/session.rs
@@ -3,7 +3,7 @@
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList};
use spark_connect::row::{Row, Value};
-use spark_connect::session::SparkSession;
+use spark_connect::session::{row_width_mismatch, SparkSession};
use spark_connect::types::{DataType, StructField};
use spark_connect::udf::{CommonInlineUserDefinedFunctionExpression,
PythonUDFPayload};
@@ -457,6 +457,9 @@ impl PySparkSession {
.or(row_names)
.or_else(|| pandas_columns.clone())
.unwrap_or_else(|| (0..values.len()).map(|i| format!("col{}",
i)).collect());
+ if names.len() != values.len() {
+ return Err(row_width_mismatch(rows.len(), values.len(),
names.len())).to_pyerr();
+ }
rows.push(Row::new(names, values));
}
diff --git a/crates/spark-connect/src/row.rs b/crates/spark-connect/src/row.rs
index 4c9cf5c..8c5c979 100644
--- a/crates/spark-connect/src/row.rs
+++ b/crates/spark-connect/src/row.rs
@@ -559,7 +559,7 @@ mod tests {
}
#[test]
- fn nested_struct_clones_share_names_and_keep_child_values_independent() {
+ fn nested_struct_clones_share_field_names() {
let fields: Arc<[String]> = vec!["items".to_string()].into();
let value = Value::Struct(
Row::with_shared_fields(
@@ -568,24 +568,15 @@ mod tests {
)
.into(),
);
- let mut cloned = value.clone();
+ let cloned = value.clone();
let Value::Struct(original_row) = &value else {
panic!("expected a struct value");
};
- let Value::Struct(cloned_row) = &mut cloned else {
+ let Value::Struct(cloned_row) = &cloned else {
panic!("expected a cloned struct value");
};
assert!(std::ptr::eq(fields.as_ref(), cloned_row.fields()));
assert!(std::ptr::eq(original_row.fields(), cloned_row.fields()));
- let Value::List(items) = &mut cloned_row.values[0] else {
- panic!("expected a list value");
- };
- items[0] = Value::String("changed".to_string());
- assert_eq!(
- original_row.get(0),
- Some(&Value::List(vec![Value::String("original".to_string())]))
- );
- assert_ne!(value, cloned);
}
#[test]
diff --git a/crates/spark-connect/src/session.rs
b/crates/spark-connect/src/session.rs
index 30a31a8..7eb29ac 100644
--- a/crates/spark-connect/src/session.rs
+++ b/crates/spark-connect/src/session.rs
@@ -596,6 +596,14 @@ impl Default for SparkSessionBuilder {
}
}
+/// Construct the shared local-input width error for Rust and Python callers.
+#[doc(hidden)]
+pub fn row_width_mismatch(row_index: usize, actual: usize, expected: usize) ->
SparkError {
+ SparkError::value_msg(format!(
+ "Row {row_index} has {actual} values, but the schema has {expected}
fields"
+ ))
+}
+
/// Convert Arrow RecordBatch data to IPC bytes.
/// Used by createDataFrame to package local data.
fn rows_to_arrow_ipc(rows: &[Row], schema: &DataType) -> Result<Vec<u8>> {
@@ -618,23 +626,19 @@ fn rows_to_arrow_ipc(rows: &[Row], schema: &DataType) ->
Result<Vec<u8>> {
_ => return Err(SparkError::connect_msg("Schema is not a struct
type")),
};
- // Coerce each value to its declared field type so the built arrays match
the
- // Arrow schema (e.g. a Python int decodes as Long but "a int" wants
Int32).
- let coerced: Vec<Row> = rows
- .iter()
- .map(|r| {
- let vals: Vec<Value> = fields
- .iter()
- .enumerate()
- .map(|(i, f)| coerce_value(r.get(i), &f.data_type))
- .collect();
- Row::new(r.fields().to_vec(), vals)
- })
- .collect();
+ // Conversion is positional, so every row must match the schema width.
+ for (index, row) in rows.iter().enumerate() {
+ if row.len() != fields.len() {
+ return Err(row_width_mismatch(index, row.len(), fields.len()));
+ }
+ }
- for field_idx in 0..fields.len() {
- let array = build_arrow_array(&coerced, field_idx,
Some(&fields[field_idx].data_type))?;
- columns.push(array);
+ for (field_idx, field) in fields.iter().enumerate() {
+ let column: Vec<Value> = rows
+ .iter()
+ .map(|row| row.values()[field_idx].clone())
+ .collect();
+ columns.push(values_to_arrow(&column, &field.data_type)?);
}
// Create RecordBatch
@@ -747,86 +751,6 @@ fn datatype_to_arrow(dt: &DataType) ->
Result<arrow::datatypes::DataType> {
})
}
-/// Coerce a value to a declared field type (numeric widening/narrowing from
the
-/// Python-inferred type). Non-numeric or already-matching values pass
through, so an
-/// explicit `createDataFrame` schema produces arrays that match the Arrow
schema.
-fn coerce_value(v: Option<&Value>, target: &DataType) -> Value {
- let v = match v {
- Some(v) => v,
- None => return Value::Null,
- };
- match (target, v) {
- (_, Value::Null) => Value::Null,
- (DataType::Byte, Value::Long(n)) => Value::Byte(*n as i8),
- (DataType::Byte, Value::Integer(n)) => Value::Byte(*n as i8),
- (DataType::Short, Value::Long(n)) => Value::Short(*n as i16),
- (DataType::Short, Value::Integer(n)) => Value::Short(*n as i16),
- (DataType::Integer, Value::Long(n)) => Value::Integer(*n as i32),
- (DataType::Long, Value::Integer(n)) => Value::Long(*n as i64),
- // Float target from any narrower numeric (a Python int decodes as
Long/Integer
- // but "a float" wants Float32); without these, an Int64Array is built
against a
- // Float32 field and RecordBatch::try_new rejects the mismatch.
- (DataType::Float, Value::Double(f)) => Value::Float(*f as f32),
- (DataType::Float, Value::Long(n)) => Value::Float(*n as f32),
- (DataType::Float, Value::Integer(n)) => Value::Float(*n as f32),
- (DataType::Double, Value::Long(n)) => Value::Double(*n as f64),
- (DataType::Double, Value::Integer(n)) => Value::Double(*n as f64),
- (DataType::Double, Value::Float(f)) => Value::Double(*f as f64),
- // Nested types: coerce children to their declared element/field/value
types so a
- // nested `Integer` field built from a Python int (Long) matches the
Arrow schema.
- (DataType::Array { element_type, .. }, Value::List(items)) =>
Value::List(
- items
- .iter()
- .map(|it| coerce_value(Some(it), element_type))
- .collect(),
- ),
- (DataType::Struct { fields }, Value::Struct(sf)) => Value::Struct(
- Row::new(
- fields.iter().map(|f| f.name.clone()).collect(),
- fields
- .iter()
- .enumerate()
- .map(|(i, f)| {
- let val = sf.get_by_name(&f.name).or_else(||
sf.get(i));
- coerce_value(val, &f.data_type)
- })
- .collect(),
- )
- .into(),
- ),
- (DataType::Map { value_type, .. }, Value::Map(m)) => Value::Map(
- m.iter()
- .map(|(k, v)| (k.clone(), coerce_value(Some(v), value_type)))
- .collect(),
- ),
- _ => v.clone(),
- }
-}
-
-/// Build an Arrow array for a specific field across all rows.
-/// Build the Arrow array for one column: gather the column's values and build
them
-/// against the declared type (schema-driven, so nested array/map/struct
recurse).
-fn build_arrow_array(
- rows: &[Row],
- field_idx: usize,
- target: Option<&DataType>,
-) -> Result<Arc<dyn arrow::array::Array>> {
- use arrow::array::NullArray;
- if rows.is_empty() {
- return Ok(Arc::new(NullArray::new(0)));
- }
- let column: Vec<Value> = rows
- .iter()
- .map(|r| r.get(field_idx).cloned().unwrap_or(Value::Null))
- .collect();
- match target {
- Some(dt) => values_to_arrow(&column, dt),
- // No declared type (shouldn't happen for createDataFrame): fall back
to the
- // value-driven scalar builder for backward compatibility.
- None => build_arrow_array_scalar(rows, field_idx),
- }
-}
-
/// Recursively build an Arrow array from a column of `Value`s against a
declared
/// `DataType`, covering scalars and the nested array/map/struct types.
fn values_to_arrow(values: &[Value], dt: &DataType) -> Result<Arc<dyn
arrow::array::Array>> {
@@ -1099,171 +1023,6 @@ fn values_to_arrow(values: &[Value], dt: &DataType) ->
Result<Arc<dyn arrow::arr
}
}
-/// Legacy value-driven scalar array builder (used only when no declared type
is available).
-fn build_arrow_array_scalar(
- rows: &[Row],
- field_idx: usize,
-) -> Result<Arc<dyn arrow::array::Array>> {
- use arrow::array::*;
-
- if rows.is_empty() {
- return Ok(Arc::new(NullArray::new(0)));
- }
-
- // Peek at the first row to determine the type
- let first_val = rows[0]
- .get(field_idx)
- .ok_or_else(|| SparkError::connect_msg("Invalid row index"))?;
-
- match first_val {
- Value::Null => Ok(Arc::new(NullArray::new(rows.len()))),
- Value::Bool(_) => {
- let values: Result<Vec<Option<bool>>> = rows
- .iter()
- .map(|r| match r.get(field_idx) {
- None | Some(Value::Null) => Ok(None),
- Some(Value::Bool(b)) => Ok(Some(*b)),
- Some(_) => Err(SparkError::connect_msg("Type mismatch in
row data")),
- })
- .collect();
- Ok(Arc::new(BooleanArray::from(values?)))
- }
- Value::Byte(_) => {
- let values: Result<Vec<_>> = rows
- .iter()
- .map(|r| match r.get(field_idx) {
- None | Some(Value::Null) => Ok(None),
- Some(Value::Byte(b)) => Ok(Some(*b)),
- Some(_) => Err(SparkError::connect_msg("Type mismatch in
row data")),
- })
- .collect();
- Ok(Arc::new(Int8Array::from(values?)))
- }
- Value::Short(_) => {
- let values: Result<Vec<_>> = rows
- .iter()
- .map(|r| match r.get(field_idx) {
- None | Some(Value::Null) => Ok(None),
- Some(Value::Short(s)) => Ok(Some(*s)),
- Some(_) => Err(SparkError::connect_msg("Type mismatch in
row data")),
- })
- .collect();
- Ok(Arc::new(Int16Array::from(values?)))
- }
- Value::Integer(_) => {
- let values: Result<Vec<_>> = rows
- .iter()
- .map(|r| match r.get(field_idx) {
- None | Some(Value::Null) => Ok(None),
- Some(Value::Integer(i)) => Ok(Some(*i)),
- Some(_) => Err(SparkError::connect_msg("Type mismatch in
row data")),
- })
- .collect();
- Ok(Arc::new(Int32Array::from(values?)))
- }
- Value::Long(_) => {
- let values: Result<Vec<_>> = rows
- .iter()
- .map(|r| match r.get(field_idx) {
- None | Some(Value::Null) => Ok(None),
- Some(Value::Long(l)) => Ok(Some(*l)),
- Some(_) => Err(SparkError::connect_msg("Type mismatch in
row data")),
- })
- .collect();
- Ok(Arc::new(Int64Array::from(values?)))
- }
- Value::Float(_) => {
- let values: Result<Vec<_>> = rows
- .iter()
- .map(|r| match r.get(field_idx) {
- None | Some(Value::Null) => Ok(None),
- Some(Value::Float(f)) => Ok(Some(*f)),
- Some(_) => Err(SparkError::connect_msg("Type mismatch in
row data")),
- })
- .collect();
- Ok(Arc::new(Float32Array::from(values?)))
- }
- Value::Double(_) => {
- let values: Result<Vec<_>> = rows
- .iter()
- .map(|r| match r.get(field_idx) {
- None | Some(Value::Null) => Ok(None),
- Some(Value::Double(d)) => Ok(Some(*d)),
- Some(_) => Err(SparkError::connect_msg("Type mismatch in
row data")),
- })
- .collect();
- Ok(Arc::new(Float64Array::from(values?)))
- }
- Value::String(_) => {
- let values: Result<Vec<_>> = rows
- .iter()
- .map(|r| match r.get(field_idx) {
- None | Some(Value::Null) => Ok(None),
- Some(Value::String(s)) => Ok(Some(s.as_str())),
- Some(_) => Err(SparkError::connect_msg("Type mismatch in
row data")),
- })
- .collect();
- Ok(Arc::new(StringArray::from(values?)))
- }
- Value::Binary(_) => {
- let values: Result<Vec<_>> = rows
- .iter()
- .map(|r| match r.get(field_idx) {
- None | Some(Value::Null) => Ok(None),
- Some(Value::Binary(b)) => Ok(Some(b.as_slice())),
- Some(_) => Err(SparkError::connect_msg("Type mismatch in
row data")),
- })
- .collect();
- Ok(Arc::new(BinaryArray::from(values?)))
- }
- Value::Date(_) => {
- let values: Result<Vec<Option<i32>>> = rows
- .iter()
- .map(|r| match r.get(field_idx) {
- None | Some(Value::Null) => Ok(None),
- Some(Value::Date(d)) => Ok(Some(*d)),
- Some(_) => Err(SparkError::connect_msg("Type mismatch in
row data")),
- })
- .collect();
- Ok(Arc::new(Date32Array::from(values?)))
- }
- Value::Timestamp(_) => {
- let values: Result<Vec<Option<i64>>> = rows
- .iter()
- .map(|r| match r.get(field_idx) {
- None | Some(Value::Null) => Ok(None),
- Some(Value::Timestamp(t)) => Ok(Some(*t)),
- Some(_) => Err(SparkError::connect_msg("Type mismatch in
row data")),
- })
- .collect();
- // Scalar fallback (no declared type): default to UTC (LTZ).
- let arr =
TimestampMicrosecondArray::from(values?).with_timezone("UTC");
- Ok(Arc::new(arr))
- }
- Value::Decimal { scale, .. } => {
- // Column-wide precision/scale come from the first value (Spark
decimals in a
- // column share one precision/scale). Each value's string is
parsed to an
- // unscaled i128 at that scale.
- let col_scale = scale.unwrap_or(0);
- let values: Result<Vec<Option<i128>>> = rows
- .iter()
- .map(|r| match r.get(field_idx) {
- None | Some(Value::Null) => Ok(None),
- Some(Value::Decimal { value, .. }) => {
- decimal_str_to_unscaled(value, col_scale).map(Some)
- }
- Some(_) => Err(SparkError::connect_msg("Type mismatch in
row data")),
- })
- .collect();
- let arr = Decimal128Array::from(values?)
- .with_precision_and_scale(38, col_scale as i8)
- .map_err(|e| SparkError::connect_msg(format!("decimal build:
{e}")))?;
- Ok(Arc::new(arr))
- }
- _ => Err(SparkError::connect_msg("Unsupported value type")),
- }
-}
-
/// Parse a decimal string (e.g. "-1.50") into an unscaled `i128` at the given
scale
/// (e.g. scale 2 → -150). Pads/truncates the fractional part to `scale`
digits.
fn decimal_str_to_unscaled(s: &str, scale: i32) -> Result<i128> {
@@ -1300,6 +1059,12 @@ fn decimal_str_to_unscaled(s: &str, scale: i32) ->
Result<i128> {
#[cfg(test)]
mod tests {
use super::*;
+ use arrow::array::{
+ BinaryArray, BooleanArray, Float32Array, Int16Array, Int32Array,
StringArray, StructArray,
+ };
+ use arrow::ipc::reader::StreamReader;
+ use arrow::record_batch::RecordBatch;
+ use std::io::Cursor;
fn struct_field(name: &str, data_type: DataType) ->
crate::types::StructField {
crate::types::StructField {
@@ -1310,88 +1075,6 @@ mod tests {
}
}
- #[test]
- fn
nested_struct_coercion_preserves_first_name_position_and_numeric_types() {
- let source = Value::Struct(
- Row::with_shared_fields(
- vec!["dup", "dup", "other", "nested", "items", "map"]
- .into_iter()
- .map(str::to_string)
- .collect::<Arc<[String]>>(),
- vec![
- Value::Long(257),
- Value::Long(999),
- Value::Long(-2),
- Value::Struct(Row::new(vec!["n".to_string()],
vec![Value::Long(300)]).into()),
- Value::List(vec![Value::Long(65_537), Value::Null]),
- Value::Map([("k".to_string(),
Value::Long(65_538))].into()),
- ],
- )
- .into(),
- );
- let target = DataType::Struct {
- fields: vec![
- struct_field("dup", DataType::Byte),
- struct_field("positional", DataType::Short),
- struct_field("dup", DataType::Long),
- struct_field(
- "nested",
- DataType::Struct {
- fields: vec![struct_field("n", DataType::Byte)],
- },
- ),
- struct_field(
- "items",
- DataType::Array {
- element_type: Box::new(DataType::Short),
- contains_null: true,
- },
- ),
- struct_field(
- "map",
- DataType::Map {
- key_type: Box::new(DataType::String {
- collation: "UTF8_BINARY".to_string(),
- }),
- value_type: Box::new(DataType::Short),
- value_contains_null: true,
- },
- ),
- struct_field("missing", DataType::Integer),
- ],
- };
- let actual = coerce_value(Some(&source), &target);
- let expected = Value::Struct(
- Row::new(
- vec![
- "dup",
- "positional",
- "dup",
- "nested",
- "items",
- "map",
- "missing",
- ]
- .into_iter()
- .map(str::to_string)
- .collect(),
- vec![
- Value::Byte(1),
- Value::Short(999),
- Value::Long(257),
- Value::Struct(Row::new(vec!["n".to_string()],
vec![Value::Byte(44)]).into()),
- Value::List(vec![Value::Short(1), Value::Null]),
- Value::Map([("k".to_string(), Value::Short(2))].into()),
- Value::Null,
- ],
- )
- .into(),
- );
- assert_eq!(actual, expected);
- assert_eq!(coerce_value(Some(&Value::Null), &target), Value::Null);
- assert_eq!(coerce_value(None, &target), Value::Null);
- }
-
#[test]
fn struct_arrow_encoding_preserves_lookup_map_and_null_semantics() {
use arrow::array::{Array, Int16Array, Int32Array, Int64Array,
StructArray};
@@ -1499,7 +1182,8 @@ mod tests {
)
.into(),
);
- let array = values_to_arrow(&[source], &target).unwrap();
+ let batch = read_batch(&[row(vec![source])], &schema(vec![target]));
+ let array = batch.column(0);
let array = array.as_any().downcast_ref::<StructArray>().unwrap();
let items = array
.column(0)
@@ -1612,4 +1296,281 @@ mod tests {
let tags = spark2.get_tags();
assert!(tags.contains(&"test".to_string()));
}
+
+ // createDataFrame input IPC fixtures.
+ fn text_type() -> DataType {
+ DataType::String {
+ collation: "UTF8_BINARY".to_string(),
+ }
+ }
+
+ fn schema(types: Vec<DataType>) -> DataType {
+ DataType::Struct {
+ fields: types
+ .into_iter()
+ .enumerate()
+ .map(|(i, dt)| struct_field(&format!("schema_{i}"), dt))
+ .collect(),
+ }
+ }
+
+ fn row(values: Vec<Value>) -> Row {
+ // Top-level conversion is positional, even when input and schema
names differ.
+ Row::new(
+ (0..values.len()).map(|i| format!("input_{i}")).collect(),
+ values,
+ )
+ }
+
+ fn read_batch(rows: &[Row], schema: &DataType) -> RecordBatch {
+ let data = rows_to_arrow_ipc(rows, schema).unwrap();
+ let mut reader = StreamReader::try_new(Cursor::new(data),
None).unwrap();
+ let batch = reader.next().unwrap().unwrap();
+ assert!(reader.next().is_none());
+ assert_eq!(batch.num_rows(), rows.len());
+ batch
+ }
+
+ #[test]
+ fn input_ipc_preserves_scalar_coercion_and_nulls() {
+ let schema = schema(vec![
+ DataType::Integer,
+ DataType::Float,
+ DataType::Boolean,
+ text_type(),
+ DataType::Binary,
+ ]);
+ let rows = vec![
+ row(vec![
+ Value::Long(3),
+ Value::Double(1.25),
+ Value::Bool(true),
+ Value::String("é\0名前".to_string()),
+ Value::Binary(vec![0, 255, 1]),
+ ]),
+ row(vec![Value::Null; 5]),
+ row(vec![
+ Value::Long(-2_147_483_648),
+ Value::Long(16_777_217),
+ Value::Bool(false),
+ Value::String(String::new()),
+ Value::Binary(vec![]),
+ ]),
+ ];
+ let batch = read_batch(&rows, &schema);
+ assert_eq!(batch.num_columns(), 5);
+ assert_eq!(batch.schema().field(0).name(), "schema_0");
+ macro_rules! column_eq {
+ ($index:expr, $array:ty, $values:expr) => {
+ assert_eq!(
+ batch
+ .column($index)
+ .as_any()
+ .downcast_ref::<$array>()
+ .unwrap(),
+ &<$array>::from($values)
+ );
+ };
+ }
+ column_eq!(0, Int32Array, vec![Some(3), None, Some(i32::MIN)]);
+ column_eq!(1, Float32Array, vec![Some(1.25), None,
Some(16_777_216.0)]);
+ column_eq!(2, BooleanArray, vec![Some(true), None, Some(false)]);
+ column_eq!(3, StringArray, vec![Some("é\0名前"), None, Some("")]);
+ column_eq!(
+ 4,
+ BinaryArray,
+ vec![Some(&[0_u8, 255, 1][..]), None, Some(&[][..])]
+ );
+ }
+
+ #[test]
+ fn input_ipc_resolves_duplicate_names_before_numeric_conversion() {
+ let nested_type = DataType::Struct {
+ fields: vec![
+ struct_field("dup", DataType::Float),
+ struct_field("positional", DataType::Short),
+ struct_field("dup", DataType::Double),
+ ],
+ };
+ let nested = Value::Struct(
+ Row::new(
+ vec!["dup".to_string(), "dup".to_string(),
"other".to_string()],
+ vec![
+ Value::Double(16_777_217.0),
+ Value::Long(999),
+ Value::Long(-2),
+ ],
+ )
+ .into(),
+ );
+ let batch = read_batch(&[row(vec![nested])],
&schema(vec![nested_type]));
+ let array = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .unwrap();
+ assert_eq!(
+ array
+ .column(0)
+ .as_any()
+ .downcast_ref::<Float32Array>()
+ .unwrap()
+ .value(0),
+ 16_777_216.0
+ );
+ assert_eq!(
+ array
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int16Array>()
+ .unwrap()
+ .value(0),
+ 999
+ );
+ // Both fields select the first source name, but each conversion must
+ // use the original value: Float rounding must not affect Double.
+ assert_eq!(
+ array
+ .column(2)
+ .as_any()
+ .downcast_ref::<arrow::array::Float64Array>()
+ .unwrap()
+ .value(0),
+ 16_777_217.0
+ );
+ }
+
+ #[test]
+ fn input_ipc_reports_schema_and_column_errors() {
+ let array = DataType::Array {
+ element_type: Box::new(DataType::Long),
+ contains_null: true,
+ };
+ let map = DataType::Map {
+ key_type: Box::new(text_type()),
+ value_type: Box::new(DataType::Long),
+ value_contains_null: true,
+ };
+ for (rows, schema, message) in [
+ (
+ vec![row(vec![Value::Long(1)])],
+ DataType::Long,
+ "Schema is not a struct",
+ ),
+ (
+ vec![row(vec![Value::String("bad".to_string())])],
+ schema(vec![DataType::Integer]),
+ "Type mismatch in row data",
+ ),
+ (
+ vec![row(vec![Value::Long(1), Value::Long(2)])],
+ schema(vec![array, map]),
+ "Type mismatch: expected array",
+ ),
+ // Schema conversion still precedes the width check.
+ (
+ vec![row(vec![Value::Long(1)])],
+ schema(vec![DataType::CalendarInterval, DataType::Long]),
+ "Unsupported type for Arrow createDataFrame conversion",
+ ),
+ ] {
+ let error = rows_to_arrow_ipc(&rows,
&schema).unwrap_err().to_string();
+ assert!(error.contains(message), "{error}");
+ }
+ }
+
+ #[test]
+ fn input_ipc_preserves_empty_input_and_zero_column_error() {
+ assert!(rows_to_arrow_ipc(&[], &DataType::Long).unwrap().is_empty());
+ let error = rows_to_arrow_ipc(&[Row::empty()], &schema(vec![]))
+ .unwrap_err()
+ .to_string();
+ assert!(error.contains("Failed to create Arrow batch"), "{error}");
+ }
+
+ #[test]
+ fn input_ipc_checks_all_widths_before_encoding() {
+ let rows = [
+ row(vec![
+ Value::String("bad integer".to_string()),
+ Value::Bool(true),
+ ]),
+ row(vec![Value::Long(1)]),
+ ];
+ let error = rows_to_arrow_ipc(&rows, &schema(vec![DataType::Integer,
DataType::Boolean]))
+ .unwrap_err()
+ .to_string();
+ assert!(
+ error.contains("Row 1 has 1 values, but the schema has 2 fields"),
+ "{error}"
+ );
+ }
+
+ #[test]
+ fn create_dataframe_rejects_short_and_long_rows() {
+ let spark = session();
+ for values in [vec![Value::Long(1)], vec![Value::Long(1); 3]] {
+ let width = values.len();
+ let result = spark.create_dataframe(
+ vec![row(values)],
+ schema(vec![DataType::Long, DataType::Long]),
+ );
+ let error = result.err().expect("expected a width
error").to_string();
+ assert!(
+ error.contains(&format!(
+ "Row 0 has {width} values, but the schema has 2 fields"
+ )),
+ "{error}"
+ );
+ }
+ }
+
+ #[test]
+ fn input_ipc_roundtrips_through_batch_decoder() {
+ let nested_type = DataType::Struct {
+ fields: vec![
+ struct_field("count", DataType::Long),
+ struct_field("text", text_type()),
+ ],
+ };
+ let nested_value = Value::Struct(
+ Row::new(
+ vec!["count".to_string(), "text".to_string()],
+ vec![Value::Long(3), Value::String("名前".to_string())],
+ )
+ .into(),
+ );
+ for (types, rows) in [
+ (
+ vec![DataType::Long, DataType::Long],
+ vec![
+ row(vec![Value::Long(i64::MIN), Value::Null]),
+ row(vec![Value::Long(i64::MAX), Value::Long(0)]),
+ ],
+ ),
+ (
+ vec![DataType::Long, nested_type],
+ vec![
+ row(vec![Value::Long(-1), nested_value]),
+ row(vec![Value::Null, Value::Null]),
+ ],
+ ),
+ ] {
+ let schema = schema(types);
+ let data = rows_to_arrow_ipc(&rows, &schema).unwrap();
+ let response = proto::execute_plan_response::ArrowBatch {
+ data: data.into(),
+ ..Default::default()
+ };
+ let decoded =
crate::dataframe::decode_arrow_batch(&response).unwrap();
+ assert_eq!(decoded.len(), rows.len());
+ for (actual, expected) in decoded.iter().zip(&rows) {
+ assert_eq!(
+ actual.fields(),
+ &["schema_0".to_string(), "schema_1".to_string()]
+ );
+ assert_eq!(actual.values(), expected.values());
+ }
+ }
+ }
}
diff --git a/crates/spark-connect/tests/e2e_coverage.rs
b/crates/spark-connect/tests/e2e_coverage.rs
index ee73920..801449d 100644
--- a/crates/spark-connect/tests/e2e_coverage.rs
+++ b/crates/spark-connect/tests/e2e_coverage.rs
@@ -1200,7 +1200,7 @@ fn create_dataframe_typed_schemas() {
1
);
- // Float target from an int value (coerce_value must widen Long/Integer ->
Float).
+ // The Arrow encoder converts integer input to the declared Float type.
let f_schema = DataType::Struct {
fields: vec![field("a", DataType::Float)],
};
diff --git a/python/tests/test_dropin_offline.py
b/python/tests/test_dropin_offline.py
index 6c74762..67ccaa4 100644
--- a/python/tests/test_dropin_offline.py
+++ b/python/tests/test_dropin_offline.py
@@ -2729,3 +2729,29 @@ def test_structtype_tree_string_max_depth():
def test_variant_type_to_internal_accepts_variant_kwarg():
# Reference names the VariantType.toInternal parameter `variant`, not
`obj`.
assert T.VariantType().toInternal(variant=None) is None
+
+
+# ---------------------------------------------------------------------------
createDataFrame
+
+def test_create_dataframe_width_mismatch_is_a_normal_exception():
+ # Isolate the active session from the remaining offline tests. Session
+ # construction and local input validation do not need a running server.
+ import subprocess
+ import sys
+
+ subprocess.run([sys.executable, "-c", """
+from pyspark.sql import SparkSession, types as T
+from pyspark.errors import PySparkValueError
+spark = SparkSession.builder.remote("sc://localhost:15002").getOrCreate()
+schemas = [None, ["a", "b"], T.StructType([
+ T.StructField("a", T.LongType()), T.StructField("b", T.LongType())
+])]
+for schema in schemas:
+ try:
+ spark.createDataFrame([(1, 2), (3,)], schema=schema)
+ except PySparkValueError as error:
+ assert type(error) is PySparkValueError, type(error)
+ assert str(error) == "Row 1 has 1 values, but the schema has 2
fields", error
+ else:
+ raise AssertionError("expected a width validation error")
+"""], check=True)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]