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 fbd6af5 [SPARK-60585][SPARK-60588][CONNECT] Share field names across
Arrow rows and structs
fbd6af5 is described below
commit fbd6af5763bae5cd572ba4438026e675165ba70a
Author: DB Tsai <[email protected]>
AuthorDate: Tue Oct 27 14:30:18 2026 -0700
[SPARK-60585][SPARK-60588][CONNECT] Share field names across Arrow rows and
structs
JIRA: [SPARK-60585](https://issues.apache.org/jira/browse/SPARK-60585),
[SPARK-60588](https://issues.apache.org/jira/browse/SPARK-60588)
### What changes were proposed in this pull request?
Share immutable names across Arrow-decoded rows and nested structs,
including structs inside arrays/maps. Local Row constructors retain owned
names; Catalog and DataFrame use one decoder.
### Why are the changes needed?
Apple M3 Max/macOS arm64; Apache master `a85e474` → combined PR head
`834402b` (top-level and nested sharing; no incremental nested-only
comparison). The [benchmark
harness](https://github.com/dbtsai/pyspark-connect-benchmark) uses canned
Arrow/gRPC fixtures, 23 frozen cases, 30 fresh processes per case/client and
three rotating clients. screen: 2,070 recorded, 2,010 successful, 60 failed, 0
missing; confirm: 2,070 recorded, 2,010 successful, 60 failed, 0 missing.
Selected confirm medians, baseline → candidate; scans fully consume output.
Held RSS includes allocator history.
| Workload | CPU ms | Wall ms | Held RSS MiB |
| --- | ---: | ---: | ---: |
| 16 × 256, nested d4/inner16, full scan | 94.08 → 88.01 | 96.52 → 90.42 |
124.95 → 120.66 |
| 16 × 256, nested d16/inner8, full scan | 202.87 → 185.94 | 207.18 →
190.45 | 145.41 → 140.02 |
| 100,000 × 50 collect | 242.12 → 127.38 | 234.48 → 118.91 | 680.85 →
447.13 |
| 4,096 × 512 collect | 87.18 → 39.68 | 83.02 → 35.42 | 297.91 → 214.72 |
**Overall frozen verdict: rejected.** Arrow `Duration(Microsecond)` is
unsupported in both Rust builds (30 failed mixed-record samples each per run);
timing guards also remain uncertain. No accepted overall CPU or memory
improvement is claimed. Independent full confirmation is included.
### Does this PR introduce _any_ user-facing change?
Rust API break: `Value::Struct` now holds public `StructValue` instead of
`Vec<(String, Value)>`, sharing names while owning values. Migrate construction
to `Value::Struct(Row::new(fields, values).into())` and use its field/value
accessors. Tested Python Row behavior is unchanged. Measured arm64 `Row` and
`Value` both remain 48 bytes.
### How was this patch tested?
412 Rust checks passed; all core integration targets compiled with and
without `wasm-udf`. The 108-class Python API snapshot and 11 behavior/fixture
comparisons pass, as do formatting and core Clippy. Local workspace Clippy
lacks the wasm standard library; local live-server parity was unavailable.
Exact-head CI: 12 checks passed; Connect parity is still running; docs
deployment was skipped.
### Was this patch authored or co-authored using generative AI tooling?
Yes. Generated-by: OpenAI Codex CLI 0.163.0-alpha.10.2.
This pull request and its description were written by Isaac.
Closes #105 from dbtsai/codex/share-arrow-row-field-names.
Authored-by: DB Tsai <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>
---
crates/pyspark-rs/src/row.rs | 11 +-
crates/pyspark-rs/src/session.rs | 18 +-
crates/spark-connect/src/catalog.rs | 99 ++++---
crates/spark-connect/src/dataframe.rs | 479 ++++++++++++++++++++++++++++++++--
crates/spark-connect/src/row.rs | 424 +++++++++++++++++++++++++++++-
crates/spark-connect/src/session.rs | 263 +++++++++++++++++--
6 files changed, 1167 insertions(+), 127 deletions(-)
diff --git a/crates/pyspark-rs/src/row.rs b/crates/pyspark-rs/src/row.rs
index 524bb73..dfdf506 100644
--- a/crates/pyspark-rs/src/row.rs
+++ b/crates/pyspark-rs/src/row.rs
@@ -70,15 +70,10 @@ pub fn value_to_py<'py>(py: Python<'py>, v: &Value) ->
PyResult<Bound<'py, PyAny
}
dict.into_any()
}
- Value::Struct(fields) => {
+ Value::Struct(row) => {
// 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());
- }
- Py::new(py, PyRow::new(spark_connect::row::Row::new(names, vals)))?
+ // rather than a plain dict. Cloning keeps shared names and
independent values.
+ Py::new(py, PyRow::new(row.clone().into()))?
.into_bound(py)
.into_any()
}
diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs
index 437cc80..5650e60 100644
--- a/crates/pyspark-rs/src/session.rs
+++ b/crates/pyspark-rs/src/session.rs
@@ -915,9 +915,11 @@ fn value_to_datatype(v: &Value) -> DataType {
contains_null: true,
}
}
- Value::Struct(fields) => DataType::Struct {
- fields: fields
+ Value::Struct(row) => DataType::Struct {
+ fields: row
+ .fields()
.iter()
+ .zip(row.values())
.map(|(n, val)| spark_connect::types::StructField {
name: n.clone(),
data_type: value_to_datatype(val),
@@ -1004,17 +1006,7 @@ pub(crate) fn py_to_value(obj: &Bound<'_, PyAny>) ->
PyResult<Value> {
// pyspark Row -> struct: its named fields (values already core Values),
so a nested
// Row inside createDataFrame data (e.g. Row(s=Row(x=1))) converts
recursively.
if let Ok(pyrow) = obj.extract::<PyRef<crate::row::PyRow>>() {
- let fields = pyrow.row.fields();
- let values = pyrow.row.values();
- let mut out = Vec::with_capacity(values.len());
- for (i, v) in values.iter().enumerate() {
- let name = fields
- .get(i)
- .cloned()
- .unwrap_or_else(|| format!("_{}", i + 1));
- out.push((name, v.clone()));
- }
- return Ok(Value::Struct(out));
+ return Ok(Value::Struct(pyrow.row.clone().into()));
}
// dict -> map (keys coerced to their string form, values recursively
converted)
diff --git a/crates/spark-connect/src/catalog.rs
b/crates/spark-connect/src/catalog.rs
index 5de93f8..7ce672a 100644
--- a/crates/spark-connect/src/catalog.rs
+++ b/crates/spark-connect/src/catalog.rs
@@ -6,7 +6,7 @@ use spark_connect_core::error::{Result, SparkError};
use spark_connect_core::runtime::block_on;
use spark_connect_proto as proto;
-use crate::dataframe::DataFrame;
+use crate::dataframe::{decode_arrow_batch, DataFrame};
use crate::row::{Row, Value};
use crate::session::SparkSession;
@@ -1083,55 +1083,54 @@ impl Catalog {
}
}
-/// Decode an Arrow batch into rows.
-pub(crate) fn decode_arrow_batch(
- batch: &proto::execute_plan_response::ArrowBatch,
-) -> Result<Vec<Row>> {
- use arrow::ipc::reader::StreamReader;
- use std::io::Cursor;
-
- if batch.data.is_empty() {
- return Ok(vec![]);
- }
-
- let cursor = Cursor::new(&batch.data);
- let mut reader = StreamReader::try_new(cursor, None).map_err(|e| {
- SparkError::connect_msg(format!("Failed to create Arrow stream reader:
{}", e))
- })?;
-
- let mut rows = vec![];
-
- while let Some(record_batch) = reader
- .next()
- .transpose()
- .map_err(|e| SparkError::connect_msg(format!("Failed to decode Arrow
batch: {}", e)))?
- {
- let schema = record_batch.schema();
- let num_rows = record_batch.num_rows();
- let num_cols = record_batch.num_columns();
-
- for row_idx in 0..num_rows {
- let mut field_names = vec![];
- let mut values = vec![];
-
- for col_idx in 0..num_cols {
- let field_name = schema.field(col_idx).name().clone();
- let column = record_batch.column(col_idx);
-
- let value = arrow_value_at(column.as_ref(), row_idx)?;
- field_names.push(field_name);
- values.push(value);
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn parse_database_preserves_metadata_and_nullable_columns() {
+ let fields = vec![
+ "name".to_string(),
+ "catalog".to_string(),
+ "description".to_string(),
+ "locationUri".to_string(),
+ ];
+ let populated = Row::new(
+ fields.clone(),
+ vec![
+ Value::String("one".to_string()),
+ Value::String("spark_catalog".to_string()),
+ Value::String("first database".to_string()),
+ Value::String("file:/one".to_string()),
+ ],
+ );
+ let nullable = Row::new(
+ fields,
+ vec![
+ Value::String("two".to_string()),
+ Value::Null,
+ Value::Null,
+ Value::String("file:/two".to_string()),
+ ],
+ );
+
+ assert_eq!(
+ parse_database(&populated),
+ Database {
+ name: "one".to_string(),
+ catalog: Some("spark_catalog".to_string()),
+ description: Some("first database".to_string()),
+ location_uri: "file:/one".to_string(),
}
-
- rows.push(Row::new(field_names, values));
- }
+ );
+ assert_eq!(
+ parse_database(&nullable),
+ Database {
+ name: "two".to_string(),
+ catalog: None,
+ description: None,
+ location_uri: "file:/two".to_string(),
+ }
+ );
}
-
- Ok(rows)
-}
-
-/// Extract a value at a row index from an Arrow array. Delegates to the
single,
-/// comprehensive decoder in `dataframe` so catalog results support every type
too.
-fn arrow_value_at(array: &dyn arrow::array::Array, index: usize) ->
Result<Value> {
- crate::dataframe::arrow_value_at(array, index)
}
diff --git a/crates/spark-connect/src/dataframe.rs
b/crates/spark-connect/src/dataframe.rs
index 0a273cc..72dff2f 100644
--- a/crates/spark-connect/src/dataframe.rs
+++ b/crates/spark-connect/src/dataframe.rs
@@ -13,7 +13,7 @@ use spark_connect_proto as proto;
use crate::column::Column;
use crate::expression::{Expression, ExpressionNode};
use crate::plan::{AggregateGroupType, JoinType, LogicalPlan, SetOpType};
-use crate::row::{Row, Value};
+use crate::row::{Row, StructValue, Value};
use crate::session::{ExecutionInfo, SparkSession};
use crate::types::DataType;
use crate::udf::CommonInlineUserDefinedFunctionExpression;
@@ -2053,7 +2053,9 @@ pub(crate) fn assign_plan_ids(
}
/// Decode an Arrow batch into rows.
-fn decode_arrow_batch(batch: &proto::execute_plan_response::ArrowBatch) ->
Result<Vec<Row>> {
+pub(crate) fn decode_arrow_batch(
+ batch: &proto::execute_plan_response::ArrowBatch,
+) -> Result<Vec<Row>> {
use arrow::ipc::reader::StreamReader;
use std::io::Cursor;
@@ -2075,22 +2077,29 @@ fn decode_arrow_batch(batch:
&proto::execute_plan_response::ArrowBatch) -> Resul
{
let schema = record_batch.schema();
let num_rows = record_batch.num_rows();
- let num_cols = record_batch.num_columns();
+ if num_rows == 0 {
+ continue;
+ }
+ let field_names = schema
+ .fields()
+ .iter()
+ .map(|field| field.name().clone())
+ .collect::<Arc<[String]>>();
+ let layouts = record_batch
+ .columns()
+ .iter()
+ .map(|column| ArrowValueLayout::new(column.data_type()))
+ .collect::<Vec<_>>();
for row_idx in 0..num_rows {
- let mut field_names = vec![];
let mut values = vec![];
- for col_idx in 0..num_cols {
- let field_name = schema.field(col_idx).name().clone();
- let column = record_batch.column(col_idx);
-
- let value = arrow_value_at(column.as_ref(), row_idx)?;
- field_names.push(field_name);
+ for (column, layout) in
record_batch.columns().iter().zip(&layouts) {
+ let value = arrow_value_at_with_layout(column.as_ref(),
row_idx, layout)?;
values.push(value);
}
- rows.push(Row::new(field_names, values));
+ rows.push(Row::with_shared_fields(Arc::clone(&field_names),
values));
}
}
@@ -2232,7 +2241,60 @@ fn map_key_to_string(v: Value) -> String {
}
}
-pub(crate) fn arrow_value_at(array: &dyn arrow::array::Array, index: usize) ->
Result<Value> {
+/// Immutable decoding metadata prepared once for each column in a record
batch.
+/// Container layouts follow their children so struct names are also shared
inside
+/// lists and maps, without retaining the Arrow arrays after decoding.
+enum ArrowValueLayout {
+ Scalar,
+ List(Box<Self>),
+ Struct {
+ names: Arc<[String]>,
+ children: Vec<Self>,
+ is_variant: bool,
+ },
+ Map(Box<(Self, Self)>),
+}
+
+impl ArrowValueLayout {
+ fn new(data_type: &arrow::datatypes::DataType) -> Self {
+ use arrow::datatypes::DataType;
+ match data_type {
+ DataType::List(field) =>
Self::List(Box::new(Self::new(field.data_type()))),
+ DataType::Struct(fields) => Self::Struct {
+ names: fields.iter().map(|field|
field.name().clone()).collect(),
+ children: fields
+ .iter()
+ .map(|field| Self::new(field.data_type()))
+ .collect(),
+ is_variant: fields.iter().any(|field| {
+ field
+ .metadata()
+ .get("variant")
+ .is_some_and(|value| value == "true")
+ }),
+ },
+ DataType::Map(entries, _) => match entries.data_type() {
+ DataType::Struct(fields) => Self::Map(Box::new((
+ Self::new(fields[0].data_type()),
+ Self::new(fields[1].data_type()),
+ ))),
+ _ => unreachable!("Arrow map entries must be a struct"),
+ },
+ _ => Self::Scalar,
+ }
+ }
+}
+
+#[cfg(test)]
+fn arrow_value_at(array: &dyn arrow::array::Array, index: usize) ->
Result<Value> {
+ arrow_value_at_with_layout(array, index,
&ArrowValueLayout::new(array.data_type()))
+}
+
+fn arrow_value_at_with_layout(
+ array: &dyn arrow::array::Array,
+ index: usize,
+ value_layout: &ArrowValueLayout,
+) -> Result<Value> {
use arrow::array::*;
if array.is_null(index) {
@@ -2332,28 +2394,37 @@ pub(crate) fn arrow_value_at(array: &dyn
arrow::array::Array, index: usize) -> R
}
// Nested: list, struct, map (recurse).
if let Some(arr) = array.as_any().downcast_ref::<ListArray>() {
+ let ArrowValueLayout::List(child_layout) = value_layout else {
+ unreachable!("layout is prepared from the array's data type");
+ };
let child = arr.value(index);
let mut items = Vec::with_capacity(child.len());
for i in 0..child.len() {
- items.push(arrow_value_at(child.as_ref(), i)?);
+ items.push(arrow_value_at_with_layout(child.as_ref(), i,
child_layout)?);
}
return Ok(Value::List(items));
}
if let Some(arr) = array.as_any().downcast_ref::<StructArray>() {
+ let ArrowValueLayout::Struct {
+ names,
+ children,
+ is_variant,
+ } = value_layout
+ else {
+ unreachable!("layout is prepared from the array's data type");
+ };
// A VARIANT column arrives as struct<value: binary, metadata: binary>
where the
// `metadata` field carries arrow metadata {"variant": "true"}.
Recognize it and
// return a Value::Variant (raw bytes) so it materializes as a
VariantVal (matching
// pyspark) rather than a plain {value, metadata} struct/dict.
- let is_variant = arr.fields().iter().any(|f| {
- f.metadata()
- .get("variant")
- .map(|v| v == "true")
- .unwrap_or(false)
- });
- if is_variant {
+ if *is_variant {
let bin_field = |name: &str| -> Result<Vec<u8>> {
- match arr.column_by_name(name) {
- Some(col) => match arrow_value_at(col.as_ref(), index)? {
+ match arr.fields().iter().position(|field| field.name() ==
name) {
+ Some(col_idx) => match arrow_value_at_with_layout(
+ arr.column(col_idx).as_ref(),
+ index,
+ &children[col_idx],
+ )? {
Value::Binary(b) => Ok(b),
Value::Null => Ok(vec![]),
_ => Err(SparkError::connect_msg("variant field is not
binary")),
@@ -2366,13 +2437,20 @@ pub(crate) fn arrow_value_at(array: &dyn
arrow::array::Array, index: usize) -> R
metadata: bin_field("metadata")?,
});
}
- let mut fields = Vec::new();
- for (f, col) in arr.fields().iter().zip(arr.columns()) {
- fields.push((f.name().clone(), arrow_value_at(col.as_ref(),
index)?));
+ let mut values = Vec::with_capacity(children.len());
+ for (col, child_layout) in arr.columns().iter().zip(children) {
+ values.push(arrow_value_at_with_layout(
+ col.as_ref(),
+ index,
+ child_layout,
+ )?);
}
- return Ok(Value::Struct(fields));
+ return Ok(Value::Struct(StructValue::new(Arc::clone(names), values)));
}
if let Some(arr) = array.as_any().downcast_ref::<MapArray>() {
+ let ArrowValueLayout::Map(children) = value_layout else {
+ unreachable!("layout is prepared from the array's data type");
+ };
let entries = arr.value(index);
let keys = entries.column(0);
let vals = entries.column(1);
@@ -2381,8 +2459,11 @@ pub(crate) fn arrow_value_at(array: &dyn
arrow::array::Array, index: usize) -> R
// A map key stringifies to its natural scalar form (e.g. `1`,
`true`,
// `1.5`), not the Rust enum's Debug output - a `map<int,string>`
key must
// be "1", never "Integer(1)".
- let k = map_key_to_string(arrow_value_at(keys.as_ref(), i)?);
- map.insert(k, arrow_value_at(vals.as_ref(), i)?);
+ let k =
map_key_to_string(arrow_value_at_with_layout(keys.as_ref(), i, &children.0)?);
+ map.insert(
+ k,
+ arrow_value_at_with_layout(vals.as_ref(), i, &children.1)?,
+ );
}
return Ok(Value::Map(map));
}
@@ -2446,6 +2527,312 @@ pub(crate) fn arrow_value_at(array: &dyn
arrow::array::Array, index: usize) -> R
)))
}
+#[cfg(test)]
+mod nested_decode_tests {
+ use super::*;
+ use arrow::array::*;
+ use arrow::buffer::{NullBuffer, OffsetBuffer};
+ use arrow::datatypes::{DataType, Field, Schema};
+ use arrow::ipc::writer::StreamWriter;
+ use arrow::record_batch::RecordBatch;
+
+ fn decode_batches(columns: Vec<(&str, ArrayRef)>, repeats: usize) ->
Vec<Row> {
+ let schema = Arc::new(Schema::new(
+ columns
+ .iter()
+ .map(|(name, array)| Field::new(*name,
array.data_type().clone(), true))
+ .collect::<Vec<_>>(),
+ ));
+ let batch = RecordBatch::try_new(
+ schema.clone(),
+ columns.into_iter().map(|(_, array)| array).collect(),
+ )
+ .unwrap();
+ let mut data = vec![];
+ {
+ let mut writer = StreamWriter::try_new(&mut data,
&schema).unwrap();
+ for _ in 0..repeats {
+ writer.write(&batch).unwrap();
+ }
+ writer.finish().unwrap();
+ }
+ decode_arrow_batch(&proto::execute_plan_response::ArrowBatch {
+ data: data.into(),
+ ..Default::default()
+ })
+ .unwrap()
+ }
+
+ fn as_struct(value: &Value) -> &StructValue {
+ let Value::Struct(row) = value else {
+ panic!("expected struct, got {value:?}");
+ };
+ row
+ }
+
+ #[test]
+ fn nested_names_preserve_duplicates_unicode_nulls_and_sibling_layouts() {
+ let names = ["dup", "dup", "", "名前"];
+ let fields = names
+ .iter()
+ .map(|name| Arc::new(Field::new(*name, DataType::Int32, true)))
+ .collect::<Vec<_>>();
+ let columns = (1..=4)
+ .map(|n| Arc::new(Int32Array::from(vec![Some(n), None, None])) as
ArrayRef)
+ .collect();
+ let nested = StructArray::new(
+ fields.into(),
+ columns,
+ Some(NullBuffer::from(vec![true, true, false])),
+ );
+ let sibling = StructArray::from(vec![(
+ Arc::new(Field::new("other", DataType::Int32, false)),
+ Arc::new(Int32Array::from(vec![7, 8, 9])) as ArrayRef,
+ )]);
+ let rows = decode_batches(
+ vec![("nested", Arc::new(nested)), ("sibling", Arc::new(sibling))],
+ 2,
+ );
+ assert_eq!(rows.len(), 6);
+ for batch in rows.chunks(3) {
+ let populated = as_struct(batch[0].get(0).unwrap());
+ let all_null = as_struct(batch[1].get(0).unwrap());
+ assert_eq!(populated.fields(), &names);
+ assert_eq!(
+ populated.values(),
+ &[
+ Value::Integer(1),
+ Value::Integer(2),
+ Value::Integer(3),
+ Value::Integer(4)
+ ]
+ );
+ assert_eq!(populated.get_by_name("dup"), Some(&Value::Integer(1)));
+ assert_eq!(all_null.values(), &vec![Value::Null; 4]);
+ assert_eq!(batch[2].get(0), Some(&Value::Null));
+ assert!(std::ptr::eq(populated.fields(), all_null.fields()));
+ let other = as_struct(batch[0].get(1).unwrap());
+ assert_eq!(other.fields(), &["other"]);
+ assert_eq!(other.values(), &[Value::Integer(7)]);
+ }
+ }
+
+ #[test]
+ fn names_are_shared_through_lists_maps_and_nested_structs() {
+ let leaf = StructArray::from(vec![(
+ Arc::new(Field::new("leaf", DataType::Int32, false)),
+ Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as ArrayRef,
+ )]);
+ let entries = StructArray::from(vec![
+ (
+ Arc::new(Field::new("key", DataType::Utf8, false)),
+ Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as
ArrayRef,
+ ),
+ (
+ Arc::new(Field::new("value", leaf.data_type().clone(), true)),
+ Arc::new(leaf) as ArrayRef,
+ ),
+ ]);
+ let map = MapArray::new(
+ Arc::new(Field::new("entries", entries.data_type().clone(),
false)),
+ OffsetBuffer::new(vec![0, 1, 2, 3, 4].into()),
+ entries,
+ None,
+ false,
+ );
+ let list = ListArray::new(
+ Arc::new(Field::new("item", map.data_type().clone(), true)),
+ OffsetBuffer::new(vec![0, 2, 4].into()),
+ Arc::new(map),
+ None,
+ );
+ let outer = StructArray::from(vec![(
+ Arc::new(Field::new("items", list.data_type().clone(), true)),
+ Arc::new(list) as ArrayRef,
+ )]);
+ let rows = decode_batches(vec![("outer", Arc::new(outer))], 1);
+ let mut leaves = vec![];
+ for row in &rows {
+ let outer = as_struct(row.get(0).unwrap());
+ let Value::List(maps) = outer.get(0).unwrap() else {
+ panic!("expected list");
+ };
+ for value in maps {
+ let Value::Map(map) = value else {
+ panic!("expected map");
+ };
+ leaves.push(as_struct(map.values().next().unwrap()));
+ }
+ }
+ assert_eq!(leaves.len(), 4);
+ for (index, leaf) in leaves.iter().enumerate() {
+ assert_eq!(leaf.fields(), &["leaf"]);
+ assert_eq!(leaf.get(0), Some(&Value::Integer(index as i32 + 1)));
+ assert!(std::ptr::eq(leaves[0].fields(), leaf.fields()));
+ }
+ }
+
+ #[test]
+ fn sliced_containers_preserve_offsets_empty_null_and_struct_values() {
+ let leaf = StructArray::new(
+ vec![Arc::new(Field::new("n", DataType::Int32, true))].into(),
+ vec![Arc::new(Int32Array::from(vec![
+ Some(0),
+ Some(1),
+ None,
+ Some(3),
+ Some(4),
+ Some(5),
+ ]))],
+ Some(NullBuffer::from(vec![true, true, false, true, true, true])),
+ );
+ let offsets = OffsetBuffer::new(vec![0, 1, 3, 3, 4, 6].into());
+ let valid = Some(NullBuffer::from(vec![true, true, true, false,
true]));
+ let list = ListArray::new(
+ Arc::new(Field::new("item", leaf.data_type().clone(), true)),
+ offsets.clone(),
+ Arc::new(leaf.clone()),
+ valid.clone(),
+ );
+ let entries = StructArray::from(vec![
+ (
+ Arc::new(Field::new("key", DataType::Utf8, false)),
+ Arc::new(StringArray::from(vec!["k0", "k1", "k2", "k3", "k4",
"k5"])) as ArrayRef,
+ ),
+ (
+ Arc::new(Field::new("value", leaf.data_type().clone(), true)),
+ Arc::new(leaf) as ArrayRef,
+ ),
+ ]);
+ let map = MapArray::new(
+ Arc::new(Field::new("entries", entries.data_type().clone(),
false)),
+ offsets,
+ entries,
+ valid,
+ false,
+ );
+ let rows = decode_batches(
+ vec![
+ ("list", Arc::new(list.slice(1, 4))),
+ ("map", Arc::new(map.slice(1, 4))),
+ ],
+ 1,
+ );
+ let value =
+ |n| Value::Struct(Row::new(vec!["n".to_string()],
vec![Value::Integer(n)]).into());
+ assert_eq!(
+ rows[0].get(0),
+ Some(&Value::List(vec![value(1), Value::Null]))
+ );
+ assert_eq!(
+ rows[0].get(1),
+ Some(&Value::Map(
+ [
+ ("k1".to_string(), value(1)),
+ ("k2".to_string(), Value::Null)
+ ]
+ .into()
+ ))
+ );
+ assert_eq!(rows[1].get(0), Some(&Value::List(vec![])));
+ assert_eq!(rows[1].get(1), Some(&Value::Map(Default::default())));
+ assert_eq!(rows[2].get(0), Some(&Value::Null));
+ assert_eq!(rows[2].get(1), Some(&Value::Null));
+ assert_eq!(rows[3].get(0), Some(&Value::List(vec![value(4),
value(5)])));
+ assert_eq!(
+ rows[3].get(1),
+ Some(&Value::Map(
+ [("k4".to_string(), value(4)), ("k5".to_string(),
value(5))].into()
+ ))
+ );
+ let (Value::List(first), Value::List(last)) =
+ (rows[0].get(0).unwrap(), rows[3].get(0).unwrap())
+ else {
+ panic!("expected lists");
+ };
+ assert!(std::ptr::eq(
+ as_struct(&first[0]).fields(),
+ as_struct(&last[0]).fields()
+ ));
+ let (Value::Map(first), Value::Map(last)) =
+ (rows[0].get(1).unwrap(), rows[3].get(1).unwrap())
+ else {
+ panic!("expected maps");
+ };
+ assert!(std::ptr::eq(
+ as_struct(&first["k1"]).fields(),
+ as_struct(&last["k4"]).fields()
+ ));
+ }
+
+ #[test]
+ fn variant_marker_precedes_struct_decoding_and_preserves_errors() {
+ let marked = |name: &str, array: ArrayRef| {
+ let field = Field::new(name, array.data_type().clone(),
true).with_metadata(
+ std::collections::HashMap::from([("variant".into(),
"true".into())]),
+ );
+ StructArray::from(vec![(Arc::new(field), array)])
+ };
+ let binary = marked(
+ "value",
+ Arc::new(BinaryArray::from(vec![Some(b"x".as_ref()), None])),
+ );
+ assert_eq!(
+ arrow_value_at(&binary, 0).unwrap(),
+ Value::Variant {
+ value: b"x".to_vec(),
+ metadata: vec![]
+ }
+ );
+ assert_eq!(
+ arrow_value_at(&binary, 1).unwrap(),
+ Value::Variant {
+ value: vec![],
+ metadata: vec![]
+ }
+ );
+ let missing = marked("unrelated", Arc::new(Int32Array::from(vec![1])));
+ assert_eq!(
+ arrow_value_at(&missing, 0).unwrap(),
+ Value::Variant {
+ value: vec![],
+ metadata: vec![]
+ }
+ );
+ let invalid = marked("metadata", Arc::new(Int32Array::from(vec![1])));
+ assert!(arrow_value_at(&invalid, 0)
+ .unwrap_err()
+ .to_string()
+ .contains("variant field is not binary"));
+ let null_parent = StructArray::new(
+ invalid.fields().clone(),
+ invalid.columns().to_vec(),
+ Some(NullBuffer::from(vec![false])),
+ );
+ assert_eq!(arrow_value_at(&null_parent, 0).unwrap(), Value::Null);
+ let ordinary = StructArray::from(vec![
+ (
+ Arc::new(Field::new("value", DataType::Int32, false)),
+ Arc::new(Int32Array::from(vec![1])) as ArrayRef,
+ ),
+ (
+ Arc::new(Field::new("metadata", DataType::Int32, false)),
+ Arc::new(Int32Array::from(vec![2])) as ArrayRef,
+ ),
+ ]);
+ assert_eq!(
+ arrow_value_at(&ordinary, 0).unwrap(),
+ Value::Struct(
+ Row::new(
+ vec!["value".to_string(), "metadata".to_string()],
+ vec![Value::Integer(1), Value::Integer(2)]
+ )
+ .into()
+ )
+ );
+ }
+}
+
/// Render microseconds-since-midnight as an ISO time string
`HH:MM:SS[.ffffff]`.
fn micros_to_time_string(micros: i64) -> String {
let total_secs = micros.div_euclid(1_000_000);
@@ -2539,6 +2926,42 @@ mod conversion_tests {
.unwrap()
}
+ #[test]
+ fn decoded_rows_share_batch_field_names_and_preserve_values() {
+ use arrow::ipc::writer::StreamWriter;
+
+ let batch = sample_batch();
+ let mut data = Vec::new();
+ {
+ let mut writer = StreamWriter::try_new(&mut data,
batch.schema().as_ref()).unwrap();
+ writer.write(&batch.slice(0, 0)).unwrap();
+ writer.write(&batch).unwrap();
+ writer.write(&batch).unwrap();
+ writer.finish().unwrap();
+ }
+ let response = proto::execute_plan_response::ArrowBatch {
+ data: data.into(),
+ ..Default::default()
+ };
+ let rows = decode_arrow_batch(&response).unwrap();
+
+ assert_eq!(rows.len(), 6);
+ for (index, row) in rows.iter().enumerate() {
+ let source_index = index % 3;
+ assert_eq!(row.fields(), &["id".to_string(), "name".to_string()]);
+ assert_eq!(row.get(0), Some(&Value::Long(source_index as i64 +
1)));
+ assert_eq!(
+ row.get_by_name("name"),
+ Some(&Value::String(["a", "b", "c"][source_index].to_string()))
+ );
+ }
+ assert!(std::ptr::eq(rows[0].fields(), rows[1].fields()));
+ assert!(std::ptr::eq(rows[1].fields(), rows[2].fields()));
+ assert!(std::ptr::eq(rows[3].fields(), rows[4].fields()));
+ assert!(std::ptr::eq(rows[4].fields(), rows[5].fields()));
+ assert!(decode_arrow_batch(&Default::default()).unwrap().is_empty());
+ }
+
#[test]
fn to_arrow_ipc_round_trips() {
use arrow::ipc::reader::FileReader;
diff --git a/crates/spark-connect/src/row.rs b/crates/spark-connect/src/row.rs
index 3a5837e..a4b3f52 100644
--- a/crates/spark-connect/src/row.rs
+++ b/crates/spark-connect/src/row.rs
@@ -5,9 +5,10 @@
use std::collections::BTreeMap;
use std::fmt;
+use std::sync::Arc;
/// A Value in a Row, supporting all Spark SQL types.
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Clone, PartialEq)]
pub enum Value {
/// NULL value
Null,
@@ -43,14 +44,65 @@ pub enum Value {
List(Vec<Value>),
/// Map of key-value pairs
Map(BTreeMap<String, Value>),
- /// Struct (nested Row)
- Struct(Vec<(String, Value)>),
+ /// Struct (nested Row), with ordered names and values.
+ /// Construct with `Value::Struct(Row::new(fields, values).into())`.
+ Struct(StructValue),
/// A VARIANT value carried as its raw (value, metadata) binary
components, matching
/// `pyspark.sql.types.VariantVal(value, metadata)`. Decoding to
JSON/Python is done
/// lazily on the Python side (VariantVal.toJson/toPython via
variant_utils).
Variant { value: Vec<u8>, metadata: Vec<u8> },
}
+impl fmt::Debug for Value {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ match self {
+ Self::Null => write!(f, "Null"),
+ Self::Bool(v) => f.debug_tuple("Bool").field(v).finish(),
+ Self::Byte(v) => f.debug_tuple("Byte").field(v).finish(),
+ Self::Short(v) => f.debug_tuple("Short").field(v).finish(),
+ Self::Integer(v) => f.debug_tuple("Integer").field(v).finish(),
+ Self::Long(v) => f.debug_tuple("Long").field(v).finish(),
+ Self::Float(v) => f.debug_tuple("Float").field(v).finish(),
+ Self::Double(v) => f.debug_tuple("Double").field(v).finish(),
+ Self::String(v) => f.debug_tuple("String").field(v).finish(),
+ Self::Binary(v) => f.debug_tuple("Binary").field(v).finish(),
+ Self::Date(v) => f.debug_tuple("Date").field(v).finish(),
+ Self::Timestamp(v) => f.debug_tuple("Timestamp").field(v).finish(),
+ Self::Decimal {
+ value,
+ precision,
+ scale,
+ } => f
+ .debug_struct("Decimal")
+ .field("value", value)
+ .field("precision", precision)
+ .field("scale", scale)
+ .finish(),
+ Self::List(v) => f.debug_tuple("List").field(v).finish(),
+ Self::Map(v) => f.debug_tuple("Map").field(v).finish(),
+ Self::Struct(row) => {
+ // Keep the previous ordered-pair representation while names
are shared.
+ struct Fields<'a>(&'a StructValue);
+
+ impl fmt::Debug for Fields<'_> {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ f.debug_list()
+
.entries(self.0.fields().iter().zip(self.0.values()))
+ .finish()
+ }
+ }
+
+ f.debug_tuple("Struct").field(&Fields(row)).finish()
+ }
+ Self::Variant { value, metadata } => f
+ .debug_struct("Variant")
+ .field("value", value)
+ .field("metadata", metadata)
+ .finish(),
+ }
+ }
+}
+
impl Value {
/// Get this value as a bool, or None if it's not a bool.
pub fn as_bool(&self) -> Option<bool> {
@@ -155,7 +207,7 @@ impl fmt::Display for Value {
}
Value::Struct(s) => {
write!(f, "(")?;
- for (i, (k, v)) in s.iter().enumerate() {
+ for (i, (k, v)) in
s.fields().iter().zip(s.values()).enumerate() {
if i > 0 {
write!(f, ", ")?;
}
@@ -175,16 +227,119 @@ impl fmt::Display for Value {
}
}
+/// An ordered struct value with shared immutable field names and owned values.
+/// Cloning shares the names while cloning each value independently.
+#[derive(Debug, Clone, PartialEq)]
+pub struct StructValue {
+ fields: Arc<[String]>,
+ values: Vec<Value>,
+}
+
+impl StructValue {
+ /// Create a struct with immutable names that can be shared with other
structs.
+ /// Panics if the numbers of names and values differ.
+ pub fn new(fields: Arc<[String]>, values: Vec<Value>) -> Self {
+ assert_eq!(
+ fields.len(),
+ values.len(),
+ "field names and values must have the same length"
+ );
+ Self { fields, values }
+ }
+
+ /// Get the number of fields.
+ pub fn len(&self) -> usize {
+ self.fields.len()
+ }
+
+ /// Check whether this struct has no fields.
+ pub fn is_empty(&self) -> bool {
+ self.fields.is_empty()
+ }
+
+ /// Get the ordered field names.
+ pub fn fields(&self) -> &[String] {
+ &self.fields
+ }
+
+ /// Get the ordered values.
+ pub fn values(&self) -> &[Value] {
+ &self.values
+ }
+
+ /// Get a field value by index.
+ pub fn get(&self, index: usize) -> Option<&Value> {
+ self.values.get(index)
+ }
+
+ /// Get the first field value with the given name.
+ pub fn get_by_name(&self, name: &str) -> Option<&Value> {
+ self.fields
+ .iter()
+ .position(|field| field == name)
+ .and_then(|index| self.values.get(index))
+ }
+}
+
+#[derive(Clone)]
+enum RowFields {
+ Owned(Vec<String>),
+ Shared(Arc<[String]>),
+}
+
+impl RowFields {
+ fn as_slice(&self) -> &[String] {
+ match self {
+ Self::Owned(fields) => fields,
+ Self::Shared(fields) => fields,
+ }
+ }
+}
+
+impl fmt::Debug for RowFields {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ fmt::Debug::fmt(self.as_slice(), f)
+ }
+}
+
+impl PartialEq for RowFields {
+ fn eq(&self, other: &Self) -> bool {
+ self.as_slice() == other.as_slice()
+ }
+}
+
/// A Row is an ordered collection of (field_name, value) pairs.
/// Supports access by index or by field name.
#[derive(Debug, Clone, PartialEq)]
pub struct Row {
/// Field names (in order)
- fields: Vec<String>,
+ fields: RowFields,
/// Values corresponding to fields (in order)
values: Vec<Value>,
}
+impl From<Row> for StructValue {
+ fn from(row: Row) -> Self {
+ let fields = match row.fields {
+ RowFields::Owned(fields) => fields.into(),
+ RowFields::Shared(fields) => fields,
+ };
+ Self {
+ fields,
+ values: row.values,
+ }
+ }
+}
+
+impl From<StructValue> for Row {
+ fn from(value: StructValue) -> Self {
+ Self {
+ fields: RowFields::Shared(value.fields),
+ values: value.values,
+ }
+ }
+}
+
impl Row {
/// Create a new Row from field names and values.
/// Panics if lengths don't match.
@@ -194,25 +349,41 @@ impl Row {
values.len(),
"field names and values must have the same length"
);
- Row { fields, values }
+ Row {
+ fields: RowFields::Owned(fields),
+ values,
+ }
+ }
+
+ /// Create a Row with immutable field names shared by rows from the same
batch.
+ pub(crate) fn with_shared_fields(fields: Arc<[String]>, values:
Vec<Value>) -> Self {
+ assert_eq!(
+ fields.len(),
+ values.len(),
+ "field names and values must have the same length"
+ );
+ Row {
+ fields: RowFields::Shared(fields),
+ values,
+ }
}
/// Create an empty Row.
pub fn empty() -> Self {
Row {
- fields: vec![],
+ fields: RowFields::Owned(vec![]),
values: vec![],
}
}
/// Get the number of fields in this Row.
pub fn len(&self) -> usize {
- self.fields.len()
+ self.fields().len()
}
/// Check if this Row is empty.
pub fn is_empty(&self) -> bool {
- self.fields.is_empty()
+ self.fields().is_empty()
}
/// Get a field value by index.
@@ -227,7 +398,7 @@ impl Row {
/// Get a field value by name.
pub fn get_by_name(&self, name: &str) -> Option<&Value> {
- self.fields
+ self.fields()
.iter()
.position(|f| f == name)
.and_then(|i| self.values.get(i))
@@ -235,7 +406,7 @@ impl Row {
/// Get field names.
pub fn fields(&self) -> &[String] {
- &self.fields
+ self.fields.as_slice()
}
/// Get values.
@@ -293,6 +464,234 @@ mod tests {
assert_eq!(row.get_by_name("nonexistent"), None);
}
+ #[test]
+ fn cloned_rows_share_field_names_and_keep_values_independent() {
+ let fields: Arc<[String]> = vec!["items".to_string()].into();
+ let row = Row::with_shared_fields(
+ Arc::clone(&fields),
+ vec![Value::List(vec![Value::String("original".to_string())])],
+ );
+ let mut cloned = row.clone();
+
+ assert!(std::ptr::eq(fields.as_ref(), row.fields()));
+ assert!(std::ptr::eq(row.fields(), cloned.fields()));
+ let Value::List(items) = &mut cloned.values[0] else {
+ panic!("expected a list value");
+ };
+ items[0] = Value::String("changed".to_string());
+ assert_eq!(
+ row.get(0),
+ Some(&Value::List(vec![Value::String("original".to_string())]))
+ );
+ assert_ne!(row.values(), cloned.values());
+ }
+
+ #[test]
+ fn constructed_rows_clone_values_independently() {
+ let row = Row::new(
+ vec!["items".to_string()],
+ vec![Value::List(vec![Value::String("original".to_string())])],
+ );
+
+ let mut cloned = row.clone();
+ assert_eq!(row.fields(), cloned.fields());
+ let Value::List(items) = &mut cloned.values[0] else {
+ panic!("expected a list value");
+ };
+ items[0] = Value::String("changed".to_string());
+ assert_eq!(
+ row.get(0),
+ Some(&Value::List(vec![Value::String("original".to_string())]))
+ );
+ assert_ne!(row.values(), cloned.values());
+ }
+
+ #[test]
+ fn shared_field_names_preserve_order_duplicates_equality_and_debug() {
+ let names = vec![
+ "second".to_string(),
+ "first".to_string(),
+ "second".to_string(),
+ ];
+ let values = vec![Value::Long(2), Value::Long(1), Value::Long(3)];
+ let shared = Row::with_shared_fields(names.clone().into(),
values.clone());
+ let owned = Row::new(names.clone(), values);
+
+ assert_eq!(shared.fields(), names.as_slice());
+ assert_eq!(shared.get_by_name("second"), Some(&Value::Long(2)));
+ assert_eq!(shared.get(2), Some(&Value::Long(3)));
+ assert_eq!(shared, owned);
+ assert_eq!(owned, shared);
+ assert_eq!(format!("{shared:?}"), format!("{owned:?}"));
+ assert_eq!(
+ format!("{shared:?}"),
+ "Row { fields: [\"second\", \"first\", \"second\"], values:
[Long(2), Long(1), Long(3)] }"
+ );
+ let reordered = Row::new(
+ vec![
+ "first".to_string(),
+ "second".to_string(),
+ "second".to_string(),
+ ],
+ owned.values().to_vec(),
+ );
+ assert_ne!(shared, reordered);
+ }
+
+ #[test]
+ fn nested_struct_names_preserve_pair_order_equality_and_formatting() {
+ let names = vec![
+ "".to_string(),
+ "a".to_string(),
+ "a".to_string(),
+ "é".to_string(),
+ ];
+ let values = vec![
+ Value::Null,
+ Value::Long(1),
+ Value::Long(2),
+ Value::List(vec![Value::String("original".to_string())]),
+ ];
+ let shared =
+ Value::Struct(Row::with_shared_fields(names.clone().into(),
values.clone()).into());
+ let owned = Value::Struct(Row::new(names, values).into());
+ assert_eq!(shared, owned);
+ assert_eq!(owned, shared);
+ let Value::Struct(row) = &shared else {
+ panic!("expected a struct value");
+ };
+ assert_eq!(row.get_by_name("a"), Some(&Value::Long(1)));
+ assert_eq!(row.get(2), Some(&Value::Long(2)));
+ assert_eq!(shared.to_string(), "(=NULL, a=1, a=2, é=[original])");
+ assert_eq!(
+ format!("{shared:?}"),
+ "Struct([(\"\", Null), (\"a\", Long(1)), (\"a\", Long(2)), (\"é\",
List([String(\"original\")]))])"
+ );
+ assert_eq!(format!("{shared:#?}"), format!("{owned:#?}"));
+ let empty = Value::Struct(Row::empty().into());
+ assert_eq!(format!("{empty:?}"), "Struct([])");
+ assert_eq!(empty.to_string(), "()");
+ let one = Value::Struct(Row::new(vec!["a".to_string()],
vec![Value::Integer(1)]).into());
+ assert_eq!(
+ format!("{one:#?}"),
+ concat!(
+ "Struct(\n",
+ " [\n",
+ " (\n",
+ " \"a\",\n",
+ " Integer(\n",
+ " 1,\n",
+ " ),\n",
+ " ),\n",
+ " ],\n",
+ ")"
+ )
+ );
+ }
+
+ #[test]
+ fn nested_struct_clones_share_names_and_keep_child_values_independent() {
+ let fields: Arc<[String]> = vec!["items".to_string()].into();
+ let value = Value::Struct(
+ Row::with_shared_fields(
+ Arc::clone(&fields),
+ vec![Value::List(vec![Value::String("original".to_string())])],
+ )
+ .into(),
+ );
+ let mut cloned = value.clone();
+ let Value::Struct(original_row) = &value else {
+ panic!("expected a struct value");
+ };
+ let Value::Struct(cloned_row) = &mut 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]
+ fn nested_struct_equality_does_not_shortcut_nan() {
+ let value = Value::Struct(
+ Row::with_shared_fields(
+ vec!["nan".to_string()].into(),
+ vec![Value::Double(f64::NAN)],
+ )
+ .into(),
+ );
+ assert_ne!(value, value.clone());
+ }
+
+ #[test]
+ fn struct_row_conversions_move_values_and_preserve_name_storage() {
+ let row = Row::new(
+ vec!["field".to_string()],
+ vec![Value::String("value".to_string())],
+ );
+ let names_text = row.fields()[0].as_ptr();
+ let values = row.values().as_ptr();
+ let value = StructValue::from(row);
+ assert_eq!(value.fields()[0].as_ptr(), names_text);
+ assert_eq!(value.values().as_ptr(), values);
+ assert_eq!(value.len(), 1);
+ assert!(!value.is_empty());
+ assert_eq!(value.get(0), Some(&Value::String("value".to_string())));
+ assert_eq!(value.get(1), None);
+ assert_eq!(value.get_by_name("missing"), None);
+
+ let names = value.fields().as_ptr();
+ let row = Row::from(value);
+ assert_eq!(row.fields().as_ptr(), names);
+ assert_eq!(row.values().as_ptr(), values);
+ let value = StructValue::from(row);
+ assert_eq!(value.fields().as_ptr(), names);
+ assert_eq!(value.values().as_ptr(), values);
+ }
+
+ #[test]
+ fn struct_value_constructor_shares_names_and_preserves_empty_values() {
+ let names: Arc<[String]> = vec!["name".to_string()].into();
+ let value = StructValue::new(Arc::clone(&names), vec![Value::Null]);
+ assert!(std::ptr::eq(names.as_ref(), value.fields()));
+ assert_eq!(value.get_by_name("name"), Some(&Value::Null));
+ let empty = StructValue::new(Vec::new().into(), vec![]);
+ assert_eq!(empty.len(), 0);
+ assert!(empty.is_empty());
+ assert!(empty.fields().is_empty());
+ assert!(empty.values().is_empty());
+ assert_eq!(Row::from(empty), Row::empty());
+ }
+
+ #[test]
+ #[should_panic(expected = "field names and values must have the same
length")]
+ fn struct_value_constructor_rejects_mismatched_lengths() {
+ let _ = StructValue::new(vec!["name".to_string()].into(), vec![]);
+ }
+
+ #[test]
+ fn shared_empty_field_names_preserve_empty_row_behavior() {
+ let fields: Arc<[String]> = Vec::new().into();
+ let row = Row::with_shared_fields(fields, vec![]);
+ assert_eq!(row, Row::empty());
+ assert!(row.fields().is_empty());
+ assert!(row.into_values().is_empty());
+ }
+
+ #[test]
+ #[should_panic(expected = "field names and values must have the same
length")]
+ fn shared_field_names_reject_mismatched_lengths() {
+ let _ = Row::with_shared_fields(vec!["field".to_string()].into(),
vec![]);
+ }
+
#[test]
fn test_value_conversions() {
let b = Value::Bool(true);
@@ -393,7 +792,8 @@ mod tests {
m.insert("k".to_string(), Value::Integer(9));
assert_eq!(Value::Map(m).to_string(), "{k: 9}");
assert_eq!(
- Value::Struct(vec![("a".to_string(),
Value::Integer(1))]).to_string(),
+ Value::Struct(Row::new(vec!["a".to_string()],
vec![Value::Integer(1)]).into())
+ .to_string(),
"(a=1)"
);
}
diff --git a/crates/spark-connect/src/session.rs
b/crates/spark-connect/src/session.rs
index eff2adf..30a31a8 100644
--- a/crates/spark-connect/src/session.rs
+++ b/crates/spark-connect/src/session.rs
@@ -781,18 +781,18 @@ fn coerce_value(v: Option<&Value>, target: &DataType) ->
Value {
.collect(),
),
(DataType::Struct { fields }, Value::Struct(sf)) => Value::Struct(
- fields
- .iter()
- .enumerate()
- .map(|(i, f)| {
- let val = sf
- .iter()
- .find(|(n, _)| n == &f.name)
- .map(|(_, v)| v)
- .or_else(|| sf.get(i).map(|(_, v)| v));
- (f.name.clone(), coerce_value(val, &f.data_type))
- })
- .collect(),
+ 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()
@@ -1002,10 +1002,9 @@ fn values_to_arrow(values: &[Value], dt: &DataType) ->
Result<Arc<dyn arrow::arr
for (i, f) in fields.iter().enumerate() {
// Prefer match by field name, else positional.
let val = sf
- .iter()
- .find(|(n, _)| n == &f.name)
- .map(|(_, val)| val.clone())
- .or_else(|| sf.get(i).map(|(_, val)|
val.clone()))
+ .get_by_name(&f.name)
+ .or_else(|| sf.get(i))
+ .cloned()
.unwrap_or(Value::Null);
cols[i].push(val);
}
@@ -1302,6 +1301,238 @@ fn decimal_str_to_unscaled(s: &str, scale: i32) ->
Result<i128> {
mod tests {
use super::*;
+ fn struct_field(name: &str, data_type: DataType) ->
crate::types::StructField {
+ crate::types::StructField {
+ name: name.to_string(),
+ data_type,
+ nullable: true,
+ metadata: Default::default(),
+ }
+ }
+
+ #[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};
+
+ let target = DataType::Struct {
+ fields: vec![
+ struct_field("dup", DataType::Integer),
+ struct_field("positional", DataType::Short),
+ struct_field("dup", DataType::Long),
+ struct_field("missing", DataType::Integer),
+ ],
+ };
+ let values = vec![
+ 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(),
+ ),
+ Value::Map(
+ [
+ ("dup".to_string(), Value::Long(5)),
+ ("other".to_string(), Value::Long(9)),
+ ]
+ .into(),
+ ),
+ Value::Null,
+ Value::Struct(Row::new(vec!["dup".to_string()],
vec![Value::Null]).into()),
+ ];
+ let array = values_to_arrow(&values, &target).unwrap();
+ let array = array.as_any().downcast_ref::<StructArray>().unwrap();
+ assert_eq!(array.len(), 4);
+ assert!(!array.is_null(0));
+ assert!(!array.is_null(1));
+ assert!(array.is_null(2));
+ assert!(!array.is_null(3));
+ let first = array
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ assert_eq!(
+ first,
+ &Int32Array::from(vec![Some(257), Some(5), None, None])
+ );
+ let positional = array
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int16Array>()
+ .unwrap();
+ assert_eq!(
+ positional,
+ &Int16Array::from(vec![Some(999), None, None, None])
+ );
+ let duplicate = array
+ .column(2)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ assert_eq!(
+ duplicate,
+ &Int64Array::from(vec![Some(257), Some(5), None, None])
+ );
+ assert_eq!(array.column(3).null_count(), 4);
+ }
+
+ #[test]
+ fn struct_arrow_encoding_recurses_through_arrays_and_maps() {
+ use arrow::array::{Int16Array, ListArray, MapArray, StructArray};
+
+ let leaf_type = DataType::Struct {
+ fields: vec![struct_field("n", DataType::Short)],
+ };
+ let leaf_value =
+ || Value::Struct(Row::new(vec!["n".to_string()],
vec![Value::Long(300)]).into());
+ let target = DataType::Struct {
+ fields: vec![
+ struct_field(
+ "items",
+ DataType::Array {
+ element_type: Box::new(leaf_type.clone()),
+ contains_null: true,
+ },
+ ),
+ struct_field(
+ "lookup",
+ DataType::Map {
+ key_type: Box::new(DataType::String {
+ collation: "UTF8_BINARY".to_string(),
+ }),
+ value_type: Box::new(leaf_type),
+ value_contains_null: true,
+ },
+ ),
+ ],
+ };
+ let source = Value::Struct(
+ Row::new(
+ vec!["items".to_string(), "lookup".to_string()],
+ vec![
+ Value::List(vec![leaf_value(), Value::Null]),
+ Value::Map([("k".to_string(), leaf_value())].into()),
+ ],
+ )
+ .into(),
+ );
+ let array = values_to_arrow(&[source], &target).unwrap();
+ let array = array.as_any().downcast_ref::<StructArray>().unwrap();
+ let items = array
+ .column(0)
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .unwrap();
+ let items = items.value(0);
+ let items = items.as_any().downcast_ref::<StructArray>().unwrap();
+ assert_eq!(
+ items
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int16Array>()
+ .unwrap(),
+ &Int16Array::from(vec![Some(300), None])
+ );
+ let lookup =
array.column(1).as_any().downcast_ref::<MapArray>().unwrap();
+ let entries = lookup.value(0);
+ let value = entries
+ .column(1)
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .unwrap();
+ assert_eq!(
+ value
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int16Array>()
+ .unwrap(),
+ &Int16Array::from(vec![Some(300)])
+ );
+ }
+
fn session() -> SparkSession {
SparkSession::builder()
.remote("sc://localhost:15002")
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
