viirya commented on code in PR #149:
URL:
https://github.com/apache/spark-connect-rust/pull/149#discussion_r4476256657
##########
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:
With a DDL string that parses to an atomic type, e.g. `createDataFrame(pdf,
"int")`, `s.repr()` is the repr of the Python string, so `data_type` comes out
as `'int'`. PySpark parses the DDL first and reports `IntegerType()`. Could
this use `str()` of the parsed type, via `data_type_to_py`? PySpark uses
`str(schema)`. A test case with a DDL atomic schema would cover it.
##########
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:
nit: `if !empty { if let Some(columns) = &pandas_columns { ... } }` reads
more directly than matching on `(false, Some(..))`.
##########
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:
This isn't quite what PySpark does for dict rows. PySpark's `to_row` looks
dicts up by name (`item.get(col)`): missing keys become None and extra keys are
ignored. Here the sorted values are matched by position. So a dict with a
missing key raises `AXIS_LENGTH_MISMATCH` instead of filling None, and
`createDataFrame([{"b": 1, "a": "x"}], "b long, a string")` puts `"x"` into
`b`. That predates this PR, but the comment and the description ("for every
kind of schema") now say this matches PySpark. Could we narrow both to plain
rows and Rows, and track dict-by-name matching separately (or fix it here,
since this loop is already being changed)?
##########
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:
This assertion doesn't depend on `schema`, so it runs once per
parametrization. Could it move to
`test_create_dataframe_pandas_errors_as_pyspark` or a separate test?
##########
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:
nit: the local relation already carries these names, so this `toDF` only
changes the plan shape and `isLocal()`. Saying that explicitly would stop a
later cleanup from removing it as redundant.
--
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]
