dbtsai commented on code in PR #120:
URL:
https://github.com/apache/spark-connect-rust/pull/120#discussion_r4419243358
##########
python/tests/test_dropin_offline.py:
##########
@@ -2755,3 +2759,335 @@ def
test_create_dataframe_width_mismatch_is_a_normal_exception():
else:
raise AssertionError("expected a width validation error")
"""], check=True)
+
+
+# ---------------------------------------------------------------------------
schema cache
+# The separate-process fixture exercises the real transport without Spark.
+# Tiny wire encoders are used because the drop-in package does not ship the
+# generated AnalyzePlan/ReleaseSession protobuf classes.
+
+
+def schema_cache_varint(value):
+ result = bytearray()
+ while value > 127:
+ result.append((value & 127) | 128)
+ value >>= 7
+ result.append(value)
+ return bytes(result)
+
+
+def schema_cache_message(field, value):
+ if isinstance(value, str):
+ value = value.encode()
+ return schema_cache_varint(field * 8 + 2) +
schema_cache_varint(len(value)) + value
+
+
+def schema_cache_schema(name="id", nested=False):
+ # types.proto: DataType.long=7, .struct=21; Struct.fields=1;
+ # StructField: name=1, data_type=2, nullable=3, metadata=4.
+ data_type = schema_cache_message(7, b"")
+ if nested:
+ inner_field = schema_cache_field("inner", data_type)
+ data_type = schema_cache_message(21, schema_cache_message(1,
inner_field))
+ return schema_cache_message(21, schema_cache_message(1,
schema_cache_field(name, data_type)))
+
+
+def schema_cache_field(name, data_type):
+ return (
+ schema_cache_message(1, name)
+ + schema_cache_message(2, data_type)
+ + b"\x18\x01"
+ + schema_cache_message(4, '{"note":"original"}')
+ )
+
+
+def schema_cache_serve(connection):
+ import grpc
+
+ # Keep the server independent of the client's interpreter and GIL.
+ class Service:
+ calls = 0
+ data_type = schema_cache_schema()
+ failure = None
+ release_calls = 0
+ block_response = False
+ response_gate = Event()
+
+ def release(self, request, context):
+ self.release_calls += 1
+ return b""
+
+ def analyze(self, request, context):
+ self.calls += 1
+ if self.block_response and not self.response_gate.wait(5):
+ context.abort(
+ grpc.StatusCode.DEADLINE_EXCEEDED,
+ "Python thread could not run while metadata lookup waited",
+ )
+ if self.failure == "rpc":
+ context.abort(grpc.StatusCode.INVALID_ARGUMENT, "schema
unavailable")
+ if self.failure == "no-result":
+ return b""
+ if self.failure == "missing-schema":
+ return schema_cache_message(2, b"")
+ if self.failure == "invalid-type":
+ return schema_cache_message(2, schema_cache_message(1, b""))
+ # base.proto: AnalyzePlanResponse.schema=2, Schema.schema=1.
+ return schema_cache_message(2, schema_cache_message(1,
self.data_type))
+
+ service = Service()
+ with ThreadPoolExecutor(max_workers=4) as executor:
+ server = grpc.server(executor)
+ handlers = {
+ "AnalyzePlan":
grpc.unary_unary_rpc_method_handler(service.analyze),
+ "ReleaseSession":
grpc.unary_unary_rpc_method_handler(service.release),
+ }
+ server.add_generic_rpc_handlers(
+
(grpc.method_handlers_generic_handler("spark.connect.SparkConnectService",
handlers),)
+ )
+ port = server.add_insecure_port("127.0.0.1:0")
+ server.start()
+ connection.send(port)
+ try:
+ while True:
+ command = connection.recv()
+ if command[0] == "stop":
+ break
+ if command[0] == "get":
+ connection.send(getattr(service, command[1]))
+ elif command[0] == "set":
+ setattr(service, command[1], command[2])
+ connection.send(None)
+ elif command[0] == "release":
+ service.response_gate.set()
+ connection.send(None)
+ finally:
+ server.stop(0).wait()
+
+
+class SchemaCacheRemoteService:
+ def __init__(self, connection):
+ object.__setattr__(self, "connection", connection)
+
+ def __getattr__(self, name):
+ self.connection.send(("get", name))
+ assert self.connection.poll(10), "Mock server did not respond"
+ return self.connection.recv()
+
+ def __setattr__(self, name, value):
+ self.connection.send(("set", name, value))
+ assert self.connection.poll(10), "Mock server did not respond"
+ self.connection.recv()
+
+ def release_response(self):
+ self.connection.send(("release",))
+ assert self.connection.poll(10), "Mock server did not respond"
+ self.connection.recv()
+
+
[email protected]
+def schema_cache_analyzed():
+ pytest.importorskip("grpc")
+ context = multiprocessing.get_context("spawn")
+ parent, child = context.Pipe()
+ process = context.Process(target=schema_cache_serve, args=(child,))
+ process.start()
+ spark = None
+ service = SchemaCacheRemoteService(parent)
+ try:
+ assert parent.poll(20), "Mock server did not start"
+ port = parent.recv()
+ spark =
SparkSession.builder.remote(f"sc://127.0.0.1:{port}").getOrCreate()
+ yield spark, service
+ finally:
+ try:
+ if spark is not None:
+ released = service.release_calls
+ spark.stop()
+ assert service.release_calls == released + 1
+ finally:
+ if process.is_alive():
+ parent.send(("stop",))
+ process.join(10)
+ if process.is_alive():
+ process.terminate()
+ process.join(10)
+ parent.close()
+ child.close()
+
+
[email protected]("first", ["schema", "columns", "dtypes"])
+def
test_schema_cache_one_analysis_shared_by_metadata_accessors(schema_cache_analyzed,
first):
+ spark, service = schema_cache_analyzed
+ df = spark.range(10)
+ assert service.calls == 0
+ getattr(df, first)
+ assert service.calls == 1
+ for _ in range(3):
+ schema = df.schema
+ assert isinstance(schema, T.StructType)
+ expected = T.StructType([T.StructField("id", T.LongType(), True,
{"note": "original"})])
+ assert schema.jsonValue() == expected.jsonValue()
+ assert df.columns == ["id"]
+ assert df.dtypes == [("id", "bigint")]
+ assert service.calls == 1
+
+
+def
test_schema_cache_returned_schema_and_lists_are_independent(schema_cache_analyzed):
+ spark, service = schema_cache_analyzed
+ service.data_type = schema_cache_schema(nested=True)
+ df = spark.range(10)
+ original = df.schema
+ expected_json = original.jsonValue()
+ modified = df.schema
+ assert original is not modified
+ # Retain nested getter results so these assertions verify actual Python
mutations.
+ nested = modified.fields[0].dataType
+ nested.add("extra_nested", T.StringType())
+ metadata = nested.fields[0].metadata
+ metadata["note"] = "changed"
+ assert nested.fieldNames() == ["inner", "extra_nested"]
+ assert metadata == {"note": "changed"}
+ modified.add("extra", T.StringType())
+ assert len(modified.fields) == 2
+ assert len(original.fields) == 1
+ assert original.jsonValue() == expected_json
+ columns = df.columns
+ columns.append("extra")
+ dtypes = df.dtypes
+ dtypes.clear()
+ assert columns == ["id", "extra"]
+ assert dtypes == []
+ assert df.schema.jsonValue() == expected_json
+ assert df.columns == ["id"]
+ assert df.dtypes == [("id", "struct<inner:bigint>")]
+ assert service.calls == 1
+
+
+def
test_schema_cache_derived_and_new_plans_resolve_independently(schema_cache_analyzed):
+ spark, service = schema_cache_analyzed
+ df = spark.range(10)
+ assert df.columns == ["id"]
+ derived = df.selectExpr("id AS renamed")
+ service.data_type = schema_cache_schema("renamed")
+ assert derived.columns == ["renamed"]
+ assert service.calls == 2
+ assert df.columns == ["id"]
+ # Rebuilding the same source must observe updated server metadata, while
the
+ # original immutable DataFrame retains its already-resolved schema.
+ fresh = spark.range(10)
+ assert fresh.columns == ["renamed"]
+ other_session = spark.newSession()
+ try:
+ assert other_session.range(10).columns == ["renamed"]
+ finally:
+ other_session.stop()
+ assert service.calls == 4
+ assert df.columns == ["id"]
+
+
[email protected]("failure", ["rpc", "no-result", "missing-schema",
"invalid-type"])
+def test_schema_cache_failures_are_not_cached(schema_cache_analyzed, failure):
+ spark, service = schema_cache_analyzed
+ df = spark.range(10)
+ service.failure = failure
+ for count in (1, 2):
+ message = {
+ "rpc": "schema unavailable",
+ "no-result": "no schema in response",
+ "missing-schema": "Schema is missing",
+ "invalid-type": "DataType kind not set",
+ }[failure]
+ with pytest.raises(Exception, match=message):
+ _ = df.schema
+ assert service.calls == count
+ service.failure = None
+ assert df.columns == ["id"]
+ assert df.dtypes == [("id", "bigint")]
+ assert service.calls == 3
+
+
[email protected]("accessor", ["schema", "columns", "dtypes"])
+def
test_schema_cache_non_struct_responses_are_not_cached(schema_cache_analyzed,
accessor):
Review Comment:
Addressed in
https://github.com/apache/spark-connect-rust/pull/120/commits/ed7efa6c1c427ef9b63ca94c3f8c578d34e1217e.
Reduced non-struct response coverage to one schema accessor case. It still
verifies that repeated invalid responses are not cached and that a later valid
response succeeds.
##########
python/tests/test_dropin_offline.py:
##########
@@ -2755,3 +2759,335 @@ def
test_create_dataframe_width_mismatch_is_a_normal_exception():
else:
raise AssertionError("expected a width validation error")
"""], check=True)
+
+
+# ---------------------------------------------------------------------------
schema cache
+# The separate-process fixture exercises the real transport without Spark.
+# Tiny wire encoders are used because the drop-in package does not ship the
+# generated AnalyzePlan/ReleaseSession protobuf classes.
+
+
+def schema_cache_varint(value):
+ result = bytearray()
+ while value > 127:
+ result.append((value & 127) | 128)
+ value >>= 7
+ result.append(value)
+ return bytes(result)
+
+
+def schema_cache_message(field, value):
+ if isinstance(value, str):
+ value = value.encode()
+ return schema_cache_varint(field * 8 + 2) +
schema_cache_varint(len(value)) + value
+
+
+def schema_cache_schema(name="id", nested=False):
+ # types.proto: DataType.long=7, .struct=21; Struct.fields=1;
+ # StructField: name=1, data_type=2, nullable=3, metadata=4.
+ data_type = schema_cache_message(7, b"")
+ if nested:
+ inner_field = schema_cache_field("inner", data_type)
+ data_type = schema_cache_message(21, schema_cache_message(1,
inner_field))
+ return schema_cache_message(21, schema_cache_message(1,
schema_cache_field(name, data_type)))
+
+
+def schema_cache_field(name, data_type):
+ return (
+ schema_cache_message(1, name)
+ + schema_cache_message(2, data_type)
+ + b"\x18\x01"
+ + schema_cache_message(4, '{"note":"original"}')
+ )
+
+
+def schema_cache_serve(connection):
+ import grpc
+
+ # Keep the server independent of the client's interpreter and GIL.
+ class Service:
+ calls = 0
+ data_type = schema_cache_schema()
+ failure = None
+ release_calls = 0
+ block_response = False
+ response_gate = Event()
+
+ def release(self, request, context):
+ self.release_calls += 1
+ return b""
+
+ def analyze(self, request, context):
+ self.calls += 1
+ if self.block_response and not self.response_gate.wait(5):
+ context.abort(
+ grpc.StatusCode.DEADLINE_EXCEEDED,
+ "Python thread could not run while metadata lookup waited",
+ )
+ if self.failure == "rpc":
+ context.abort(grpc.StatusCode.INVALID_ARGUMENT, "schema
unavailable")
+ if self.failure == "no-result":
+ return b""
+ if self.failure == "missing-schema":
+ return schema_cache_message(2, b"")
+ if self.failure == "invalid-type":
+ return schema_cache_message(2, schema_cache_message(1, b""))
+ # base.proto: AnalyzePlanResponse.schema=2, Schema.schema=1.
+ return schema_cache_message(2, schema_cache_message(1,
self.data_type))
+
+ service = Service()
+ with ThreadPoolExecutor(max_workers=4) as executor:
+ server = grpc.server(executor)
+ handlers = {
+ "AnalyzePlan":
grpc.unary_unary_rpc_method_handler(service.analyze),
+ "ReleaseSession":
grpc.unary_unary_rpc_method_handler(service.release),
+ }
+ server.add_generic_rpc_handlers(
+
(grpc.method_handlers_generic_handler("spark.connect.SparkConnectService",
handlers),)
+ )
+ port = server.add_insecure_port("127.0.0.1:0")
+ server.start()
+ connection.send(port)
+ try:
+ while True:
+ command = connection.recv()
+ if command[0] == "stop":
+ break
+ if command[0] == "get":
+ connection.send(getattr(service, command[1]))
+ elif command[0] == "set":
+ setattr(service, command[1], command[2])
+ connection.send(None)
+ elif command[0] == "release":
+ service.response_gate.set()
+ connection.send(None)
+ finally:
+ server.stop(0).wait()
+
+
+class SchemaCacheRemoteService:
+ def __init__(self, connection):
+ object.__setattr__(self, "connection", connection)
+
+ def __getattr__(self, name):
+ self.connection.send(("get", name))
+ assert self.connection.poll(10), "Mock server did not respond"
+ return self.connection.recv()
+
+ def __setattr__(self, name, value):
+ self.connection.send(("set", name, value))
+ assert self.connection.poll(10), "Mock server did not respond"
+ self.connection.recv()
+
+ def release_response(self):
+ self.connection.send(("release",))
+ assert self.connection.poll(10), "Mock server did not respond"
+ self.connection.recv()
+
+
[email protected]
+def schema_cache_analyzed():
+ pytest.importorskip("grpc")
+ context = multiprocessing.get_context("spawn")
+ parent, child = context.Pipe()
+ process = context.Process(target=schema_cache_serve, args=(child,))
+ process.start()
+ spark = None
+ service = SchemaCacheRemoteService(parent)
+ try:
+ assert parent.poll(20), "Mock server did not start"
+ port = parent.recv()
+ spark =
SparkSession.builder.remote(f"sc://127.0.0.1:{port}").getOrCreate()
+ yield spark, service
+ finally:
+ try:
+ if spark is not None:
+ released = service.release_calls
+ spark.stop()
+ assert service.release_calls == released + 1
Review Comment:
Addressed in
https://github.com/apache/spark-connect-rust/pull/120/commits/ed7efa6c1c427ef9b63ca94c3f8c578d34e1217e.
Removed the ReleaseSession count and teardown assertion. Teardown still
calls spark.stop() and always joins or terminates the mock-server process.
--
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]
