viirya commented on code in PR #151:
URL:
https://github.com/apache/spark-connect-rust/pull/151#discussion_r4477802288
##########
python/tests/test_dropin_offline.py:
##########
@@ -3786,3 +3786,28 @@ def
test_connect_registry_sends_auto_cdc_define_flow(mock_connect):
assert 4 not in plain
assert 10 not in details and 11 not in details and 12 not in details
+
+
+# ---------------------------------------------------------------------------
reader schema
Review Comment:
nit: three blank lines before this header; the rest of the file uses two.
##########
python/tests/test_dropin_offline.py:
##########
@@ -3786,3 +3786,28 @@ def
test_connect_registry_sends_auto_cdc_define_flow(mock_connect):
assert 4 not in plain
assert 10 not in details and 11 not in details and 12 not in details
+
+
+# ---------------------------------------------------------------------------
reader schema
+
[email protected]("read", [
+ lambda spark, schema: spark.read.schema(schema),
+ lambda spark, schema: spark.read.csv("unused.csv", schema=schema),
+ lambda spark, schema: spark.read.load("unused", format="json",
schema=schema),
+ lambda spark, schema: spark.readStream.schema(schema),
+ lambda spark, schema: spark.readStream.csv("unused", schema=schema),
+ lambda spark, schema: spark.readStream.load("unused", format="json",
schema=schema),
+], ids=["read.schema", "read.csv", "read.load", "readStream.schema",
"readStream.csv",
+ "readStream.load"])
+def test_reader_schema_takes_a_struct_type_or_a_string_as_pyspark(read):
+ from pyspark.errors import PySparkTypeError
+
+ spark = _offline_session()
+ # pyspark sends a StructType as its JSON and a string as is, for the
server to parse.
+ read(spark, T.StructType([T.StructField("a", T.LongType(), True,
{"comment": "key"})]))
Review Comment:
These rows only check that the call doesn't raise. The StructType case for
`read.csv` / `read.load` / `readStream.csv` / `readStream.load` already passed
before this change, because it was sent as a simple string. Asserting the sent
`DataSource.schema` through `mock_connect`'s recorded ExecutePlan request would
pin the JSON-vs-string behavior offline. Optional, since the e2e test covers it.
##########
crates/pyspark-rs/src/readwriter.rs:
##########
@@ -59,11 +59,35 @@ fn set_opt(
}
}
-/// Resolve a reader `schema` arg (a StructType/DataType or a DDL string) to
the
-/// reader's DDL schema.
+/// A reader `schema` as pyspark's `DataFrameReader.schema` and
`DataStreamReader.schema`
+/// send it: a StructType as its JSON, which keeps field metadata, and a
string as is, for
+/// the server to parse. Anything else raises `NOT_EXPECTED_TYPE`.
+pub(crate) fn reader_schema(schema: &Bound<'_, PyAny>) -> PyResult<String> {
+ if let Ok(schema) = schema.extract::<String>() {
+ return Ok(schema);
+ }
+ if schema
+ .extract::<PyRef<crate::types::PyStructType>>()
+ .is_ok()
+ {
+ return Ok(crate::types::py_to_data_type(schema)?.json());
Review Comment:
`DataType::json_value()` writes `geometry(OGC:CRS84)` /
`geography(OGC:CRS84, SPHERICAL)` for every SRID other than -1. PySpark writes
the actual CRS, e.g. `geometry(EPSG:3857)`. So `GeometryType(3857)` in a reader
schema now reaches the server as SRID 4326, with no error. Could we map the
SRID to its CRS in `json_value` here, or track it in a follow-up and mention it
in the description?
Separately, after the `PyRef<PyStructType>` extract succeeds,
`py_to_data_type` re-walks the whole type ladder. Using the extracted ref
directly is simpler:
```rust
if let Ok(st) = schema.extract::<PyRef<crate::types::PyStructType>>() {
return Ok(DataType::Struct { fields: st.fields.clone() }.json());
}
```
--
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]
