viirya commented on code in PR #120:
URL:
https://github.com/apache/spark-connect-rust/pull/120#discussion_r4418660833
##########
crates/spark-connect/src/dataframe.rs:
##########
@@ -769,14 +773,40 @@ impl DataFrame {
}
/// Get the schema of this DataFrame.
+ ///
+ /// The first successful lookup is cached and shared with exact clones.
Review Comment:
Since the lock is held across `analyze_plan`, a slow or stalled analysis
also blocks every other reader of this DataFrame and its clones, including
`GroupedData.input_columns()`. Reference PySpark doesn't synchronize here at
all, so this is new behavior on our side. Could the doc comment mention it?
##########
crates/spark-connect/src/dataframe.rs:
##########
@@ -769,14 +773,40 @@ impl DataFrame {
}
/// Get the schema of this DataFrame.
+ ///
+ /// The first successful lookup is cached and shared with exact clones.
+ /// Changes to external table or view metadata do not invalidate this
cache;
+ /// construct a new DataFrame to resolve the schema again. Each call
returns
+ /// an independent value, and failed lookups can be retried.
pub fn schema(&self) -> Result<DataType> {
+ // Keep the lock during analysis so concurrent accesses issue only one
+ // successful RPC. Errors leave the cache empty and can be retried.
+ let mut cached = self
+ .schema_cache
+ .lock()
+ // The cache contains only None or a fully decoded schema, so a
+ // panic during analysis cannot leave a partially written value.
+ .unwrap_or_else(|poisoned| {
+ self.schema_cache.clear_poison();
Review Comment:
Every other mutex in the crate uses `.lock().unwrap()`. A short note on why
this one needs special handling would help the next reader: it's the only lock
held across a blocking, fallible call. Otherwise it looks like an
inconsistency. Also, I don't think a panic while holding this lock is reachable
from the Python side. The test can only trigger one by calling the sync API
inside a tokio runtime. If keeping it is mainly defensive, that's fine, but the
comment should say so.
##########
crates/spark-connect/src/dataframe.rs:
##########
@@ -769,14 +773,40 @@ impl DataFrame {
}
/// Get the schema of this DataFrame.
+ ///
+ /// The first successful lookup is cached and shared with exact clones.
+ /// Changes to external table or view metadata do not invalidate this
cache;
+ /// construct a new DataFrame to resolve the schema again. Each call
returns
+ /// an independent value, and failed lookups can be retried.
pub fn schema(&self) -> Result<DataType> {
+ // Keep the lock during analysis so concurrent accesses issue only one
+ // successful RPC. Errors leave the cache empty and can be retried.
+ let mut cached = self
+ .schema_cache
+ .lock()
+ // The cache contains only None or a fully decoded schema, so a
+ // panic during analysis cannot leave a partially written value.
+ .unwrap_or_else(|poisoned| {
+ self.schema_cache.clear_poison();
+ poisoned.into_inner()
+ });
+ if let Some(schema) = cached.as_ref() {
+ return Ok(schema.clone());
+ }
let request = self.build_analyze_request()?;
let response = block_on(self.session.client().analyze_plan(request))?;
if let Some(proto::analyze_plan_response::Result::Schema(schema)) =
response.result {
- Ok(DataType::from_proto(&schema.schema.ok_or_else(|| {
- SparkError::connect_msg("Schema is missing")
- })?)?)
+ let schema = DataType::from_proto(
+ &schema
+ .schema
+ .ok_or_else(|| SparkError::connect_msg("Schema is
missing"))?,
+ )?;
+ if !matches!(&schema, DataType::Struct { .. }) {
Review Comment:
With this check, `schema()` only ever returns `DataType::Struct`, so the `_
=>` arms in `columns()` (line 842), `dtypes()` (line 1483), and pyspark-rs
`dtypes` (`crates/pyspark-rs/src/dataframe.rs:1387`) can't be reached anymore.
Could we drop them (e.g. `let DataType::Struct { fields } = self.schema()? else
{ unreachable!() }`) so the struct invariant is enforced in one place?
##########
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):
Review Comment:
Same idea on the Python side: `data_type_to_py` builds new Python objects on
every getter call, so returned values are independent by construction. This
test would only start catching something if we cached Python objects, which
this PR doesn't. Could it be trimmed or removed?
##########
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:
The struct check lives only in `DataFrame::schema()`, and `columns`/`dtypes`
both go through it. So running this once per accessor (3 cases), on top of
`failures_are_not_cached`, mostly repeats the same path. One case would cover
it.
##########
crates/spark-connect/src/dataframe.rs:
##########
@@ -2850,6 +2877,118 @@ mod cache_tests {
use super::*;
use prost::Message;
+ #[test]
+ fn schema_cache_follows_exact_clones_but_not_new_plans() {
+ let spark = SparkSession::builder()
+ .remote("sc://localhost:15002")
+ .get_or_create()
+ .unwrap();
+ let df = spark.range(3).unwrap();
+ let cloned_before_analysis = df.clone();
+ let schema = DataType::Struct { fields: vec![] };
+ *df.schema_cache.lock().unwrap() = Some(schema.clone());
+
+ // These accesses succeed without a server, including a clone created
+ // before the schema was populated.
+ assert_eq!(cloned_before_analysis.schema().unwrap(), schema);
+ assert_eq!(df.clone().schema().unwrap(), schema);
+ assert!(Arc::ptr_eq(
+ &df.schema_cache,
+ &cloned_before_analysis.schema_cache
+ ));
+ for derived in [
+ df.limit(1),
+ df.select(["id"]),
+ df.zip(&df).unwrap(),
+ spark.range(3).unwrap(),
+ ] {
+ assert!(!Arc::ptr_eq(&df.schema_cache, &derived.schema_cache));
+ assert!(derived.schema_cache.lock().unwrap().is_none());
+ }
+ }
+
+ #[test]
+ fn schema_cache_returns_independent_nested_values() {
Review Comment:
`DataType` and its fields own their data (`String`, `HashMap`,
`serde_json::Value`), so `clone()` is always a deep copy. The test checks
`Clone`, not the cache. I think it can be removed.
##########
crates/spark-connect/src/dataframe.rs:
##########
@@ -2850,6 +2877,118 @@ mod cache_tests {
use super::*;
use prost::Message;
+ #[test]
+ fn schema_cache_follows_exact_clones_but_not_new_plans() {
+ let spark = SparkSession::builder()
+ .remote("sc://localhost:15002")
+ .get_or_create()
+ .unwrap();
+ let df = spark.range(3).unwrap();
+ let cloned_before_analysis = df.clone();
+ let schema = DataType::Struct { fields: vec![] };
+ *df.schema_cache.lock().unwrap() = Some(schema.clone());
+
+ // These accesses succeed without a server, including a clone created
+ // before the schema was populated.
+ assert_eq!(cloned_before_analysis.schema().unwrap(), schema);
+ assert_eq!(df.clone().schema().unwrap(), schema);
+ assert!(Arc::ptr_eq(
+ &df.schema_cache,
+ &cloned_before_analysis.schema_cache
+ ));
+ for derived in [
+ df.limit(1),
Review Comment:
Asserting that `limit`/`select` start with an empty cache ties the test to
the current choice, not to a correctness requirement. Reference PySpark does
carry the cached schema through `limit` (and `filter`, `sort`, `alias`, ...).
For `select` and `zip` the assertion makes sense, since the schema can change.
For `limit`, I'd drop it so propagating schemas like the reference doesn't need
to start by deleting a test.
##########
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:
This asserts ReleaseSession behavior from a fixture teardown, which isn't
part of the schema-cache change. When it fails, the failure gets reported
against whichever schema-cache test was running. Could we drop it, or move it
into a test of its own?
--
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]
