From f87973b1e02fe500e3c4266a8ad3f353244994af Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Wed, 28 Oct 2026 21:44:42 -0700 Subject: [PATCH 1/9] [SPARK-60677][CONNECT] Cache DataFrame schemas per immutable plan --- crates/spark-connect/src/dataframe.rs | 60 ++++++- python/tests/test_schema_cache.py | 225 ++++++++++++++++++++++++++ 2 files changed, 277 insertions(+), 8 deletions(-) create mode 100644 python/tests/test_schema_cache.py diff --git a/crates/spark-connect/src/dataframe.rs b/crates/spark-connect/src/dataframe.rs index ea0e2c62..3dc034c0 100644 --- a/crates/spark-connect/src/dataframe.rs +++ b/crates/spark-connect/src/dataframe.rs @@ -3,7 +3,7 @@ //! Provides transformations and actions for working with distributed data. use std::collections::HashMap; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use spark_connect_core::client::ReattachableResponseStream; use spark_connect_core::error::{Result, SparkError}; @@ -26,6 +26,9 @@ use crate::udf::CommonInlineUserDefinedFunctionExpression; pub struct DataFrame { pub(crate) session: SparkSession, pub(crate) plan: Arc, + // Exact clones share this cache. Every new plan starts with an empty cache, + // even when its input has already been analyzed. + schema_cache: Arc>>, } /// Iterator over rows from a DataFrame, yielded lazily as the server streams results. @@ -206,6 +209,7 @@ impl DataFrame { DataFrame { session, plan: Arc::new(plan), + schema_cache: Arc::new(Mutex::new(None)), } } @@ -770,13 +774,26 @@ impl DataFrame { /// Get the schema of this DataFrame. pub fn schema(&self) -> Result { + // 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() + .map_err(|_| SparkError::connect_msg("Schema cache lock poisoned"))?; + 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"))?, + )?; + *cached = Some(schema.clone()); + Ok(schema) } else { Err(SparkError::connect_msg( "Schema analyze failed: no schema in response", @@ -1746,10 +1763,7 @@ impl DataFrame { left: Arc::clone(&self.plan), right: Arc::clone(&other.plan), }; - Ok(DataFrame { - plan: Arc::new(plan), - session: self.session.clone(), - }) + Ok(DataFrame::new(self.session.clone(), plan)) } /// Register this DataFrame as a temporary table (deprecated - use createTempView). @@ -2850,6 +2864,36 @@ 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 cache_default_is_memory_and_disk_deser() { let sl = DataFrame::memory_and_disk_deser(); diff --git a/python/tests/test_schema_cache.py b/python/tests/test_schema_cache.py new file mode 100644 index 00000000..e085e8f0 --- /dev/null +++ b/python/tests/test_schema_cache.py @@ -0,0 +1,225 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Schema caching through the real client transport, without a Spark server. + +The fixture serves only AnalyzePlan. Tiny wire encoders keep the tests independent +of the reference PySpark client and its generated protobuf modules. +""" +from concurrent.futures import ThreadPoolExecutor +import multiprocessing + +import pytest + +grpc = pytest.importorskip("grpc") + +from pyspark.sql import SparkSession, types as T + + +def _varint(value): + result = bytearray() + while value > 127: + result.append((value & 127) | 128) + value >>= 7 + result.append(value) + return bytes(result) + + +def _message(field, value): + if isinstance(value, str): + value = value.encode() + return _varint(field * 8 + 2) + _varint(len(value)) + value + + +def _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 = _message(21, _message(1, _field("inner", _message(7, b"")))) if nested else _message(7, b"") + return _message(21, _message(1, _field(name, data_type))) + + +def _field(name, data_type): + return (_message(1, name) + _message(2, data_type) + b"\x18\x01" + + _message(4, '{"note":"original"}')) + + +def _serve(connection): + # A separate process also exercises columns/dtypes, whose first access may + # hold the client's GIL while waiting for the RPC. + class Service: + calls = 0 + data_type = _schema() + failure = None + + def analyze(self, request, context): + self.calls += 1 + 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 _message(2, b"") + if self.failure == "invalid-type": + return _message(2, _message(1, b"")) + # base.proto: AnalyzePlanResponse.schema=2, Schema.schema=1. + return _message(2, _message(1, self.data_type)) + + service = Service() + with ThreadPoolExecutor(max_workers=4) as executor: + server = grpc.server(executor) + handler = grpc.unary_unary_rpc_method_handler(service.analyze) + server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler( + "spark.connect.SparkConnectService", {"AnalyzePlan": handler}),)) + 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) + finally: + server.stop(0).wait() + + +class _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() + + +@pytest.fixture +def analyzed(): + context = multiprocessing.get_context("spawn") + parent, child = context.Pipe() + process = context.Process(target=_serve, args=(child,)) + process.start() + 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, _RemoteService(parent) + finally: + if process.is_alive(): + parent.send(("stop",)) + process.join(10) + if process.is_alive(): + process.terminate() + process.join(10) + parent.close() + child.close() + + +@pytest.mark.parametrize("first", ["schema", "columns", "dtypes"]) +def test_one_analysis_shared_by_metadata_accessors(analyzed, first): + spark, service = 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_returned_schema_and_lists_are_independent(analyzed): + spark, service = analyzed + service.data_type = _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")] + assert service.calls == 1 + + +def test_derived_and_new_plans_resolve_independently(analyzed): + spark, service = analyzed + df = spark.range(10) + assert df.columns == ["id"] + derived = df.selectExpr("id AS renamed") + service.data_type = _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() + assert other_session.range(10).columns == ["renamed"] + assert service.calls == 4 + assert df.columns == ["id"] + + +@pytest.mark.parametrize("failure", ["rpc", "no-result", "missing-schema", "invalid-type"]) +def test_failures_are_not_cached(analyzed, failure): + spark, service = 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_concurrent_schema_access_issues_one_analysis(analyzed): + spark, service = analyzed + df = spark.range(10) + with ThreadPoolExecutor(max_workers=8) as pool: + results = list(pool.map(lambda _: df.schema, range(8))) + assert all(schema.jsonValue() == results[0].jsonValue() for schema in results) + assert len({id(schema) for schema in results}) == 8 + assert service.calls == 1 From e8f491b5535ff0567a1d4cb98a4046728a3e20f7 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 03:58:24 -0700 Subject: [PATCH 2/9] [SPARK-60677][CONNECT] Include schema cache regressions in selected CI suite --- python/tests/test_dropin_offline.py | 226 ++++++++++++++++++++++++++++ python/tests/test_schema_cache.py | 225 --------------------------- 2 files changed, 226 insertions(+), 225 deletions(-) delete mode 100644 python/tests/test_schema_cache.py diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 67ccaa4d..4336551e 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -2755,3 +2755,229 @@ def test_create_dataframe_width_mismatch_is_a_normal_exception(): else: raise AssertionError("expected a width validation error") """], check=True) + +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + + +# --------------------------------------------------------------------------- schema cache +# Schema RPC regressions live here because coverage.yml explicitly selects this +# file. The separate-process fixture needs no live Spark server and leaves other +# offline tests available even if optional grpcio is not installed. +from concurrent.futures import ThreadPoolExecutor as schema_cache_ThreadPoolExecutor +import multiprocessing as schema_cache_multiprocessing + +from pyspark.sql import SparkSession as schema_cache_SparkSession + + +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(21, schema_cache_message(1, schema_cache_field("inner", schema_cache_message(7, b"")))) if nested else schema_cache_message(7, b"") + 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 + + # A separate process also exercises columns/dtypes, whose first access may + # hold the client's GIL while waiting for the RPC. + class Service: + calls = 0 + data_type = schema_cache_schema() + failure = None + + def analyze(self, request, context): + self.calls += 1 + 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 schema_cache_ThreadPoolExecutor(max_workers=4) as executor: + server = schema_cache_grpc.server(executor) + handler = schema_cache_grpc.unary_unary_rpc_method_handler(service.analyze) + server.add_generic_rpc_handlers((schema_cache_grpc.method_handlers_generic_handler( + "spark.connect.SparkConnectService", {"AnalyzePlan": handler}),)) + 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) + 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() + + +@pytest.fixture +def schema_cache_analyzed(): + pytest.importorskip("grpc") + context = schema_cache_multiprocessing.get_context("spawn") + parent, child = context.Pipe() + process = context.Process(target=schema_cache_serve, args=(child,)) + process.start() + try: + assert parent.poll(20), "Mock server did not start" + port = parent.recv() + spark = schema_cache_SparkSession.builder.remote(f"sc://127.0.0.1:{port}").getOrCreate() + yield spark, schema_cache_RemoteService(parent) + finally: + if process.is_alive(): + parent.send(("stop",)) + process.join(10) + if process.is_alive(): + process.terminate() + process.join(10) + parent.close() + child.close() + + +@pytest.mark.parametrize("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")] + 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() + assert other_session.range(10).columns == ["renamed"] + assert service.calls == 4 + assert df.columns == ["id"] + + +@pytest.mark.parametrize("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 schema_cache_ThreadPoolExecutor(max_workers=8) as pool: + results = list(pool.map(lambda _: df.schema, range(8))) + assert all(schema.jsonValue() == results[0].jsonValue() for schema in results) + assert len({id(schema) for schema in results}) == 8 + assert service.calls == 1 diff --git a/python/tests/test_schema_cache.py b/python/tests/test_schema_cache.py deleted file mode 100644 index e085e8f0..00000000 --- a/python/tests/test_schema_cache.py +++ /dev/null @@ -1,225 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one or more -# contributor license agreements. See the NOTICE file distributed with -# this work for additional information regarding copyright ownership. -# The ASF licenses this file to You under the Apache License, Version 2.0 -# (the "License"); you may not use this file except in compliance with -# the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Schema caching through the real client transport, without a Spark server. - -The fixture serves only AnalyzePlan. Tiny wire encoders keep the tests independent -of the reference PySpark client and its generated protobuf modules. -""" -from concurrent.futures import ThreadPoolExecutor -import multiprocessing - -import pytest - -grpc = pytest.importorskip("grpc") - -from pyspark.sql import SparkSession, types as T - - -def _varint(value): - result = bytearray() - while value > 127: - result.append((value & 127) | 128) - value >>= 7 - result.append(value) - return bytes(result) - - -def _message(field, value): - if isinstance(value, str): - value = value.encode() - return _varint(field * 8 + 2) + _varint(len(value)) + value - - -def _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 = _message(21, _message(1, _field("inner", _message(7, b"")))) if nested else _message(7, b"") - return _message(21, _message(1, _field(name, data_type))) - - -def _field(name, data_type): - return (_message(1, name) + _message(2, data_type) + b"\x18\x01" - + _message(4, '{"note":"original"}')) - - -def _serve(connection): - # A separate process also exercises columns/dtypes, whose first access may - # hold the client's GIL while waiting for the RPC. - class Service: - calls = 0 - data_type = _schema() - failure = None - - def analyze(self, request, context): - self.calls += 1 - 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 _message(2, b"") - if self.failure == "invalid-type": - return _message(2, _message(1, b"")) - # base.proto: AnalyzePlanResponse.schema=2, Schema.schema=1. - return _message(2, _message(1, self.data_type)) - - service = Service() - with ThreadPoolExecutor(max_workers=4) as executor: - server = grpc.server(executor) - handler = grpc.unary_unary_rpc_method_handler(service.analyze) - server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler( - "spark.connect.SparkConnectService", {"AnalyzePlan": handler}),)) - 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) - finally: - server.stop(0).wait() - - -class _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() - - -@pytest.fixture -def analyzed(): - context = multiprocessing.get_context("spawn") - parent, child = context.Pipe() - process = context.Process(target=_serve, args=(child,)) - process.start() - 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, _RemoteService(parent) - finally: - if process.is_alive(): - parent.send(("stop",)) - process.join(10) - if process.is_alive(): - process.terminate() - process.join(10) - parent.close() - child.close() - - -@pytest.mark.parametrize("first", ["schema", "columns", "dtypes"]) -def test_one_analysis_shared_by_metadata_accessors(analyzed, first): - spark, service = 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_returned_schema_and_lists_are_independent(analyzed): - spark, service = analyzed - service.data_type = _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")] - assert service.calls == 1 - - -def test_derived_and_new_plans_resolve_independently(analyzed): - spark, service = analyzed - df = spark.range(10) - assert df.columns == ["id"] - derived = df.selectExpr("id AS renamed") - service.data_type = _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() - assert other_session.range(10).columns == ["renamed"] - assert service.calls == 4 - assert df.columns == ["id"] - - -@pytest.mark.parametrize("failure", ["rpc", "no-result", "missing-schema", "invalid-type"]) -def test_failures_are_not_cached(analyzed, failure): - spark, service = 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_concurrent_schema_access_issues_one_analysis(analyzed): - spark, service = analyzed - df = spark.range(10) - with ThreadPoolExecutor(max_workers=8) as pool: - results = list(pool.map(lambda _: df.schema, range(8))) - assert all(schema.jsonValue() == results[0].jsonValue() for schema in results) - assert len({id(schema) for schema in results}) == 8 - assert service.calls == 1 From ee971071b70961db85e0152ff3eb9824cda41ddc Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 09:30:29 -0700 Subject: [PATCH 3/9] [SPARK-60677][CONNECT] Recover schema cache after panics and release metadata GIL --- crates/pyspark-rs/src/dataframe.rs | 8 +- crates/spark-connect/src/dataframe.rs | 37 +++++++- python/tests/test_dropin_offline.py | 120 ++++++++++++++++++-------- 3 files changed, 122 insertions(+), 43 deletions(-) diff --git a/crates/pyspark-rs/src/dataframe.rs b/crates/pyspark-rs/src/dataframe.rs index 7101aadf..3d980202 100644 --- a/crates/pyspark-rs/src/dataframe.rs +++ b/crates/pyspark-rs/src/dataframe.rs @@ -1302,8 +1302,8 @@ impl PyDataFrame { /// Column names (a property, mirroring `DataFrame.columns`). #[getter] - fn columns(&self) -> PyResult> { - self.dataframe.columns().to_pyerr() + fn columns(&self, py: Python<'_>) -> PyResult> { + py.detach(|| self.dataframe.columns()).to_pyerr() } /// Get the last `num` rows as a list of Rows (eager, like pyspark `tail`). @@ -1374,8 +1374,8 @@ impl PyDataFrame { /// Get the schema as a list of (name, type) tuples. #[getter] - fn dtypes(&self) -> PyResult> { - let schema = self.dataframe.schema().to_pyerr()?; + fn dtypes(&self, py: Python<'_>) -> PyResult> { + let schema = py.detach(|| self.dataframe.schema()).to_pyerr()?; match schema { spark_connect::types::DataType::Struct { fields } => { let dtypes = fields diff --git a/crates/spark-connect/src/dataframe.rs b/crates/spark-connect/src/dataframe.rs index 3dc034c0..18cfe8bb 100644 --- a/crates/spark-connect/src/dataframe.rs +++ b/crates/spark-connect/src/dataframe.rs @@ -773,13 +773,20 @@ 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 { // 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() - .map_err(|_| SparkError::connect_msg("Schema cache lock poisoned"))?; + // The cache contains only None or a fully decoded schema, so a + // panic during analysis cannot leave a partially written value. + .unwrap_or_else(std::sync::PoisonError::into_inner); if let Some(schema) = cached.as_ref() { return Ok(schema.clone()); } @@ -2894,6 +2901,34 @@ mod cache_tests { } } + #[test] + fn schema_cache_recovers_after_analysis_panics() { + let spark = SparkSession::builder() + .remote("sc://localhost:15002") + .get_or_create() + .unwrap(); + let df = spark.range(3).unwrap(); + let cloned = df.clone(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + + // Calling the synchronous API inside a runtime panics during analysis. + // A clone must still attempt analysis afterward, rather than returning + // a permanent cache-poisoning error. + for frame in [&df, &cloned] { + let result = runtime.block_on(async { + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| frame.schema())) + }); + assert!(result.is_err()); + } + let mut cached = df.schema_cache.lock().unwrap_err().into_inner(); + assert!(cached.is_none()); + let schema = DataType::Struct { fields: vec![] }; + *cached = Some(schema.clone()); + drop(cached); + assert_eq!(df.schema().unwrap(), schema); + assert_eq!(cloned.schema().unwrap(), schema); + } + #[test] fn cache_default_is_memory_and_disk_deser() { let sl = DataFrame::memory_and_disk_deser(); diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 4336551e..ad189bfa 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -6,10 +6,14 @@ that require a live server live in test_dropin_e2e.py and are gated on SPARK_REMOTE. """ +from concurrent.futures import ThreadPoolExecutor as schema_cache_ThreadPoolExecutor +import multiprocessing as schema_cache_multiprocessing import pickle +from threading import Event as schema_cache_Event import pytest +from pyspark.sql import SparkSession as schema_cache_SparkSession from pyspark.sql import types as T @@ -2756,31 +2760,11 @@ def test_create_dataframe_width_mismatch_is_a_normal_exception(): raise AssertionError("expected a width validation error") """], check=True) -# Licensed to the Apache Software Foundation (ASF) under one or more -# contributor license agreements. See the NOTICE file distributed with -# this work for additional information regarding copyright ownership. -# The ASF licenses this file to You under the Apache License, Version 2.0 -# (the "License"); you may not use this file except in compliance with -# the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - - # --------------------------------------------------------------------------- schema cache -# Schema RPC regressions live here because coverage.yml explicitly selects this -# file. The separate-process fixture needs no live Spark server and leaves other -# offline tests available even if optional grpcio is not installed. -from concurrent.futures import ThreadPoolExecutor as schema_cache_ThreadPoolExecutor -import multiprocessing as schema_cache_multiprocessing - -from pyspark.sql import SparkSession as schema_cache_SparkSession +# 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): @@ -2801,7 +2785,10 @@ def schema_cache_message(field, 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(21, schema_cache_message(1, schema_cache_field("inner", schema_cache_message(7, b"")))) if nested else schema_cache_message(7, b"") + 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))) @@ -2813,15 +2800,26 @@ def schema_cache_field(name, data_type): def schema_cache_serve(connection): import grpc as schema_cache_grpc - # A separate process also exercises columns/dtypes, whose first access may - # hold the client's GIL while waiting for the RPC. + # 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 = schema_cache_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": @@ -2836,9 +2834,12 @@ def analyze(self, request, context): service = Service() with schema_cache_ThreadPoolExecutor(max_workers=4) as executor: server = schema_cache_grpc.server(executor) - handler = schema_cache_grpc.unary_unary_rpc_method_handler(service.analyze) + 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", {"AnalyzePlan": handler}),)) + "spark.connect.SparkConnectService", handlers),)) port = server.add_insecure_port("127.0.0.1:0") server.start() connection.send(port) @@ -2850,7 +2851,10 @@ def analyze(self, request, context): if command[0] == "get": connection.send(getattr(service, command[1])) elif command[0] == "set": - setattr(service, command[1], command[2]) + if command[1] == "allow_response": + service.response_gate.set() + else: + setattr(service, command[1], command[2]) connection.send(None) finally: server.stop(0).wait() @@ -2878,20 +2882,28 @@ def schema_cache_analyzed(): 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 = schema_cache_SparkSession.builder.remote(f"sc://127.0.0.1:{port}").getOrCreate() - yield spark, schema_cache_RemoteService(parent) + yield spark, service finally: - if process.is_alive(): - parent.send(("stop",)) - process.join(10) - if process.is_alive(): - process.terminate() + 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) - parent.close() - child.close() + if process.is_alive(): + process.terminate() + process.join(10) + parent.close() + child.close() @pytest.mark.parametrize("first", ["schema", "columns", "dtypes"]) @@ -2947,7 +2959,10 @@ def test_schema_cache_derived_and_new_plans_resolve_independently(schema_cache_a fresh = spark.range(10) assert fresh.columns == ["renamed"] other_session = spark.newSession() - assert other_session.range(10).columns == ["renamed"] + try: + assert other_session.range(10).columns == ["renamed"] + finally: + other_session.stop() assert service.calls == 4 assert df.columns == ["id"] @@ -2981,3 +2996,32 @@ def test_schema_cache_concurrent_schema_access_issues_one_analysis(schema_cache_ assert all(schema.jsonValue() == results[0].jsonValue() for schema in results) assert len({id(schema) for schema in results}) == 8 assert service.calls == 1 + + +@pytest.mark.parametrize("accessor", ["schema", "columns", "dtypes"]) +def test_schema_cache_metadata_wait_releases_gil(schema_cache_analyzed, accessor): + spark, service = schema_cache_analyzed + df = spark.range(10) + service.block_response = True + + def allow_response_after_request(): + for _ in range(500): + if service.calls: + service.allow_response = True + return + schema_cache_Event().wait(0.01) + raise AssertionError("Metadata request never reached the server") + + # This Python thread must run while the main thread waits for the RPC. The + # server bounds its wait, so retaining the GIL fails instead of deadlocking. + with schema_cache_ThreadPoolExecutor(max_workers=1) as pool: + allowed = pool.submit(allow_response_after_request) + result = getattr(df, accessor) + allowed.result(timeout=10) + if accessor == "schema": + assert isinstance(result, T.StructType) + elif accessor == "columns": + assert result == ["id"] + else: + assert result == [("id", "bigint")] + assert service.calls == 1 From 7248d33506c25c110305d40a91dcfbba97685667 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 10:01:10 -0700 Subject: [PATCH 4/9] [SPARK-60677][CONNECT] Simplify schema cache tests and clear recovered poison --- crates/spark-connect/src/dataframe.rs | 7 ++++- python/tests/test_dropin_offline.py | 39 +++++++++++++++------------ 2 files changed, 28 insertions(+), 18 deletions(-) diff --git a/crates/spark-connect/src/dataframe.rs b/crates/spark-connect/src/dataframe.rs index 18cfe8bb..e91465e0 100644 --- a/crates/spark-connect/src/dataframe.rs +++ b/crates/spark-connect/src/dataframe.rs @@ -786,7 +786,10 @@ impl DataFrame { .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(std::sync::PoisonError::into_inner); + .unwrap_or_else(|poisoned| { + self.schema_cache.clear_poison(); + poisoned.into_inner() + }); if let Some(schema) = cached.as_ref() { return Ok(schema.clone()); } @@ -2925,7 +2928,9 @@ mod cache_tests { let schema = DataType::Struct { fields: vec![] }; *cached = Some(schema.clone()); drop(cached); + assert!(df.schema_cache.is_poisoned()); assert_eq!(df.schema().unwrap(), schema); + assert!(!df.schema_cache.is_poisoned()); assert_eq!(cloned.schema().unwrap(), schema); } diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index ad189bfa..fbe819eb 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -6,15 +6,15 @@ that require a live server live in test_dropin_e2e.py and are gated on SPARK_REMOTE. """ -from concurrent.futures import ThreadPoolExecutor as schema_cache_ThreadPoolExecutor -import multiprocessing as schema_cache_multiprocessing +import multiprocessing import pickle -from threading import Event as schema_cache_Event +import time +from concurrent.futures import ThreadPoolExecutor +from threading import Event import pytest -from pyspark.sql import SparkSession as schema_cache_SparkSession -from pyspark.sql import types as T +from pyspark.sql import SparkSession, types as T # --------------------------------------------------------------------------- types @@ -2807,7 +2807,7 @@ class Service: failure = None release_calls = 0 block_response = False - response_gate = schema_cache_Event() + response_gate = Event() def release(self, request, context): self.release_calls += 1 @@ -2832,7 +2832,7 @@ def analyze(self, request, context): return schema_cache_message(2, schema_cache_message(1, self.data_type)) service = Service() - with schema_cache_ThreadPoolExecutor(max_workers=4) as executor: + 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), @@ -2851,10 +2851,10 @@ def analyze(self, request, context): if command[0] == "get": connection.send(getattr(service, command[1])) elif command[0] == "set": - if command[1] == "allow_response": - service.response_gate.set() - else: - setattr(service, command[1], command[2]) + 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() @@ -2874,11 +2874,16 @@ def __setattr__(self, 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() + @pytest.fixture def schema_cache_analyzed(): pytest.importorskip("grpc") - context = schema_cache_multiprocessing.get_context("spawn") + context = multiprocessing.get_context("spawn") parent, child = context.Pipe() process = context.Process(target=schema_cache_serve, args=(child,)) process.start() @@ -2887,7 +2892,7 @@ def schema_cache_analyzed(): try: assert parent.poll(20), "Mock server did not start" port = parent.recv() - spark = schema_cache_SparkSession.builder.remote(f"sc://127.0.0.1:{port}").getOrCreate() + spark = SparkSession.builder.remote(f"sc://127.0.0.1:{port}").getOrCreate() yield spark, service finally: try: @@ -2991,7 +2996,7 @@ def test_schema_cache_failures_are_not_cached(schema_cache_analyzed, failure): def test_schema_cache_concurrent_schema_access_issues_one_analysis(schema_cache_analyzed): spark, service = schema_cache_analyzed df = spark.range(10) - with schema_cache_ThreadPoolExecutor(max_workers=8) as pool: + with ThreadPoolExecutor(max_workers=8) as pool: results = list(pool.map(lambda _: df.schema, range(8))) assert all(schema.jsonValue() == results[0].jsonValue() for schema in results) assert len({id(schema) for schema in results}) == 8 @@ -3007,14 +3012,14 @@ def test_schema_cache_metadata_wait_releases_gil(schema_cache_analyzed, accessor def allow_response_after_request(): for _ in range(500): if service.calls: - service.allow_response = True + service.release_response() return - schema_cache_Event().wait(0.01) + time.sleep(0.01) raise AssertionError("Metadata request never reached the server") # This Python thread must run while the main thread waits for the RPC. The # server bounds its wait, so retaining the GIL fails instead of deadlocking. - with schema_cache_ThreadPoolExecutor(max_workers=1) as pool: + with ThreadPoolExecutor(max_workers=1) as pool: allowed = pool.submit(allow_response_after_request) result = getattr(df, accessor) allowed.result(timeout=10) From d0183ef47f22bd6d213b2a94b8d7d5fbd97bef27 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 13:15:27 -0700 Subject: [PATCH 5/9] [SPARK-60677][CONNECT] Reject malformed schemas and strengthen cache regressions Reject non-struct analysis responses before caching, verify nested value isolation, and force overlapping metadata reads in the concurrency regression. Convert errors inside detached Python metadata calls to avoid adding Clippy warnings. --- crates/pyspark-rs/src/dataframe.rs | 4 +- crates/spark-connect/src/dataframe.rs | 49 +++++++++++++ python/tests/test_dropin_offline.py | 100 +++++++++++++++++++------- 3 files changed, 127 insertions(+), 26 deletions(-) diff --git a/crates/pyspark-rs/src/dataframe.rs b/crates/pyspark-rs/src/dataframe.rs index 3d980202..b97596a4 100644 --- a/crates/pyspark-rs/src/dataframe.rs +++ b/crates/pyspark-rs/src/dataframe.rs @@ -1303,7 +1303,7 @@ impl PyDataFrame { /// Column names (a property, mirroring `DataFrame.columns`). #[getter] fn columns(&self, py: Python<'_>) -> PyResult> { - py.detach(|| self.dataframe.columns()).to_pyerr() + py.detach(|| self.dataframe.columns().to_pyerr()) } /// Get the last `num` rows as a list of Rows (eager, like pyspark `tail`). @@ -1375,7 +1375,7 @@ impl PyDataFrame { /// Get the schema as a list of (name, type) tuples. #[getter] fn dtypes(&self, py: Python<'_>) -> PyResult> { - let schema = py.detach(|| self.dataframe.schema()).to_pyerr()?; + let schema = py.detach(|| self.dataframe.schema().to_pyerr())?; match schema { spark_connect::types::DataType::Struct { fields } => { let dtypes = fields diff --git a/crates/spark-connect/src/dataframe.rs b/crates/spark-connect/src/dataframe.rs index e91465e0..34a48fef 100644 --- a/crates/spark-connect/src/dataframe.rs +++ b/crates/spark-connect/src/dataframe.rs @@ -802,6 +802,9 @@ impl DataFrame { .schema .ok_or_else(|| SparkError::connect_msg("Schema is missing"))?, )?; + if !matches!(&schema, DataType::Struct { .. }) { + return Err(SparkError::connect_msg("Schema is not a struct type")); + } *cached = Some(schema.clone()); Ok(schema) } else { @@ -2904,6 +2907,52 @@ mod cache_tests { } } + #[test] + fn schema_cache_returns_independent_nested_values() { + let spark = SparkSession::builder() + .remote("sc://localhost:15002") + .get_or_create() + .unwrap(); + let df = spark.range(3).unwrap(); + let cloned = df.clone(); + let schema = DataType::from_json(&serde_json::json!({ + "type": "struct", + "fields": [{ + "name": "outer", + "type": { + "type": "struct", + "fields": [{ + "name": "inner", + "type": "long", + "nullable": true, + "metadata": {"details": {"labels": ["original"]}} + }] + }, + "nullable": true, + "metadata": {} + }] + })) + .unwrap(); + *df.schema_cache.lock().unwrap() = Some(schema.clone()); + + let mut changed = df.schema().unwrap(); + let unchanged = df.schema().unwrap(); + let DataType::Struct { fields } = &mut changed else { + panic!("Expected outer struct"); + }; + let DataType::Struct { fields } = &mut fields[0].data_type else { + panic!("Expected nested struct"); + }; + fields[0].name = "changed".to_string(); + fields[0].metadata.get_mut("details").unwrap()["labels"][0] = serde_json::json!("changed"); + assert_eq!(fields[0].name, "changed"); + assert_eq!(fields[0].metadata["details"]["labels"][0], "changed"); + assert_ne!(changed, schema); + assert_eq!(unchanged, schema); + assert_eq!(df.schema_cache.lock().unwrap().as_ref(), Some(&schema)); + assert_eq!(cloned.schema().unwrap(), schema); + } + #[test] fn schema_cache_recovers_after_analysis_panics() { let spark = SparkSession::builder() diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index fbe819eb..fb60a87e 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -10,7 +10,7 @@ import pickle import time from concurrent.futures import ThreadPoolExecutor -from threading import Event +from threading import Barrier, Event import pytest @@ -2793,12 +2793,16 @@ def schema_cache_schema(name="id", nested=False): 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"}')) + 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 + import grpc # Keep the server independent of the client's interpreter and GIL. class Service: @@ -2817,11 +2821,11 @@ 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, + 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") + context.abort(grpc.StatusCode.INVALID_ARGUMENT, "schema unavailable") if self.failure == "no-result": return b"" if self.failure == "missing-schema": @@ -2833,13 +2837,14 @@ def analyze(self, request, context): service = Service() with ThreadPoolExecutor(max_workers=4) as executor: - server = schema_cache_grpc.server(executor) + server = 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), + "AnalyzePlan": grpc.unary_unary_rpc_method_handler(service.analyze), + "ReleaseSession": grpc.unary_unary_rpc_method_handler(service.release), } - server.add_generic_rpc_handlers((schema_cache_grpc.method_handlers_generic_handler( - "spark.connect.SparkConnectService", handlers),)) + 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) @@ -2860,7 +2865,7 @@ def analyze(self, request, context): server.stop(0).wait() -class schema_cache_RemoteService: +class SchemaCacheRemoteService: def __init__(self, connection): object.__setattr__(self, "connection", connection) @@ -2888,7 +2893,7 @@ def schema_cache_analyzed(): process = context.Process(target=schema_cache_serve, args=(child,)) process.start() spark = None - service = schema_cache_RemoteService(parent) + service = SchemaCacheRemoteService(parent) try: assert parent.poll(20), "Mock server did not start" port = parent.recv() @@ -2935,15 +2940,15 @@ def test_schema_cache_returned_schema_and_lists_are_independent(schema_cache_ana 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() + columns = df.columns + columns.append("extra") + dtypes = df.dtypes + dtypes.clear() + assert columns == ["id", "extra"] + assert dtypes == [] assert df.schema.jsonValue() == original.jsonValue() assert df.columns == ["id"] assert df.dtypes == [("id", "struct")] @@ -2993,13 +2998,60 @@ def test_schema_cache_failures_are_not_cached(schema_cache_analyzed, failure): assert service.calls == 3 -def test_schema_cache_concurrent_schema_access_issues_one_analysis(schema_cache_analyzed): +@pytest.mark.parametrize("accessor", ["schema", "columns", "dtypes"]) +def test_schema_cache_non_struct_responses_are_not_cached(schema_cache_analyzed, accessor): 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))) - assert all(schema.jsonValue() == results[0].jsonValue() for schema in results) - assert len({id(schema) for schema in results}) == 8 + service.data_type = schema_cache_message(7, b"") + for count in (1, 2): + with pytest.raises(Exception, match="Schema is not a struct type"): + getattr(df, accessor) + assert service.calls == count + service.data_type = schema_cache_schema() + assert df.columns == ["id"] + assert df.dtypes == [("id", "bigint")] + assert service.calls == 3 + + +def test_schema_cache_concurrent_metadata_access_issues_one_analysis(schema_cache_analyzed): + spark, service = schema_cache_analyzed + df = spark.range(10) + accessors = ("schema", "columns", "dtypes") * 3 + start = Barrier(len(accessors) + 1) + service.block_response = True + + def read_metadata(accessor): + start.wait(timeout=10) + return getattr(df, accessor) + + with ThreadPoolExecutor(max_workers=len(accessors)) as pool: + pending = [pool.submit(read_metadata, accessor) for accessor in accessors] + try: + start.wait(timeout=10) + for _ in range(500): + if service.calls: + break + time.sleep(0.01) + else: + raise AssertionError("Metadata request never reached the server") + # Keep the first RPC in flight while the other workers attempt their + # cold lookups. An implementation without miss synchronization must + # not pass just because the first response arrives immediately. + time.sleep(0.1) + assert not any(future.done() for future in pending) + finally: + service.release_response() + results = [future.result(timeout=10) for future in pending] + schemas = [] + for accessor, result in zip(accessors, results): + if accessor == "schema": + schemas.append(result) + elif accessor == "columns": + assert result == ["id"] + else: + assert result == [("id", "bigint")] + assert all(schema.jsonValue() == schemas[0].jsonValue() for schema in schemas) + assert len({id(schema) for schema in schemas}) == len(schemas) assert service.calls == 1 From 76ced15357b24337e7e874d9950fbb8dbc931d3a Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 13:31:51 -0700 Subject: [PATCH 6/9] [SPARK-60677][CONNECT] Restore nested Python schema mutation checks Retain nested getter results and assert their mutations. Compare original and cached schemas with a pre-mutation snapshot so shared changes cannot mask isolation failures. --- python/tests/test_dropin_offline.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index fb60a87e..3ea9e3db 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -2938,18 +2938,27 @@ def test_schema_cache_returned_schema_and_lists_are_independent(schema_cache_ana 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() == original.jsonValue() + assert df.schema.jsonValue() == expected_json assert df.columns == ["id"] assert df.dtypes == [("id", "struct")] assert service.calls == 1 From 2ee1b3acc8975f2976791c66f4bf2604fdf00096 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 13:46:17 -0700 Subject: [PATCH 7/9] [SPARK-60677][CONNECT] Make schema analysis lock boundary explicit Assert that the shared cache mutex is poisoned after each analysis panic, so a lock-scope regression fails at the analysis boundary with a clear message. The existing final poison check already rejected an early unlock; this makes the intended invariant explicit without scheduler assumptions. --- crates/spark-connect/src/dataframe.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/crates/spark-connect/src/dataframe.rs b/crates/spark-connect/src/dataframe.rs index 34a48fef..c59f6a1c 100644 --- a/crates/spark-connect/src/dataframe.rs +++ b/crates/spark-connect/src/dataframe.rs @@ -2971,6 +2971,12 @@ mod cache_tests { std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| frame.schema())) }); assert!(result.is_err()); + // Poisoning proves the lock spans the analysis boundary without + // relying on thread scheduling to produce contended cache misses. + assert!( + frame.schema_cache.is_poisoned(), + "schema analysis must hold the shared cache lock" + ); } let mut cached = df.schema_cache.lock().unwrap_err().into_inner(); assert!(cached.is_none()); From ed7efa6c1c427ef9b63ca94c3f8c578d34e1217e Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 17:17:47 -0700 Subject: [PATCH 8/9] [SPARK-60677][CONNECT] Propagate resolved DataFrame schemas Snapshot warm schemas through Spark 4.2's selected preserving transforms and equal-schema set operations. Keep derived cache locks independent and share only immutable schema values; nonblocking snapshots avoid waiting for an in-flight analysis while Python holds the GIL. Seed declared struct outputs for to and stateless partition/group UDFs. Extract the reusable offline Connect server, simplify redundant tests and branches, and document blocking readers and defensive panic recovery. Validate with 398 Rust library tests, 318 offline Python tests, and 19 live Spark 4.2 comparison cases. The prior native build and a blocking-snapshot mutation fail the new regressions as intended. Clippy matches the base's 923 diagnostics with no additions. --- crates/pyspark-rs/src/dataframe.rs | 22 +- crates/spark-connect/src/dataframe.rs | 285 +++++++++++++------- crates/spark-connect/src/group.rs | 44 +++- python/tests/_mock_connect_server.py | 160 +++++++++++ python/tests/test_dropin_offline.py | 364 +++++++++++++------------- 5 files changed, 583 insertions(+), 292 deletions(-) create mode 100644 python/tests/_mock_connect_server.py diff --git a/crates/pyspark-rs/src/dataframe.rs b/crates/pyspark-rs/src/dataframe.rs index b97596a4..162bc96e 100644 --- a/crates/pyspark-rs/src/dataframe.rs +++ b/crates/pyspark-rs/src/dataframe.rs @@ -1375,19 +1375,15 @@ impl PyDataFrame { /// Get the schema as a list of (name, type) tuples. #[getter] fn dtypes(&self, py: Python<'_>) -> PyResult> { - let schema = py.detach(|| self.dataframe.schema().to_pyerr())?; - match schema { - spark_connect::types::DataType::Struct { fields } => { - let dtypes = fields - .iter() - .map(|f| (f.name.clone(), f.data_type.simple_string())) - .collect(); - Ok(dtypes) - } - _ => Err(PyErr::new::( - "Schema is not a struct type", - )), - } + let spark_connect::types::DataType::Struct { fields } = + py.detach(|| self.dataframe.schema().to_pyerr())? + else { + unreachable!("schema() validates struct responses") + }; + Ok(fields + .into_iter() + .map(|field| (field.name, field.data_type.simple_string())) + .collect()) } /// Additional set operation: unionAll (alias for union). diff --git a/crates/spark-connect/src/dataframe.rs b/crates/spark-connect/src/dataframe.rs index c59f6a1c..0802fe9a 100644 --- a/crates/spark-connect/src/dataframe.rs +++ b/crates/spark-connect/src/dataframe.rs @@ -26,9 +26,10 @@ use crate::udf::CommonInlineUserDefinedFunctionExpression; pub struct DataFrame { pub(crate) session: SparkSession, pub(crate) plan: Arc, - // Exact clones share this cache. Every new plan starts with an empty cache, - // even when its input has already been analyzed. - schema_cache: Arc>>, + // Exact clones share this cache. Derived plans can snapshot an already + // resolved schema, but have their own cache for subsequent analysis. + // The resolved value is immutable, so propagation need not copy wide schemas. + schema_cache: Arc>>>, } /// Iterator over rows from a DataFrame, yielded lazily as the server streams results. @@ -206,11 +207,57 @@ fn spawn_prefetch( impl DataFrame { /// Create a new DataFrame. pub(crate) fn new(session: SparkSession, plan: LogicalPlan) -> Self { + Self::new_with_schema(session, plan, None) + } + + pub(crate) fn new_with_schema( + session: SparkSession, + plan: LogicalPlan, + schema: Option, + ) -> Self { + Self::new_with_cached_schema( + session, + plan, + schema + .filter(|schema| matches!(schema, DataType::Struct { .. })) + .map(Arc::new), + ) + } + + fn new_with_cached_schema( + session: SparkSession, + plan: LogicalPlan, + schema: Option>, + ) -> Self { DataFrame { session, plan: Arc::new(plan), - schema_cache: Arc::new(Mutex::new(None)), + schema_cache: Arc::new(Mutex::new(schema)), + } + } + + fn cached_schema(&self) -> Option> { + // Propagation snapshots resolved metadata; it must not wait for an + // in-flight analysis. Python transformations can still hold the GIL. + self.schema_cache.try_lock().ok()?.clone() + } + + fn with_cached_schema(&self, plan: LogicalPlan) -> Self { + Self::new_with_cached_schema(self.session.clone(), plan, self.cached_schema()) + } + + fn with_merged_schema(&self, other: &DataFrame, plan: LogicalPlan) -> Self { + // A right-hand plan from another session is analyzed in our session; + // its previously resolved schema is not evidence about that resolution. + if self.session.client().session_id() != other.session.client().session_id() { + return Self::new(self.session.clone(), plan); } + // Snapshot separately, including for self-unions; never lock both inputs. + let schema = match (self.cached_schema(), other.cached_schema()) { + (Some(left), Some(right)) if left == right => Some(left), + _ => None, + }; + Self::new_with_cached_schema(self.session.clone(), plan, schema) } /// Get the underlying logical plan. @@ -235,7 +282,7 @@ impl DataFrame { input: Arc::clone(&self.plan), condition, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Alias for filter(). @@ -304,7 +351,7 @@ impl DataFrame { input: Arc::clone(&self.plan), limit: n, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Skip the first n rows. @@ -333,7 +380,7 @@ impl DataFrame { column_names: vec![], within_watermark: false, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Remove duplicate rows, optionally on specific columns. @@ -349,7 +396,7 @@ impl DataFrame { column_names: cols, within_watermark: false, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Sort rows. @@ -359,7 +406,7 @@ impl DataFrame { order: columns, is_global: true, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Alias for sort(). @@ -455,7 +502,7 @@ impl DataFrame { by_name: false, allow_missing_columns: false, }; - DataFrame::new(self.session.clone(), plan) + self.with_merged_schema(other, plan) } /// Union by name. @@ -475,7 +522,7 @@ impl DataFrame { by_name: true, allow_missing_columns, }; - DataFrame::new(self.session.clone(), plan) + self.with_merged_schema(other, plan) } /// Intersect with another DataFrame. @@ -488,7 +535,7 @@ impl DataFrame { by_name: false, allow_missing_columns: false, }; - DataFrame::new(self.session.clone(), plan) + self.with_merged_schema(other, plan) } /// Subtract (except) another DataFrame. @@ -501,7 +548,7 @@ impl DataFrame { by_name: false, allow_missing_columns: false, }; - DataFrame::new(self.session.clone(), plan) + self.with_merged_schema(other, plan) } /// Repartition. @@ -511,7 +558,7 @@ impl DataFrame { num_partitions, shuffle: true, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Coalesce. @@ -521,7 +568,7 @@ impl DataFrame { num_partitions, shuffle: false, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Add a hint. @@ -535,7 +582,7 @@ impl DataFrame { name: name.to_string(), parameters: parameters.into_iter().map(Into::into).collect(), }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Marks a DataFrame as eligible for broadcast join (smaller table). @@ -560,7 +607,7 @@ impl DataFrame { input: Arc::clone(&self.plan), alias: alias.to_string(), }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Map over each partition with a pandas UDF (`DataFrame.mapInPandas`). @@ -590,12 +637,13 @@ impl DataFrame { func: CommonInlineUserDefinedFunctionExpression, is_barrier: bool, ) -> DataFrame { + let schema = func.python_udf.output_type.clone(); let plan = LogicalPlan::MapPartitions { input: Arc::clone(&self.plan), func, is_barrier, }; - DataFrame::new(self.session.clone(), plan) + Self::new_with_schema(self.session.clone(), plan, Some(schema)) } /// Apply a function to each row for its side effects (`DataFrame.foreach`). @@ -633,7 +681,7 @@ impl DataFrame { with_replacement, seed, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Group by columns for aggregation. @@ -775,23 +823,31 @@ impl DataFrame { /// Get the schema of this DataFrame. /// /// The first successful lookup is cached and shared with exact clones. + /// Selected schema-preserving transformations snapshot an already resolved + /// schema; set operations do so only when both cached input schemas match + /// exactly in the same session. Declared struct output schemas also seed + /// the cache. Other plans resolve independently. + /// Concurrent lookups on the same DataFrame or its clones wait for a single + /// analysis RPC, so a stalled analysis blocks those metadata readers. /// 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. + /// recreate the DataFrame from its source to resolve the schema again. + /// Each call returns an independent value, and failed lookups can be retried. pub fn schema(&self) -> Result { // 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. + // Unlike the other cache locks, this guard spans a blocking RPC. + // Recover defensively from Rust callers that panic during analysis + // (for example, by calling the sync API inside a Tokio runtime). + // The slot contains only None or a complete schema, so it is safe. .unwrap_or_else(|poisoned| { self.schema_cache.clear_poison(); poisoned.into_inner() }); if let Some(schema) = cached.as_ref() { - return Ok(schema.clone()); + return Ok(schema.as_ref().clone()); } let request = self.build_analyze_request()?; let response = block_on(self.session.client().analyze_plan(request))?; @@ -805,7 +861,7 @@ impl DataFrame { if !matches!(&schema, DataType::Struct { .. }) { return Err(SparkError::connect_msg("Schema is not a struct type")); } - *cached = Some(schema.clone()); + *cached = Some(Arc::new(schema.clone())); Ok(schema) } else { Err(SparkError::connect_msg( @@ -836,11 +892,10 @@ impl DataFrame { /// Get column names. pub fn columns(&self) -> Result> { - let schema = self.schema()?; - match schema { - DataType::Struct { fields } => Ok(fields.iter().map(|f| f.name.clone()).collect()), - _ => Err(SparkError::connect_msg("Schema is not a struct type")), - } + let DataType::Struct { fields } = self.schema()? else { + unreachable!("schema() validates struct responses") + }; + Ok(fields.into_iter().map(|field| field.name).collect()) } /// Build an ExecutePlanRequest from the logical plan. @@ -1119,7 +1174,7 @@ impl DataFrame { num_partitions: Some(num_partitions), partition_exprs: columns, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Repartition into `num_partitions` by hashing the given column expressions. @@ -1134,7 +1189,7 @@ impl DataFrame { num_partitions, expressions: columns, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Alias for to_df(). @@ -1191,7 +1246,7 @@ impl DataFrame { name: name.to_string(), exprs, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Get stat functions. @@ -1386,7 +1441,7 @@ impl DataFrame { order: columns, is_global: false, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Drop duplicates within a watermark. @@ -1430,7 +1485,7 @@ impl DataFrame { with_replacement: false, seed, }; - results.push(DataFrame::new(self.session.clone(), plan)); + results.push(self.with_cached_schema(plan)); cumulative = upper; } @@ -1471,17 +1526,13 @@ impl DataFrame { /// Get dtypes (column names and types). pub fn dtypes(&self) -> Result> { - let schema = self.schema()?; - match schema { - DataType::Struct { fields } => { - let dtypes = fields - .iter() - .map(|f| (f.name.clone(), f.data_type.to_string())) - .collect(); - Ok(dtypes) - } - _ => Err(SparkError::connect_msg("Schema is not a struct type")), - } + let DataType::Struct { fields } = self.schema()? else { + unreachable!("schema() validates struct responses") + }; + Ok(fields + .into_iter() + .map(|field| (field.name, field.data_type.to_string())) + .collect()) } /// Compute the server-side semantic hash of this DataFrame's logical plan, @@ -1569,7 +1620,7 @@ impl DataFrame { by_name: false, allow_missing_columns: false, }; - DataFrame::new(self.session.clone(), plan) + self.with_merged_schema(other, plan) } /// Except all rows. @@ -1582,7 +1633,7 @@ impl DataFrame { by_name: false, allow_missing_columns: false, }; - DataFrame::new(self.session.clone(), plan) + self.with_merged_schema(other, plan) } /// Intersect all rows. @@ -1595,7 +1646,7 @@ impl DataFrame { by_name: false, allow_missing_columns: false, }; - DataFrame::new(self.session.clone(), plan) + self.with_merged_schema(other, plan) } /// Unpivot columns (like melt). @@ -1729,9 +1780,9 @@ impl DataFrame { pub fn to(&self, schema: DataType) -> DataFrame { let plan = LogicalPlan::ToSchema { input: Arc::clone(&self.plan), - schema, + schema: schema.clone(), }; - DataFrame::new(self.session.clone(), plan) + Self::new_with_schema(self.session.clone(), plan, Some(schema)) } /// Check if the DataFrame exists (is not empty). @@ -2878,7 +2929,7 @@ mod cache_tests { use prost::Message; #[test] - fn schema_cache_follows_exact_clones_but_not_new_plans() { + fn schema_cache_shares_exact_clones_but_not_changed_plans_or_new_sources() { let spark = SparkSession::builder() .remote("sc://localhost:15002") .get_or_create() @@ -2886,7 +2937,7 @@ mod cache_tests { 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()); + *df.schema_cache.lock().unwrap() = Some(Arc::new(schema.clone())); // These accesses succeed without a server, including a clone created // before the schema was populated. @@ -2897,7 +2948,6 @@ mod cache_tests { &cloned_before_analysis.schema_cache )); for derived in [ - df.limit(1), df.select(["id"]), df.zip(&df).unwrap(), spark.range(3).unwrap(), @@ -2908,49 +2958,100 @@ mod cache_tests { } #[test] - fn schema_cache_returns_independent_nested_values() { + fn schema_propagation_snapshots_without_waiting_for_analysis() { let spark = SparkSession::builder() .remote("sc://localhost:15002") .get_or_create() .unwrap(); let df = spark.range(3).unwrap(); - let cloned = df.clone(); - let schema = DataType::from_json(&serde_json::json!({ + let early_child = df.limit(1); + let schema = DataType::Struct { fields: vec![] }; + *df.schema_cache.lock().unwrap() = Some(Arc::new(schema.clone())); + let warm_child = df.limit(1); + assert_eq!(warm_child.schema().unwrap(), schema); + assert!(!Arc::ptr_eq(&df.schema_cache, &warm_child.schema_cache)); + assert!(early_child.cached_schema().is_none()); + + // Even when an input lock is held, constructing unary and binary plans + // must finish. Release it before joining so a regression fails boundedly. + let guard = df.schema_cache.lock().unwrap(); + let input = df.clone(); + let (tx, rx) = std::sync::mpsc::channel(); + let worker = std::thread::spawn(move || { + tx.send((input.limit(1), input.union(&input))).unwrap(); + }); + let result = rx.recv_timeout(std::time::Duration::from_secs(5)); + drop(guard); + worker.join().unwrap(); + let (unary, binary) = result.expect("plan construction must not wait for schema analysis"); + assert!(unary.cached_schema().is_none()); + assert!(binary.cached_schema().is_none()); + } + + #[test] + fn schema_propagation_merges_only_exact_schemas_from_the_same_session() { + let spark = SparkSession::builder() + .remote("sc://localhost:15002") + .get_or_create() + .unwrap(); + let df = spark.range(3).unwrap(); + let schema = serde_json::json!({ "type": "struct", - "fields": [{ - "name": "outer", - "type": { - "type": "struct", - "fields": [{ - "name": "inner", - "type": "long", - "nullable": true, - "metadata": {"details": {"labels": ["original"]}} - }] - }, - "nullable": true, - "metadata": {} - }] - })) - .unwrap(); - *df.schema_cache.lock().unwrap() = Some(schema.clone()); - - let mut changed = df.schema().unwrap(); - let unchanged = df.schema().unwrap(); - let DataType::Struct { fields } = &mut changed else { - panic!("Expected outer struct"); - }; - let DataType::Struct { fields } = &mut fields[0].data_type else { - panic!("Expected nested struct"); - }; - fields[0].name = "changed".to_string(); - fields[0].metadata.get_mut("details").unwrap()["labels"][0] = serde_json::json!("changed"); - assert_eq!(fields[0].name, "changed"); - assert_eq!(fields[0].metadata["details"]["labels"][0], "changed"); - assert_ne!(changed, schema); - assert_eq!(unchanged, schema); - assert_eq!(df.schema_cache.lock().unwrap().as_ref(), Some(&schema)); - assert_eq!(cloned.schema().unwrap(), schema); + "fields": [{"name": "id", "type": "long", "nullable": true, + "metadata": {"note": "original"}}] + }); + let left = df.to(DataType::from_json(&schema).unwrap()); + assert_eq!(left.union(&left).schema().unwrap(), left.schema().unwrap()); + let mut changes = vec![]; + for (key, value) in [ + ("name", serde_json::json!("renamed")), + ("type", serde_json::json!("double")), + ("nullable", serde_json::json!(false)), + ("metadata", serde_json::json!({"note": "changed"})), + ] { + let mut changed = schema.clone(); + changed["fields"][0][key] = value; + changes.push(DataType::from_json(&changed).unwrap()); + } + for changed in changes { + assert!(left.union(&df.to(changed)).cached_schema().is_none()); + } + assert!(left.union(&df).cached_schema().is_none()); + let other_session = spark.new_session(); + let other = other_session.range(3).unwrap().to(left.schema().unwrap()); + assert!(left.union(&other).cached_schema().is_none()); + } + + #[test] + fn schema_propagation_seeds_only_struct_output_schemas() { + use crate::udf::PythonUDFPayload; + + let spark = SparkSession::builder() + .remote("sc://localhost:15002") + .get_or_create() + .unwrap(); + let df = spark.range(3).unwrap(); + let schema = DataType::Struct { fields: vec![] }; + let udf = |output_type| { + CommonInlineUserDefinedFunctionExpression::new( + "map".to_string(), + true, + vec![], + PythonUDFPayload::new(output_type, 205, vec![], "3.12".to_string()), + ) + }; + for output in [ + df.to(schema.clone()), + df.map_in_pandas(udf(schema.clone()), false), + df.map_in_arrow(udf(schema.clone()), false), + ] { + assert_eq!(output.schema().unwrap(), schema); + } + assert!(df.to(DataType::Long).cached_schema().is_none()); + assert!(df + .map_in_arrow(udf(DataType::Long), false) + .cached_schema() + .is_none()); } #[test] @@ -2981,7 +3082,7 @@ mod cache_tests { let mut cached = df.schema_cache.lock().unwrap_err().into_inner(); assert!(cached.is_none()); let schema = DataType::Struct { fields: vec![] }; - *cached = Some(schema.clone()); + *cached = Some(Arc::new(schema.clone())); drop(cached); assert!(df.schema_cache.is_poisoned()); assert_eq!(df.schema().unwrap(), schema); diff --git a/crates/spark-connect/src/group.rs b/crates/spark-connect/src/group.rs index 0047a677..54c9eead 100644 --- a/crates/spark-connect/src/group.rs +++ b/crates/spark-connect/src/group.rs @@ -141,6 +141,7 @@ impl GroupedData { } fn group_map(&self, func: CommonInlineUserDefinedFunctionExpression) -> DataFrame { + let schema = func.python_udf.output_type.clone(); let plan = LogicalPlan::GroupMap { input: Arc::clone(&self.dataframe.plan), grouping_expressions: self.grouping_expressions(), @@ -154,7 +155,7 @@ impl GroupedData { state_schema: None, transform_with_state_info: None, }; - DataFrame::new(self.dataframe.session.clone(), plan) + DataFrame::new_with_schema(self.dataframe.session.clone(), plan, Some(schema)) } /// Apply a stateful pandas UDF to each group (`GroupedData.applyInPandasWithState`). @@ -525,6 +526,7 @@ impl CoGroupedData { } fn cogroup_map(&self, func: CommonInlineUserDefinedFunctionExpression) -> DataFrame { + let schema = func.python_udf.output_type.clone(); let plan = LogicalPlan::CoGroupMap { input: Arc::clone(&self.left.dataframe.plan), input_grouping_expressions: self.left.grouping_expressions(), @@ -532,7 +534,7 @@ impl CoGroupedData { other_grouping_expressions: self.right.grouping_expressions(), func, }; - DataFrame::new(self.left.dataframe.session.clone(), plan) + DataFrame::new_with_schema(self.left.dataframe.session.clone(), plan, Some(schema)) } } @@ -571,6 +573,44 @@ mod tests { assert!(Arc::ptr_eq(input, &df.plan)); } + #[test] + fn grouped_udfs_reuse_declared_output_schema() { + use crate::udf::{eval_type, PythonUDFPayload}; + + let spark = session(); + let grouped = spark.range(10).unwrap().group_by(Vec::::new()); + let cogrouped = grouped.cogroup(&grouped); + let schema = DataType::from_json(&serde_json::json!({ + "type": "struct", + "fields": [{ + "name": "result", + "type": "long", + "nullable": false, + "metadata": {"description": "declared output"} + }] + })) + .unwrap(); + let udf = |eval_type| { + CommonInlineUserDefinedFunctionExpression::new( + "output".to_string(), + true, + vec![], + PythonUDFPayload::new(schema.clone(), eval_type, vec![], "3.12".to_string()), + ) + }; + + // These accesses must work without a server and use the supplied output + // schema, which differs from the unresolved input's `id` column. + for result in [ + grouped.apply_in_pandas(udf(eval_type::SQL_GROUPED_MAP_PANDAS_UDF)), + grouped.apply_in_arrow(udf(eval_type::SQL_GROUPED_MAP_ARROW_UDF)), + cogrouped.apply_in_pandas(udf(eval_type::SQL_COGROUPED_MAP_PANDAS_UDF)), + cogrouped.apply_in_arrow(udf(eval_type::SQL_COGROUPED_MAP_ARROW_UDF)), + ] { + assert_eq!(result.schema().unwrap(), schema); + } + } + #[test] fn group_count_plan() { let spark = session(); diff --git a/python/tests/_mock_connect_server.py b/python/tests/_mock_connect_server.py new file mode 100644 index 00000000..555f829a --- /dev/null +++ b/python/tests/_mock_connect_server.py @@ -0,0 +1,160 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Small, separately hosted Spark Connect service for transport-level tests.""" + +import multiprocessing +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from threading import Event + + +# The drop-in package does not ship generated service protobuf classes, so these +# encoders provide just the messages needed by this test server. +def varint(value): + result = bytearray() + while value > 127: + result.append((value & 127) | 128) + value >>= 7 + result.append(value) + return bytes(result) + + +def message_field(field, value): + if isinstance(value, str): + value = value.encode() + return varint(field * 8 + 2) + varint(len(value)) + value + + +def struct_schema(name="id"): + # types.proto: DataType.long=7, .struct=21; Struct.fields=1; + # StructField: name=1, data_type=2, nullable=3, metadata=4. + field = ( + message_field(1, name) + + message_field(2, message_field(7, b"")) + + b"\x18\x01" + + message_field(4, '{"note":"original"}') + ) + return message_field(21, message_field(1, field)) + + +def serve(connection): + import grpc + + # A separate interpreter lets tests detect clients retaining the GIL while + # they wait for a response. + class Service: + calls = 0 + data_type = struct_schema() + failure = None + block_response = False + response_gate = Event() + + 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 message_field(2, b"") + if self.failure == "invalid-type": + return message_field(2, message_field(1, b"")) + # base.proto: AnalyzePlanResponse.schema=2, Schema.schema=1. + return message_field(2, message_field(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(lambda request, context: b""), + } + 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 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() + + +@contextmanager +def mock_connect_session(): + from pyspark.sql import SparkSession + + context = multiprocessing.get_context("spawn") + parent, child = context.Pipe() + process = context.Process(target=serve, args=(child,)) + process.start() + spark = None + 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, RemoteService(parent) + finally: + try: + if spark is not None: + spark.stop() + finally: + if process.is_alive(): + parent.send(("stop",)) + process.join(10) + if process.is_alive(): + process.terminate() + process.join(10) + parent.close() + child.close() diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 3ea9e3db..9988f655 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -6,14 +6,14 @@ that require a live server live in test_dropin_e2e.py and are gated on SPARK_REMOTE. """ -import multiprocessing import pickle import time from concurrent.futures import ThreadPoolExecutor -from threading import Barrier, Event +from threading import Barrier import pytest +from _mock_connect_server import message_field, mock_connect_session, struct_schema from pyspark.sql import SparkSession, types as T @@ -2762,163 +2762,18 @@ def test_create_dataframe_width_mismatch_is_a_normal_exception(): # --------------------------------------------------------------------------- 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() @pytest.fixture -def schema_cache_analyzed(): +def mock_connect(): 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() + with mock_connect_session() as session: + yield session @pytest.mark.parametrize("first", ["schema", "columns", "dtypes"]) -def test_schema_cache_one_analysis_shared_by_metadata_accessors(schema_cache_analyzed, first): - spark, service = schema_cache_analyzed +def test_schema_cache_one_analysis_shared_by_metadata_accessors(mock_connect, first): + spark, service = mock_connect df = spark.range(10) assert service.calls == 0 getattr(df, first) @@ -2933,43 +2788,29 @@ def test_schema_cache_one_analysis_shared_by_metadata_accessors(schema_cache_ana 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) +def test_schema_cache_returned_schema_and_lists_are_independent(mock_connect): + spark, service = mock_connect df = spark.range(10) - original = df.schema - expected_json = original.jsonValue() + expected_json = df.schema.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")] + assert df.dtypes == [("id", "bigint")] assert service.calls == 1 -def test_schema_cache_derived_and_new_plans_resolve_independently(schema_cache_analyzed): - spark, service = schema_cache_analyzed +def test_schema_cache_derived_and_new_plans_resolve_independently(mock_connect): + spark, service = mock_connect df = spark.range(10) assert df.columns == ["id"] derived = df.selectExpr("id AS renamed") - service.data_type = schema_cache_schema("renamed") + service.data_type = struct_schema("renamed") assert derived.columns == ["renamed"] assert service.calls == 2 assert df.columns == ["id"] @@ -2986,9 +2827,163 @@ def test_schema_cache_derived_and_new_plans_resolve_independently(schema_cache_a assert df.columns == ["id"] +@pytest.mark.parametrize("transform", [ + pytest.param(lambda df: df.filter("id > 0"), id="filter"), + pytest.param(lambda df: df.limit(2), id="limit"), + pytest.param(lambda df: df.alias("source"), id="alias"), + pytest.param(lambda df: df.distinct(), id="distinct"), + pytest.param(lambda df: df.dropDuplicates(["id"]), id="dropDuplicates"), + pytest.param(lambda df: df.coalesce(1), id="coalesce"), + pytest.param(lambda df: df.repartition(2), id="repartition"), + pytest.param(lambda df: df.repartition(2, "id"), id="repartition-expression"), + pytest.param(lambda df: df.repartition("id"), id="repartition-column"), + pytest.param(lambda df: df.repartitionByRange(2, "id"), id="repartitionByRange"), + pytest.param(lambda df: df.repartitionById(2, "id"), id="repartitionById"), + pytest.param(lambda df: df.sort("id"), id="sort"), + pytest.param(lambda df: df.sortWithinPartitions("id"), id="sortWithinPartitions"), + pytest.param(lambda df: df.sample(fraction=0.5, seed=1), id="sample"), + pytest.param(lambda df: df.hint("broadcast"), id="hint"), +]) +def test_schema_cache_preserving_transforms_reuse_warm_schema(mock_connect, transform): + spark, service = mock_connect + df = spark.range(10) + expected = df.schema.jsonValue() + assert service.calls == 1 + # A new analysis would return different metadata, making accidental misses + # observable as well as countable. + service.data_type = struct_schema("unexpected") + derived = transform(df) + assert derived.schema.jsonValue() == expected + assert derived.limit(1).columns == ["id"] + assert service.calls == 1 + + +def test_schema_cache_random_split_and_observe_reuse_warm_schema(mock_connect): + from pyspark.sql import functions as F + + spark, service = mock_connect + df = spark.range(10) + assert df.columns == ["id"] + service.data_type = struct_schema("unexpected") + for split in df.randomSplit([1.0, 2.0], seed=1): + assert split.columns == ["id"] + observed = df.observe("rows", F.count("*").alias("count")) + assert observed.columns == ["id"] + assert service.calls == 1 + + +def test_schema_cache_cold_children_keep_independent_caches(mock_connect): + spark, service = mock_connect + df = spark.range(10) + before_lookup = df.limit(1) + service.block_response = True + with ThreadPoolExecutor(max_workers=1) as pool: + lookup = pool.submit(lambda: df.schema) + try: + for _ in range(500): + if service.calls: + break + time.sleep(0.01) + else: + raise AssertionError("Metadata request never reached the server") + # Constructing a child also runs with the GIL. It must not wait for + # the in-flight cache miss, or this thread cannot release the server. + during_lookup = df.filter("id > 0") + assert not lookup.done() + finally: + service.release_response() + assert lookup.result(timeout=10).fieldNames() == ["id"] + assert service.calls == 1 + # Neither child picks up a cache populated after the child was constructed. + for calls, child in enumerate((before_lookup, during_lookup), start=2): + service.data_type = struct_schema(f"child_{calls}") + assert child.columns == [f"child_{calls}"] + assert service.calls == calls + assert df.columns == ["id"] + + +@pytest.mark.parametrize("transform", [ + pytest.param(lambda df: df.selectExpr("id AS renamed"), id="select"), + pytest.param(lambda df: df.zipWithIndex(), id="zipWithIndex"), + # Reference PySpark does not propagate through offset, despite its unchanged + # output columns. Match its explicit propagation rules rather than guessing. + pytest.param(lambda df: df.offset(1), id="offset"), +]) +def test_schema_cache_other_transforms_analyze_independently(mock_connect, transform): + spark, service = mock_connect + df = spark.range(10) + assert df.columns == ["id"] + derived = transform(df) + assert service.calls == 1 + service.data_type = struct_schema("derived") + assert derived.columns == ["derived"] + assert service.calls == 2 + assert df.columns == ["id"] + + +@pytest.mark.parametrize("operation", [ + "union", "unionAll", "unionByName", "subtract", "exceptAll", "intersect", "intersectAll", +]) +def test_schema_cache_set_operations_require_equal_warm_schemas(mock_connect, operation): + spark, service = mock_connect + left, right = spark.range(10), spark.range(10) + assert left.columns == right.columns == ["id"] + assert service.calls == 2 + equal = getattr(left, operation)(right) + assert equal.columns == ["id"] + assert service.calls == 2 + + cold = spark.range(10) + missing = getattr(left, operation)(cold) + assert service.calls == 2 + service.data_type = struct_schema("resolved") + assert missing.columns == ["resolved"] + assert service.calls == 3 + + different = spark.range(10) + assert different.columns == ["resolved"] + assert service.calls == 4 + mismatched = getattr(left, operation)(different) + service.data_type = struct_schema("coerced") + assert mismatched.columns == ["coerced"] + assert service.calls == 5 + assert left.columns == ["id"] + + +def test_schema_cache_to_uses_supplied_struct_schema(mock_connect): + spark, service = mock_connect + df = spark.range(10) + target = T.StructType([T.StructField("converted", T.StringType(), True)]) + expected = target.jsonValue() + converted = df.to(target) + target.add("extra", T.LongType()) + assert converted.schema.jsonValue() == expected + assert converted.limit(1).dtypes == [("converted", "string")] + assert service.calls == 0 + assert df.columns == ["id"] + assert service.calls == 1 + + +@pytest.mark.parametrize("method", ["mapInPandas", "mapInArrow"]) +def test_schema_cache_map_uses_declared_struct_schema(mock_connect, method): + spark, service = mock_connect + target = T.StructType([ + T.StructField("output", T.StringType(), True, {"note": "declared"}), + ]) + mapped = getattr(spark.range(10), method)(lambda batches: batches, schema=target) + # The binding resolves the input columns while building the UDF. Reading its + # declared output schema must not issue another analysis request. + assert service.calls == 1 + service.data_type = struct_schema("unexpected") + assert mapped.schema.jsonValue() == target.jsonValue() + assert mapped.columns == ["output"] + assert mapped.dtypes == [("output", "string")] + assert service.calls == 1 + + @pytest.mark.parametrize("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 +def test_schema_cache_failures_are_not_cached(mock_connect, failure): + spark, service = mock_connect df = spark.range(10) service.failure = failure for count in (1, 2): @@ -3007,23 +3002,22 @@ def test_schema_cache_failures_are_not_cached(schema_cache_analyzed, failure): assert service.calls == 3 -@pytest.mark.parametrize("accessor", ["schema", "columns", "dtypes"]) -def test_schema_cache_non_struct_responses_are_not_cached(schema_cache_analyzed, accessor): - spark, service = schema_cache_analyzed +def test_schema_cache_non_struct_responses_are_not_cached(mock_connect): + spark, service = mock_connect df = spark.range(10) - service.data_type = schema_cache_message(7, b"") + service.data_type = message_field(7, b"") for count in (1, 2): with pytest.raises(Exception, match="Schema is not a struct type"): - getattr(df, accessor) + _ = df.schema assert service.calls == count - service.data_type = schema_cache_schema() + service.data_type = struct_schema() assert df.columns == ["id"] assert df.dtypes == [("id", "bigint")] assert service.calls == 3 -def test_schema_cache_concurrent_metadata_access_issues_one_analysis(schema_cache_analyzed): - spark, service = schema_cache_analyzed +def test_schema_cache_concurrent_metadata_access_issues_one_analysis(mock_connect): + spark, service = mock_connect df = spark.range(10) accessors = ("schema", "columns", "dtypes") * 3 start = Barrier(len(accessors) + 1) @@ -3065,8 +3059,8 @@ def read_metadata(accessor): @pytest.mark.parametrize("accessor", ["schema", "columns", "dtypes"]) -def test_schema_cache_metadata_wait_releases_gil(schema_cache_analyzed, accessor): - spark, service = schema_cache_analyzed +def test_schema_cache_metadata_wait_releases_gil(mock_connect, accessor): + spark, service = mock_connect df = spark.range(10) service.block_response = True From 1f9e52530cf242dc6fdc158d5e647147add93b1a Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 30 Oct 2026 09:24:21 -0700 Subject: [PATCH 9/9] [SPARK-60677][CONNECT] Release schema cache lock before copying schemas Clone the cached Arc under the mutex and copy the schema after releasing it, on both cache hits and successful analysis. This lets concurrent schema-preserving transformations snapshot the cache while metadata readers copy wide schemas. Move propagation details beside the cache helpers to keep schema() documentation focused on metadata access. Co-authored-by: Isaac --- crates/spark-connect/src/dataframe.rs | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/crates/spark-connect/src/dataframe.rs b/crates/spark-connect/src/dataframe.rs index 0802fe9a..4ce38983 100644 --- a/crates/spark-connect/src/dataframe.rs +++ b/crates/spark-connect/src/dataframe.rs @@ -210,6 +210,7 @@ impl DataFrame { Self::new_with_schema(session, plan, None) } + /// Seed declared struct output schemas without analyzing the plan. pub(crate) fn new_with_schema( session: SparkSession, plan: LogicalPlan, @@ -242,10 +243,13 @@ impl DataFrame { self.schema_cache.try_lock().ok()?.clone() } + /// Selected schema-preserving transformations snapshot the resolved schema, + /// mirroring Spark 4.2 `_cached_schema` propagation. Other plans resolve independently. fn with_cached_schema(&self, plan: LogicalPlan) -> Self { Self::new_with_cached_schema(self.session.clone(), plan, self.cached_schema()) } + /// Set operations propagate only exactly equal cached schemas from the same session. fn with_merged_schema(&self, other: &DataFrame, plan: LogicalPlan) -> Self { // A right-hand plan from another session is analyzed in our session; // its previously resolved schema is not evidence about that resolution. @@ -823,10 +827,6 @@ impl DataFrame { /// Get the schema of this DataFrame. /// /// The first successful lookup is cached and shared with exact clones. - /// Selected schema-preserving transformations snapshot an already resolved - /// schema; set operations do so only when both cached input schemas match - /// exactly in the same session. Declared struct output schemas also seed - /// the cache. Other plans resolve independently. /// Concurrent lookups on the same DataFrame or its clones wait for a single /// analysis RPC, so a stalled analysis blocks those metadata readers. /// Changes to external table or view metadata do not invalidate this cache; @@ -846,7 +846,9 @@ impl DataFrame { self.schema_cache.clear_poison(); poisoned.into_inner() }); - if let Some(schema) = cached.as_ref() { + if let Some(schema) = cached.clone() { + // Copy the schema after unlocking so propagation can snapshot it. + drop(cached); return Ok(schema.as_ref().clone()); } let request = self.build_analyze_request()?; @@ -861,8 +863,10 @@ impl DataFrame { if !matches!(&schema, DataType::Struct { .. }) { return Err(SparkError::connect_msg("Schema is not a struct type")); } - *cached = Some(Arc::new(schema.clone())); - Ok(schema) + let schema = Arc::new(schema); + *cached = Some(Arc::clone(&schema)); + drop(cached); + Ok(schema.as_ref().clone()) } else { Err(SparkError::connect_msg( "Schema analyze failed: no schema in response",