dbtsai commented on code in PR #109:
URL:
https://github.com/apache/spark-connect-rust/pull/109#discussion_r4404606000
##########
crates/spark-connect/src/session.rs:
##########
@@ -1613,3 +1611,277 @@ mod tests {
assert!(tags.contains(&"test".to_string()));
}
}
+
+#[cfg(test)]
+mod input_ipc_tests {
+ // Keep IPC encoding fixtures separate from the session API tests above.
Review Comment:
Done in f26ea7e: moved the IPC fixtures into the existing `mod tests` and
reused `struct_field`; the separate module and duplicate helper are gone.
##########
crates/spark-connect/src/session.rs:
##########
@@ -1613,3 +1611,277 @@ mod tests {
assert!(tags.contains(&"test".to_string()));
}
}
+
+#[cfg(test)]
+mod input_ipc_tests {
+ // Keep IPC encoding fixtures separate from the session API tests above.
+ use super::*;
+ use arrow::array::{
+ BinaryArray, BooleanArray, Float32Array, Int16Array, Int32Array,
Int64Array, Int8Array,
+ StringArray, StructArray,
+ };
+ use arrow::ipc::reader::StreamReader;
+ use arrow::record_batch::RecordBatch;
+ use std::io::Cursor;
+
+ fn field(name: &str, data_type: DataType) -> crate::types::StructField {
+ crate::types::StructField {
+ name: name.to_string(),
+ data_type,
+ nullable: true,
+ metadata: Default::default(),
+ }
+ }
+
+ 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)| 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() {
+ // Compatibility check for this optimization: integer overflow
currently
+ // wraps instead of rejecting out-of-range input as PySpark does. These
+ // expectations pin existing behavior, not the desired validation
policy.
+ let schema = schema(vec![
+ DataType::Integer,
+ DataType::Float,
+ DataType::Boolean,
+ text_type(),
+ DataType::Binary,
+ ]);
+ let rows = vec![
+ row(vec![
+ Value::Long((1_i64 << 32) + 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_649),
+ 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(2_147_483_647)]);
+ 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_composes_nested_coercion_and_duplicate_name_lookup() {
+ let nested_type = DataType::Struct {
+ fields: vec![
+ field("dup", DataType::Byte),
+ field("positional", DataType::Short),
+ field("dup", DataType::Long),
+ ],
+ };
+ let nested = Value::Struct(
+ Row::new(
+ vec!["dup".to_string(), "dup".to_string(),
"other".to_string()],
+ vec![Value::Long(257), 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::<Int8Array>()
+ .unwrap()
+ .value(0),
+ 1
+ );
+ assert_eq!(
+ array
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int16Array>()
+ .unwrap()
+ .value(0),
+ 999
+ );
+ // Compatibility check for an existing bug, not desired semantics:
+ // encoding resolves both duplicate names to the first coerced child,
+ // so the Long field receives the value already narrowed to Byte.
+ assert_eq!(
+ array
+ .column(2)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap()
+ .value(0),
+ 1
+ );
+ }
+
+ #[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("must either specify a row count or at least one
column"),
+ "{error}"
+ );
+ }
+
+ #[test]
+ #[should_panic(expected = "field names and values must have the same
length")]
Review Comment:
Kept the ordering test in f26ea7e and converted it to `unwrap_err()` plus a
width-error message assertion. The earlier row still contains an invalid
integer, so the test checks that the later width mismatch is reported before
encoding starts. Added public Rust API coverage for both short and long rows as
well.
##########
crates/spark-connect/src/session.rs:
##########
@@ -1613,3 +1611,277 @@ mod tests {
assert!(tags.contains(&"test".to_string()));
}
}
+
+#[cfg(test)]
+mod input_ipc_tests {
+ // Keep IPC encoding fixtures separate from the session API tests above.
+ use super::*;
+ use arrow::array::{
+ BinaryArray, BooleanArray, Float32Array, Int16Array, Int32Array,
Int64Array, Int8Array,
+ StringArray, StructArray,
+ };
+ use arrow::ipc::reader::StreamReader;
+ use arrow::record_batch::RecordBatch;
+ use std::io::Cursor;
+
+ fn field(name: &str, data_type: DataType) -> crate::types::StructField {
+ crate::types::StructField {
+ name: name.to_string(),
+ data_type,
+ nullable: true,
+ metadata: Default::default(),
+ }
+ }
+
+ 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)| 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() {
+ // Compatibility check for this optimization: integer overflow
currently
+ // wraps instead of rejecting out-of-range input as PySpark does. These
+ // expectations pin existing behavior, not the desired validation
policy.
+ let schema = schema(vec![
+ DataType::Integer,
+ DataType::Float,
+ DataType::Boolean,
+ text_type(),
+ DataType::Binary,
+ ]);
+ let rows = vec![
+ row(vec![
+ Value::Long((1_i64 << 32) + 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_649),
+ 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(2_147_483_647)]);
+ 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_composes_nested_coercion_and_duplicate_name_lookup() {
+ let nested_type = DataType::Struct {
+ fields: vec![
+ field("dup", DataType::Byte),
+ field("positional", DataType::Short),
+ field("dup", DataType::Long),
+ ],
+ };
+ let nested = Value::Struct(
+ Row::new(
+ vec!["dup".to_string(), "dup".to_string(),
"other".to_string()],
+ vec![Value::Long(257), 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::<Int8Array>()
+ .unwrap()
+ .value(0),
+ 1
+ );
+ assert_eq!(
+ array
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int16Array>()
+ .unwrap()
+ .value(0),
+ 999
+ );
+ // Compatibility check for an existing bug, not desired semantics:
Review Comment:
Agreed. Removed the coercion-only test along with the helper, and replaced
this composition test in f26ea7e with an in-range precision regression:
duplicate Float/Double fields read 16,777,217 independently, yielding
16,777,216 as Float and 16,777,217 as Double. This checks that one field's
rounding cannot affect the other.
##########
crates/spark-connect/src/session.rs:
##########
@@ -815,12 +809,16 @@ fn build_arrow_array(
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),
+ Some(dt) => {
+ // Coerce only this column, avoiding a second complete row matrix
and
+ // a deep clone of each coerced value before Arrow construction.
+ let column: Vec<Value> = rows
+ .iter()
+ .map(|r| coerce_value(r.get(field_idx), dt))
+ .collect();
+ 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),
Review Comment:
Done in f26ea7e: inlined column gathering into `rows_to_arrow_ipc` and
removed `build_arrow_array`, the unused scalar builder, and the redundant
empty-input branch.
--
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]
