dbtsai commented on code in PR #120:
URL:
https://github.com/apache/spark-connect-rust/pull/120#discussion_r4417364603
##########
python/tests/test_dropin_offline.py:
##########
@@ -2755,3 +2759,274 @@ 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 as schema_cache_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(
+ schema_cache_grpc.StatusCode.DEADLINE_EXCEEDED,
+ "Python thread could not run while metadata lookup waited",
+ )
+ if self.failure == "rpc":
+ context.abort(schema_cache_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 = schema_cache_grpc.server(executor)
+ handlers = {
+ "AnalyzePlan":
schema_cache_grpc.unary_unary_rpc_method_handler(service.analyze),
+ "ReleaseSession":
schema_cache_grpc.unary_unary_rpc_method_handler(service.release),
+ }
+
server.add_generic_rpc_handlers((schema_cache_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 schema_cache_RemoteService:
+ 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 = schema_cache_RemoteService(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
+ modified = df.schema
+ assert original is not modified
+ modified.fields[0].metadata["note"] = "changed"
+ nested = modified.fields[0].dataType
+ nested.add("extra_nested", T.StringType())
+ assert len(nested.fields) == 2
+ modified.add("extra", T.StringType())
+ assert len(modified.fields) == 2
+ assert len(original.fields) == 1
+ df.columns.append("extra")
+ df.dtypes.clear()
+ assert df.schema.jsonValue() == original.jsonValue()
+ 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
+
+
+def
test_schema_cache_concurrent_schema_access_issues_one_analysis(schema_cache_analyzed):
+ spark, service = schema_cache_analyzed
+ df = spark.range(10)
+ with ThreadPoolExecutor(max_workers=8) as pool:
+ results = list(pool.map(lambda _: df.schema, range(8)))
Review Comment:
Agreed—the immediate response allowed the test to pass without exercising
contention. I’ve updated it locally to start nine mixed
`schema`/`columns`/`dtypes` readers behind a barrier, hold the first response
open, and release the gate in `finally`. It checks one analysis RPC and equal,
distinct schema objects, with bounded waits.
I also tested its sensitivity against a build that deliberately drops the
cache lock before analysis: the revised test failed in all three runs,
observing nine RPCs each time. The final passing rerun is still queued behind
the shared CPU lock; these changes are not pushed yet.
--
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]
