dbtsai commented on code in PR #109:
URL:
https://github.com/apache/spark-connect-rust/pull/109#discussion_r4404606858
##########
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"),
Review Comment:
Changed in f26ea7e to assert only our `Failed to create Arrow batch` prefix.
--
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]
