HyukjinKwon commented on code in PR #120:
URL:
https://github.com/apache/spark-connect-rust/pull/120#discussion_r4415685655
##########
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:
**Nit (P3):** The mock's AnalyzePlan handler replies immediately here
because `service.block_response` is never set. Whether threads 2..8 reach
`DataFrame::schema()` while the first RPC is still in flight is therefore left
to the scheduler. Suppose a future change dropped the guard before
`block_on(analyze_plan)` and re-checked the cache afterwards. This test could
still see `service.calls == 1` and pass, so its signal for the single-flight
guarantee, the main reason for holding the mutex across the RPC, is only
probabilistic. The fixture already has the gate that would make it
deterministic, as used by `test_schema_cache_metadata_wait_releases_gil`. Set
`block_response = True`, start the 8 reads, wait until `service.calls` is
non-zero and give the other readers a moment to block on the cache, then call
`release_response()` and assert one call.
**Verification:**
- **Regression:** With the first analysis response held until several
readers have entered the schema accessor, concurrent schema reads on one
DataFrame produce exactly one analysis request on the server and equal,
distinct schema objects.
- **Behavior:** The gated test stays bounded: if the response is never
released, or a reader keeps the GIL, the test fails within its timeouts and
does not hang.
--
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]
