dbtsai commented on code in PR #120:
URL:
https://github.com/apache/spark-connect-rust/pull/120#discussion_r4419242267
##########
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:
Addressed in
https://github.com/apache/spark-connect-rust/pull/120/commits/ed7efa6c1c427ef9b63ca94c3f8c578d34e1217e.
Added this trade-off to schema() rustdoc: concurrent readers of the
DataFrame and its exact clones wait for one analysis RPC, so a stalled analysis
blocks those metadata readers. Propagation itself uses a nonblocking snapshot
and skips an in-flight lookup.
##########
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:
Addressed in
https://github.com/apache/spark-connect-rust/pull/120/commits/ed7efa6c1c427ef9b63ca94c3f8c578d34e1217e.
Clarified that this is defensive recovery for Rust callers: the lock spans a
blocking RPC, and calling the sync API inside a Tokio runtime can panic there.
The slot contains only None or a complete schema. The test also verifies
recovery clears the poisoned flag.
##########
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:
Addressed in
https://github.com/apache/spark-connect-rust/pull/120/commits/ed7efa6c1c427ef9b63ca94c3f8c578d34e1217e.
Removed the duplicate non-struct error branches from Rust columns/dtypes and
the Python dtypes binding. These now destructure the Struct invariant enforced
by schema(); declared output schemas are also filtered to Struct before seeding.
##########
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:
Addressed in
https://github.com/apache/spark-connect-rust/pull/120/commits/ed7efa6c1c427ef9b63ca94c3f8c578d34e1217e.
Removed the limit-empty assertion and implemented propagation here. Warm
limit/filter/sort/alias and the other selected Spark 4.2 transformations
snapshot the resolved schema. Cold children keep independent caches. Tests
cover both cases and ensure construction never waits for an in-flight parent
analysis.
##########
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:
Addressed in
https://github.com/apache/spark-connect-rust/pull/120/commits/ed7efa6c1c427ef9b63ca94c3f8c578d34e1217e.
Removed the Rust deep-clone test. The remaining tests focus on cache
sharing, construction-time snapshots, exact schema equality for set operations,
declared outputs, and panic recovery.
##########
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:
Addressed in
https://github.com/apache/spark-connect-rust/pull/120/commits/ed7efa6c1c427ef9b63ca94c3f8c578d34e1217e.
Trimmed this to a small public behavior check that mutating a returned
schema or metadata list does not affect later reads. Removed the nested
mutation assertions and helper encoding for that case.
--
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]
