HyukjinKwon commented on code in PR #149:
URL:
https://github.com/apache/spark-connect-rust/pull/149#discussion_r4476724101
##########
crates/pyspark-rs/src/session.rs:
##########
@@ -449,89 +455,154 @@ 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 = true;
+
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 let (false, Some(columns)) = (empty, &pandas_columns) {
+ if atomic_schema {
+ let data_type = match schema {
+ Some(s) => s.repr()?.extract::<String>()?,
Review Comment:
Done in b4e41cb. The error now reports the parsed type through
`data_type_to_py`, so `createDataFrame(pdf, "int")` gives `IntegerType()` as
PySpark does. It takes the type's repr, which is what PySpark's `str()` of a
data type returns. There is a test case for the atomic DDL schema.
##########
crates/pyspark-rs/src/session.rs:
##########
@@ -449,89 +455,154 @@ 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 = true;
+
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 let (false, Some(columns)) = (empty, &pandas_columns) {
+ if atomic_schema {
+ let data_type = match schema {
+ Some(s) => s.repr()?.extract::<String>()?,
+ None => String::new(),
+ };
+ 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![];
+ let mut items: Vec<(Vec<Value>, Option<Vec<String>>)> =
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 parsed = 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<_>>()?;
- (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<_>>()?,
+ 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
+ .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)
+ };
+ items.push(parsed);
+ }
+
+ // Plain rows (no names of their own) 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(|(_, names)| names.is_none())
+ .map(|(values, _)| values.len())
+ .max()
+ .unwrap_or(0);
+ 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);
+
+ let mut rows: Vec<Row> = Vec::with_capacity(items.len());
+ for (values, row_names) in items {
+ let names = match (&spec, &to_df_names) {
+ (Spec::Struct(fields), _) => fields.iter().map(|f|
f.name.clone()).collect(),
+ (_, Some(names)) => names.clone(),
+ _ => row_names.unwrap_or_else(|| default_names.clone()),
+ };
if names.len() != values.len() {
- return Err(row_width_mismatch(rows.len(), values.len(),
names.len())).to_pyerr();
+ // As pyspark, which checks every row against the width of the
schema.
Review Comment:
Fixed here in b4e41cb, since the loop was changing anyway. Dict rows, and
other objects through their `__dict__` as in PySpark's `_infer_schema`, are now
read by name. A struct schema takes their values by its field names, so a
missing key is null and other keys are ignored. Without a schema, the fields
are the rows' sorted keys, in the order first seen. Missing keys also needed
inference past nulls, so each inferred field now takes the type of its first
non-null value, and a field that is null in every row raises
`CANNOT_DETERMINE_TYPE`. New live tests cover nine dict and object cases and
the null cases, and they match official PySpark 4.3.0. PySpark's full type
merging across rows goes further: strings absorb other atomic types, nested
nulls must resolve, merge conflicts get their own errors, and `Decimal` becomes
`decimal(38,18)`. I filed SPARK-60908 for that and listed it in
`dev/parity/known-differences.md`. The width comment and the description now
say plain rows and Rows.
##########
crates/pyspark-rs/src/session.rs:
##########
@@ -449,89 +455,154 @@ 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 = true;
+
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 let (false, Some(columns)) = (empty, &pandas_columns) {
Review Comment:
Done in b4e41cb.
##########
crates/pyspark-rs/src/session.rs:
##########
@@ -554,7 +625,14 @@ impl PySparkSession {
let df = py
.detach(|| self.session.create_dataframe(rows, schema_dtype))
.to_pyerr()?;
- Ok(PyDataFrame::new(df))
+ // As pyspark, column names are applied with toDF, which the server
resolves into a
Review Comment:
Done in b4e41cb. The comment now says the local relation already carries the
names, so `toDF` only changes the plan and keeps `isLocal()` as in PySpark; it
is not redundant.
##########
python/tests/test_dropin_offline.py:
##########
@@ -2903,13 +2903,87 @@ def
test_create_dataframe_width_mismatch_is_a_normal_exception():
try:
spark.createDataFrame([(1, 2), (3,)], schema=schema)
except PySparkValueError as error:
+ # The same error as pyspark, for every kind of schema.
assert type(error) is PySparkValueError, type(error)
- assert str(error) == "Row 1 has 1 values, but the schema has 2
fields", error
+ assert error.getCondition() == "AXIS_LENGTH_MISMATCH", error
+ assert error.getMessageParameters() == {
+ "expected_length": "2", "actual_length": "1"}, error
else:
raise AssertionError("expected a width validation error")
"""], check=True)
+def _local_create_dataframe(*args):
+ # The local relation and its toDF are built without a server.
+ return _offline_session().createDataFrame(*args)
+
+
[email protected]("data,schema,expected", [
+ # Column names are applied with toDF, as in pyspark, so the result is not
local.
+ ([(1, 2)], ["a", "b"], False),
+ ([(1, 2)], ("a", "b"), False),
+ ([(1, 2)], ["a"], False),
+ ([(1, 2)], [], False),
+ ([T.Row(x=1, y=2)], ["a", "b"], False),
+ ([{"x": 1, "y": 2}], ["a", "b"], False),
+ # Without names to apply, the local relation itself is local.
+ ([(1, 2)], None, True),
+ ([T.Row(x=1, y=2)], None, True),
+ ([(1, 2)], "a long, b long", True),
+ ([1, 2], T.LongType(), True),
+ ([], "a long", True),
+])
+def test_create_dataframe_rows_is_local_as_pyspark(data, schema, expected):
+ assert _local_create_dataframe(data, schema).isLocal() is expected
+
+
[email protected]("schema", [None, ["c", "d"], ["c"], "a long, b
string",
+ T.StructType([T.StructField("a",
T.LongType()),
+ T.StructField("b",
T.StringType())])])
+def test_create_dataframe_pandas_is_not_local_as_pyspark(schema):
+ pd = pytest.importorskip("pandas")
+ # pyspark applies the columns, or the field names of the schema, with toDF.
+ df = _local_create_dataframe(pd.DataFrame({"a": [1, 2], "b": ["x", "y"]}),
schema)
+ assert not df.isLocal()
+ # Empty input returns the relation of its schema, without toDF.
Review Comment:
Done in b4e41cb: it is
`test_create_dataframe_empty_pandas_is_local_as_pyspark` now.
--
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]
