viirya commented on code in PR #149:
URL:
https://github.com/apache/spark-connect-rust/pull/149#discussion_r4477708284
##########
crates/pyspark-rs/src/session.rs:
##########
@@ -449,112 +455,236 @@ impl PySparkSession {
// column (matching pyspark createDataFrame(data,
AtomicType())).
match crate::types::py_to_data_type(s)? {
DataType::Struct { fields } => Spec::Struct(fields),
- atomic =>
Spec::Struct(vec![spark_connect::types::StructField {
- name: "value".to_string(),
- data_type: atomic,
- nullable: true,
- metadata: std::collections::BTreeMap::new(),
- }]),
+ atomic => {
+ atomic_schema = Some(atomic.clone());
+
Spec::Struct(vec![spark_connect::types::StructField {
+ name: "value".to_string(),
+ data_type: atomic,
+ nullable: true,
+ metadata: std::collections::BTreeMap::new(),
+ }])
+ }
}
}
}
};
- let spec_names: Option<Vec<String>> = match &spec {
- Spec::Names(n) => Some(n.clone()),
- Spec::Struct(f) => Some(f.iter().map(|x|
x.name.clone()).collect()),
- Spec::None => None,
- };
+ // As pyspark, empty input needs a schema to type its relation, and
gets no toDF.
+ let empty = data.is_empty();
+ if empty && !matches!(spec, Spec::Struct(_)) {
+ return Err(crate::column::raise_pyspark(
+ py,
+ "PySparkValueError",
+ "CANNOT_INFER_EMPTY_SCHEMA",
+ &[],
+ ));
+ }
+ if !empty {
+ if let Some(columns) = &pandas_columns {
+ if let Some(atomic) = &atomic_schema {
+ // pyspark reports str() of the parsed type, which is its
repr.
+ let data_type = crate::types::data_type_to_py(py, atomic)?
+ .bind(py)
+ .repr()?
+ .extract::<String>()?;
+ return Err(crate::column::raise_pyspark(
+ py,
+ "PySparkTypeError",
+ "UNSUPPORTED_DATA_TYPE_FOR_ARROW",
+ &[("data_type", &data_type)],
+ ));
+ }
+ if matches!(spec, Spec::None) && columns.is_empty() {
+ return Err(crate::column::raise_pyspark(
+ py,
+ "PySparkValueError",
+ "CANNOT_INFER_EMPTY_SCHEMA",
+ &[],
+ ));
+ }
+ }
+ }
- // Parse each item into values; a list/tuple carries no names, a Row
carries its own.
- let mut rows: Vec<Row> = vec![];
+ // How a row carries its values, as in pyspark: positionally (lists,
tuples and
+ // scalars), positionally with names of its own (Rows), or by name
(dicts, and other
+ // objects through their `__dict__`, whose keys are sorted as
`_infer_schema` does).
+ enum Item {
Review Comment:
namedtuples are missing here. PySpark names them from `_fields`; this treats
them as plain tuples (`_1`, `_2`). Could they go to `Item::Named`, or into
`known-differences.md`?
##########
crates/pyspark-rs/src/session.rs:
##########
@@ -449,112 +455,236 @@ impl PySparkSession {
// column (matching pyspark createDataFrame(data,
AtomicType())).
match crate::types::py_to_data_type(s)? {
DataType::Struct { fields } => Spec::Struct(fields),
- atomic =>
Spec::Struct(vec![spark_connect::types::StructField {
- name: "value".to_string(),
- data_type: atomic,
- nullable: true,
- metadata: std::collections::BTreeMap::new(),
- }]),
+ atomic => {
+ atomic_schema = Some(atomic.clone());
+
Spec::Struct(vec![spark_connect::types::StructField {
+ name: "value".to_string(),
+ data_type: atomic,
+ nullable: true,
+ metadata: std::collections::BTreeMap::new(),
+ }])
+ }
}
}
}
};
- let spec_names: Option<Vec<String>> = match &spec {
- Spec::Names(n) => Some(n.clone()),
- Spec::Struct(f) => Some(f.iter().map(|x|
x.name.clone()).collect()),
- Spec::None => None,
- };
+ // As pyspark, empty input needs a schema to type its relation, and
gets no toDF.
+ let empty = data.is_empty();
+ if empty && !matches!(spec, Spec::Struct(_)) {
+ return Err(crate::column::raise_pyspark(
+ py,
+ "PySparkValueError",
+ "CANNOT_INFER_EMPTY_SCHEMA",
+ &[],
+ ));
+ }
+ if !empty {
+ if let Some(columns) = &pandas_columns {
+ if let Some(atomic) = &atomic_schema {
+ // pyspark reports str() of the parsed type, which is its
repr.
+ let data_type = crate::types::data_type_to_py(py, atomic)?
+ .bind(py)
+ .repr()?
+ .extract::<String>()?;
+ return Err(crate::column::raise_pyspark(
+ py,
+ "PySparkTypeError",
+ "UNSUPPORTED_DATA_TYPE_FOR_ARROW",
+ &[("data_type", &data_type)],
+ ));
+ }
+ if matches!(spec, Spec::None) && columns.is_empty() {
+ return Err(crate::column::raise_pyspark(
+ py,
+ "PySparkValueError",
+ "CANNOT_INFER_EMPTY_SCHEMA",
+ &[],
+ ));
+ }
+ }
+ }
- // Parse each item into values; a list/tuple carries no names, a Row
carries its own.
- let mut rows: Vec<Row> = vec![];
+ // How a row carries its values, as in pyspark: positionally (lists,
tuples and
+ // scalars), positionally with names of its own (Rows), or by name
(dicts, and other
+ // objects through their `__dict__`, whose keys are sorted as
`_infer_schema` does).
+ enum Item {
+ Plain(Vec<Value>),
+ Named(Vec<String>, Vec<Value>),
+ Record(Vec<(String, Value)>),
+ }
+ let mut items: Vec<Item> = Vec::with_capacity(data.len());
for item in data.iter() {
Review Comment:
nit: this match repeats the Row/list/tuple checks of the `if` chain below. A
`__dict__` branch at the end of that chain, before the scalar one, would avoid
the duplication.
##########
crates/pyspark-rs/src/session.rs:
##########
@@ -449,112 +455,236 @@ impl PySparkSession {
// column (matching pyspark createDataFrame(data,
AtomicType())).
match crate::types::py_to_data_type(s)? {
DataType::Struct { fields } => Spec::Struct(fields),
- atomic =>
Spec::Struct(vec![spark_connect::types::StructField {
- name: "value".to_string(),
- data_type: atomic,
- nullable: true,
- metadata: std::collections::BTreeMap::new(),
- }]),
+ atomic => {
+ atomic_schema = Some(atomic.clone());
+
Spec::Struct(vec![spark_connect::types::StructField {
+ name: "value".to_string(),
+ data_type: atomic,
+ nullable: true,
+ metadata: std::collections::BTreeMap::new(),
+ }])
+ }
}
}
}
};
- let spec_names: Option<Vec<String>> = match &spec {
- Spec::Names(n) => Some(n.clone()),
- Spec::Struct(f) => Some(f.iter().map(|x|
x.name.clone()).collect()),
- Spec::None => None,
- };
+ // As pyspark, empty input needs a schema to type its relation, and
gets no toDF.
+ let empty = data.is_empty();
+ if empty && !matches!(spec, Spec::Struct(_)) {
+ return Err(crate::column::raise_pyspark(
+ py,
+ "PySparkValueError",
+ "CANNOT_INFER_EMPTY_SCHEMA",
+ &[],
+ ));
+ }
+ if !empty {
+ if let Some(columns) = &pandas_columns {
+ if let Some(atomic) = &atomic_schema {
+ // pyspark reports str() of the parsed type, which is its
repr.
+ let data_type = crate::types::data_type_to_py(py, atomic)?
+ .bind(py)
+ .repr()?
+ .extract::<String>()?;
+ return Err(crate::column::raise_pyspark(
+ py,
+ "PySparkTypeError",
+ "UNSUPPORTED_DATA_TYPE_FOR_ARROW",
+ &[("data_type", &data_type)],
+ ));
+ }
+ if matches!(spec, Spec::None) && columns.is_empty() {
+ return Err(crate::column::raise_pyspark(
+ py,
+ "PySparkValueError",
+ "CANNOT_INFER_EMPTY_SCHEMA",
+ &[],
+ ));
+ }
+ }
+ }
- // Parse each item into values; a list/tuple carries no names, a Row
carries its own.
- let mut rows: Vec<Row> = vec![];
+ // How a row carries its values, as in pyspark: positionally (lists,
tuples and
+ // scalars), positionally with names of its own (Rows), or by name
(dicts, and other
+ // objects through their `__dict__`, whose keys are sorted as
`_infer_schema` does).
+ enum Item {
+ Plain(Vec<Value>),
+ Named(Vec<String>, Vec<Value>),
+ Record(Vec<(String, Value)>),
+ }
+ let mut items: Vec<Item> = Vec::with_capacity(data.len());
for item in data.iter() {
- let (values, row_names): (Vec<Value>, Option<Vec<String>>) =
- if let Ok(pyrow) = item.extract::<PyRef<PyRow>>() {
- (
- pyrow.row.values().to_vec(),
- Some(pyrow.row.fields().to_vec()),
- )
- } else if let Ok(lst) = item.cast::<PyList>() {
- (
- lst.iter()
- .map(|v| py_to_value(&v))
- .collect::<PyResult<_>>()?,
- None,
- )
- } else if let Ok(tuple) = item.cast::<pyo3::types::PyTuple>() {
- (
- tuple
- .iter()
- .map(|v| py_to_value(&v))
- .collect::<PyResult<_>>()?,
- None,
- )
- } else if let Ok(dict) = item.cast::<PyDict>() {
- // A dict row: field names are its keys (sorted, mirroring
pyspark's
- // `_infer_schema` for dicts), values in that key order.
- let mut pairs: Vec<(String, Bound<'_, PyAny>)> =
Vec::with_capacity(dict.len());
- for (k, v) in dict.iter() {
- pairs.push((k.str()?.to_string(), v));
- }
- pairs.sort_by(|a, b| a.0.cmp(&b.0));
- let names: Vec<String> = pairs.iter().map(|(k, _)|
k.clone()).collect();
- let values: Vec<Value> = pairs
+ let record = match item.cast::<PyDict>() {
+ Ok(dict) => Some(dict.clone()),
+ Err(_) if item.extract::<PyRef<PyRow>>().is_ok() => None,
+ Err(_) if item.cast::<PyList>().is_ok() => None,
+ Err(_) if item.cast::<pyo3::types::PyTuple>().is_ok() => None,
+ Err(_) => match item.getattr("__dict__") {
+ Ok(attributes) => attributes.cast_into::<PyDict>().ok(),
+ Err(_) => None,
+ },
+ };
+ items.push(if let Some(dict) = record {
+ let mut pairs: Vec<(String, Value)> =
Vec::with_capacity(dict.len());
+ for (k, v) in dict.iter() {
+ pairs.push((k.str()?.to_string(), py_to_value(&v)?));
+ }
+ pairs.sort_by(|a, b| a.0.cmp(&b.0));
+ Item::Record(pairs)
+ } else if let Ok(pyrow) = item.extract::<PyRef<PyRow>>() {
+ Item::Named(pyrow.row.fields().to_vec(),
pyrow.row.values().to_vec())
+ } else if let Ok(lst) = item.cast::<PyList>() {
+ Item::Plain(
+ lst.iter()
+ .map(|v| py_to_value(&v))
+ .collect::<PyResult<_>>()?,
+ )
+ } else if let Ok(tuple) = item.cast::<pyo3::types::PyTuple>() {
+ Item::Plain(
+ tuple
.iter()
- .map(|(_, v)| py_to_value(v))
- .collect::<PyResult<_>>()?;
- (values, Some(names))
- } else {
- // A scalar row (int/str/float/...): a single-field row,
matching
- // `createDataFrame([1, 2, 3], IntegerType())`. The field
name comes from
- // the schema when given, else defaults below.
- (vec![py_to_value(&item)?], None)
- };
- let names = spec_names
- .clone()
- .or(row_names)
- .or_else(|| pandas_columns.clone())
- .unwrap_or_else(|| (0..values.len()).map(|i| format!("col{}",
i)).collect());
+ .map(|v| py_to_value(&v))
+ .collect::<PyResult<_>>()?,
+ )
+ } else {
+ // A scalar row (int/str/float/...): a single-field row,
matching
+ // `createDataFrame([1, 2, 3], IntegerType())`. The field name
comes from
+ // the schema when given, else defaults below.
+ Item::Plain(vec![py_to_value(&item)?])
+ });
+ }
+
+ // Plain rows are as wide as the widest of them: pyspark's schema
inference names
+ // them `_1`, `_2`, ..., and extends a shorter list of names the same
way.
+ let plain_width = items
+ .iter()
+ .filter_map(|item| match item {
+ Item::Plain(values) => Some(values.len()),
+ _ => None,
+ })
+ .max()
+ .unwrap_or(0);
+ // The keys of the dict rows, in the order pyspark merges their
schemas: each row's
+ // sorted keys that no earlier row had.
+ let mut record_names: Vec<String> = Vec::new();
+ for item in &items {
+ if let Item::Record(pairs) = item {
+ for (key, _) in pairs {
+ if !record_names.contains(key) {
+ record_names.push(key.clone());
+ }
+ }
+ }
+ }
+ let extend = |names: &[String], width: usize| -> Vec<String> {
+ let mut names = names.to_vec();
+ names.extend((names.len()..width).map(|i| format!("_{}", i + 1)));
+ names
+ };
+ // The names pyspark applies with toDF: a list of names, and the
columns of a pandas
+ // DataFrame or the field names of its schema.
+ let to_df_names: Option<Vec<String>> = match (&spec, &pandas_columns) {
+ _ if empty => None,
+ (Spec::Names(names), Some(columns)) => Some(extend(names,
columns.len())),
+ (Spec::Names(names), None) => Some(extend(names, plain_width)),
+ (Spec::None, Some(columns)) => Some(columns.clone()),
+ (Spec::Struct(fields), Some(_)) => {
+ Some(fields.iter().map(|f| f.name.clone()).collect())
+ }
+ (_, None) => None,
+ };
+ let default_names = extend(&[], plain_width);
+ // A dict row's value for `name`, as pyspark looks it up: null when
the key is missing.
+ let lookup = |pairs: &[(String, Value)], name: &String| {
Review Comment:
`record_names.contains` and `lookup` are linear scans, so this is
O(rows·keys²) for wide dicts. A `HashSet` for `record_names`, and a binary
search on the sorted `pairs`, would keep it linear.
--
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]
