viirya commented on code in PR #127:
URL:
https://github.com/apache/spark-connect-rust/pull/127#discussion_r4431031769
##########
python/tests/test_dropin_offline.py:
##########
@@ -3085,3 +3085,162 @@ def allow_response_after_request():
else:
assert result == [("id", "bigint")]
assert service.calls == 1
+
+
+# Exact built-in string conversion and full outbound Arrow IPC guards.
+
+
[email protected]("text", [
+ "", "1.25", "-42", "NaN", "a\x00b", "é雪😀", "e\u0301", "\uffff",
+])
+def test_exact_string_roundtrip(text):
+ actual = Row(value=text).value
+ assert type(actual) is str
+ assert actual == text
+
+
[email protected]("text", ["\ud800", "\udfff", "a\ud800b"])
+def test_exact_string_surrogate_error_preserves_generic_fallback(text):
+ with pytest.raises(TypeError):
+ Row(value=text)
+
+
+def test_exact_string_subclass_protocol_order():
+ # Pin existing behavior with PyO3 0.28, not a desired contract for str
+ # subclasses. Revisit conversion order when upgrading PyO3.
+ trace = []
+
+ class IndexString(str):
+ def __index__(self):
+ trace.append("index")
+ return 17
+
+ def __float__(self):
+ raise AssertionError("index must take precedence")
+
+ class FloatString(str):
+ def __float__(self):
+ trace.append("float")
+ return 2.5
+
+ def __str__(self):
+ raise AssertionError("numeric conversion must take precedence")
+
+ class FallbackString(str):
+ def __index__(self):
+ trace.append("bad index")
+ raise ValueError("index rejected")
+
+ def __float__(self):
+ trace.append("bad float")
+ raise ValueError("float rejected")
+
+ def __str__(self):
+ raise AssertionError("string extraction must retain the underlying
text")
+
+ class PlainString(str):
+ def __str__(self):
+ raise AssertionError("string extraction must not call __str__")
+
+ cases = [
+ (IndexString("x"), 17),
+ (FloatString("x"), 2.5),
+ (FallbackString("original"), "original"),
+ (PlainString("plain"), "plain"),
+ ]
+ for value, expected in cases:
+ actual = Row(value=value).value
+ assert type(actual) is type(expected)
+ assert actual == expected
+ assert trace == ["index", "float", "bad index", "bad float"]
+
+
+def test_exact_string_nested_strings_and_nulls():
+ value = Row(
+ items=["雪", None, ""],
+ mapping={"x": "😀", "null": None},
+ )
+ actual = Row(value=value).value
+ assert list(actual.items) == ["雪", None, ""]
+ assert actual.mapping == {"x": "😀", "null": None}
+
+
[email protected]("source", ["tuples", "pandas"])
[email protected]("schema_mode", ["inferred", "names", "explicit"])
+def test_exact_string_input_ipc(mock_connect, source, schema_mode):
Review Comment:
These tests (also `test_exact_string_nested_input_ipc` and
`test_exact_string_batch_input_ipc`) check the Value → Arrow encoding, which
this PR doesn't change. Core already covers it, including `é\0…`, `None`, and
`""` (`crates/spark-connect/src/session.rs:1378`). The 256-row case doesn't
match any batching boundary in the builder. Please drop these three tests.
##########
python/tests/_mock_connect_server.py:
##########
@@ -39,6 +39,71 @@ def message_field(field, value):
return varint(field * 8 + 2) + varint(len(value)) + value
+# Field numbers from crates/spark-connect-proto/proto/spark/connect/:
+_ANALYZE_SCHEMA = 4 # base.proto: AnalyzePlanRequest.schema
+_SCHEMA_PLAN = 1 # base.proto: AnalyzePlanRequest.Schema.plan
+_PLAN_ROOT = 1 # base.proto: Plan.root
+_RELATION_LOCAL_RELATION = 11 # relations.proto: Relation.local_relation
+_LOCAL_RELATION_DATA = 1 # relations.proto: LocalRelation.data
+
+
+def _message_fields(data):
+ """Read length-delimited protobuf fields, skipping scalar fields."""
+ def read_varint(position):
+ value = 0
+ for shift in range(0, 64, 7):
+ if position >= len(data):
+ raise ValueError("truncated protobuf varint")
+ byte = data[position]
+ position += 1
+ if shift == 63 and byte > 1:
+ raise ValueError("protobuf varint exceeds 64 bits")
+ value |= (byte & 127) << shift
+ if byte < 128:
+ return value, position
+ raise ValueError("protobuf varint exceeds 64 bits")
+
+ fields = {}
+ position = 0
+ while position < len(data):
+ tag, position = read_varint(position)
+ field, wire_type = tag >> 3, tag & 7
+ if not field:
+ raise ValueError("protobuf field number must be positive")
+ if wire_type == 0:
+ value, position = read_varint(position)
+ else:
+ if wire_type == 1:
+ size = 8
+ elif wire_type == 2:
+ size, position = read_varint(position)
+ elif wire_type == 5:
+ size = 4
+ else:
+ raise ValueError(f"unsupported protobuf wire type {wire_type}")
+ end = position + size
+ if end > len(data):
+ raise ValueError("truncated protobuf field")
+ value, position = data[position:end], end
+ if wire_type == 2:
+ fields[field] = value
+ return fields
+
+
+def _local_relation_data(request):
+ try:
+ for field in (
+ _ANALYZE_SCHEMA, _SCHEMA_PLAN, _PLAN_ROOT,
+ _RELATION_LOCAL_RELATION, _LOCAL_RELATION_DATA,
+ ):
+ request = _message_fields(request).get(field)
+ if request is None:
+ return None
+ except (ValueError, TypeError):
Review Comment:
If this stays: a parse failure returns `None`, and the test then fails
inside `pa.ipc.open_stream(None)`, which hides the cause. Let the error raise.
##########
python/tests/test_dropin_offline.py:
##########
@@ -3085,3 +3085,162 @@ def allow_response_after_request():
else:
assert result == [("id", "bigint")]
assert service.calls == 1
+
+
+# Exact built-in string conversion and full outbound Arrow IPC guards.
+
+
[email protected]("text", [
+ "", "1.25", "-42", "NaN", "a\x00b", "é雪😀", "e\u0301", "\uffff",
+])
+def test_exact_string_roundtrip(text):
+ actual = Row(value=text).value
+ assert type(actual) is str
+ assert actual == text
+
+
[email protected]("text", ["\ud800", "\udfff", "a\ud800b"])
+def test_exact_string_surrogate_error_preserves_generic_fallback(text):
+ with pytest.raises(TypeError):
+ Row(value=text)
+
+
+def test_exact_string_subclass_protocol_order():
Review Comment:
By its own comment, this test locks in behavior that isn't a contract, so a
PyO3 upgrade or a fix to str-subclass handling would mean rewriting it. One
`__index__` subclass is enough to guard that the fast path stays `cast_exact`.
Separately, reference PySpark sends a `str` subclass to Arrow as text when
the schema is StringType. Should a follow-up switch this to `cast` and treat
every `str` instance as a string? That would also cover `numpy.str_`.
##########
python/tests/_mock_connect_server.py:
##########
@@ -39,6 +39,71 @@ def message_field(field, value):
return varint(field * 8 + 2) + varint(len(value)) + value
+# Field numbers from crates/spark-connect-proto/proto/spark/connect/:
+_ANALYZE_SCHEMA = 4 # base.proto: AnalyzePlanRequest.schema
+_SCHEMA_PLAN = 1 # base.proto: AnalyzePlanRequest.Schema.plan
+_PLAN_ROOT = 1 # base.proto: Plan.root
+_RELATION_LOCAL_RELATION = 11 # relations.proto: Relation.local_relation
+_LOCAL_RELATION_DATA = 1 # relations.proto: LocalRelation.data
+
+
+def _message_fields(data):
Review Comment:
This protobuf reader is only needed for the IPC tests below, which cover
unchanged code. Please drop it along with them.
##########
crates/pyspark-rs/src/session.rs:
##########
@@ -955,6 +955,14 @@ pub(crate) fn py_to_value(obj: &Bound<'_, PyAny>) ->
PyResult<Value> {
return Ok(Value::Null);
}
+ // Exact built-in strings cannot implement numeric conversion protocols.
+ // Keep subclasses and strings that fail UTF-8 extraction on the generic
path.
+ if let Ok(s) = obj.cast_exact::<pyo3::types::PyString>() {
+ if let Ok(s) = s.extract::<String>() {
Review Comment:
Nit: when UTF-8 extraction fails on an exact `str`, the generic path can't
succeed. It ends with "unsupported value type for createDataFrame: str", which
is misleading. This isn't new and doesn't block the PR, but returning the
extraction error here would be simpler and clearer.
--
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]
