diff --git a/crates/pyspark-rs/src/dataframe.rs b/crates/pyspark-rs/src/dataframe.rs index 7101aadf..162bc96e 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,20 +1374,16 @@ impl PyDataFrame { /// Get the schema as a list of (name, type) tuples. #[getter] - fn dtypes(&self) -> PyResult> { - let schema = 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", - )), - } + fn dtypes(&self, py: Python<'_>) -> PyResult> { + 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 ea0e2c62..4ce38983 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,10 @@ use crate::udf::CommonInlineUserDefinedFunctionExpression; pub struct DataFrame { pub(crate) session: SparkSession, pub(crate) plan: 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. @@ -203,10 +207,61 @@ 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) + } + + /// Seed declared struct output schemas without analyzing the plan. + 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(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() + } + + /// 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. + 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. @@ -231,7 +286,7 @@ impl DataFrame { input: Arc::clone(&self.plan), condition, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Alias for filter(). @@ -300,7 +355,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. @@ -329,7 +384,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. @@ -345,7 +400,7 @@ impl DataFrame { column_names: cols, within_watermark: false, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Sort rows. @@ -355,7 +410,7 @@ impl DataFrame { order: columns, is_global: true, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Alias for sort(). @@ -451,7 +506,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. @@ -471,7 +526,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. @@ -484,7 +539,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. @@ -497,7 +552,7 @@ impl DataFrame { by_name: false, allow_missing_columns: false, }; - DataFrame::new(self.session.clone(), plan) + self.with_merged_schema(other, plan) } /// Repartition. @@ -507,7 +562,7 @@ impl DataFrame { num_partitions, shuffle: true, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Coalesce. @@ -517,7 +572,7 @@ impl DataFrame { num_partitions, shuffle: false, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Add a hint. @@ -531,7 +586,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). @@ -556,7 +611,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`). @@ -586,12 +641,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`). @@ -629,7 +685,7 @@ impl DataFrame { with_replacement, seed, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Group by columns for aggregation. @@ -769,14 +825,48 @@ impl DataFrame { } /// Get the schema of this DataFrame. + /// + /// The first successful lookup is cached and shared with exact clones. + /// 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; + /// 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() + // 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.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()?; let response = block_on(self.session.client().analyze_plan(request))?; if let Some(proto::analyze_plan_response::Result::Schema(schema)) = response.result { - Ok(DataType::from_proto(&schema.schema.ok_or_else(|| { - SparkError::connect_msg("Schema is missing") - })?)?) + let schema = DataType::from_proto( + &schema + .schema + .ok_or_else(|| SparkError::connect_msg("Schema is missing"))?, + )?; + if !matches!(&schema, DataType::Struct { .. }) { + return Err(SparkError::connect_msg("Schema is not a struct type")); + } + 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", @@ -806,11 +896,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. @@ -1089,7 +1178,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. @@ -1104,7 +1193,7 @@ impl DataFrame { num_partitions, expressions: columns, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Alias for to_df(). @@ -1161,7 +1250,7 @@ impl DataFrame { name: name.to_string(), exprs, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Get stat functions. @@ -1356,7 +1445,7 @@ impl DataFrame { order: columns, is_global: false, }; - DataFrame::new(self.session.clone(), plan) + self.with_cached_schema(plan) } /// Drop duplicates within a watermark. @@ -1400,7 +1489,7 @@ impl DataFrame { with_replacement: false, seed, }; - results.push(DataFrame::new(self.session.clone(), plan)); + results.push(self.with_cached_schema(plan)); cumulative = upper; } @@ -1441,17 +1530,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, @@ -1539,7 +1624,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. @@ -1552,7 +1637,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. @@ -1565,7 +1650,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). @@ -1699,9 +1784,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). @@ -1746,10 +1831,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 +2932,168 @@ mod cache_tests { use super::*; use prost::Message; + #[test] + fn schema_cache_shares_exact_clones_but_not_changed_plans_or_new_sources() { + 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(Arc::new(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.select(["id"]), + df.zip(&df).unwrap(), + spark.range(3).unwrap(), + ] { + assert!(!Arc::ptr_eq(&df.schema_cache, &derived.schema_cache)); + assert!(derived.schema_cache.lock().unwrap().is_none()); + } + } + + #[test] + fn schema_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 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": "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] + 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()); + // 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()); + let schema = DataType::Struct { fields: vec![] }; + *cached = Some(Arc::new(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); + } + #[test] fn cache_default_is_memory_and_disk_deser() { let sl = DataFrame::memory_and_disk_deser(); 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 67ccaa4d..9988f655 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -7,10 +7,14 @@ """ import pickle +import time +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier import pytest -from pyspark.sql import types as T +from _mock_connect_server import message_field, mock_connect_session, struct_schema +from pyspark.sql import SparkSession, types as T # --------------------------------------------------------------------------- types @@ -2755,3 +2759,329 @@ def test_create_dataframe_width_mismatch_is_a_normal_exception(): else: raise AssertionError("expected a width validation error") """], check=True) + + +# --------------------------------------------------------------------------- schema cache + + +@pytest.fixture +def mock_connect(): + pytest.importorskip("grpc") + 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(mock_connect, first): + spark, service = mock_connect + 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(mock_connect): + spark, service = mock_connect + df = spark.range(10) + expected_json = df.schema.jsonValue() + modified = df.schema + modified.add("extra", T.StringType()) + assert len(modified.fields) == 2 + columns = df.columns + columns.append("extra") + dtypes = df.dtypes + dtypes.clear() + assert df.schema.jsonValue() == expected_json + assert df.columns == ["id"] + assert df.dtypes == [("id", "bigint")] + assert service.calls == 1 + + +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 = struct_schema("renamed") + assert derived.columns == ["renamed"] + assert service.calls == 2 + assert df.columns == ["id"] + # Rebuilding the same source must observe updated server metadata, while the + # original immutable DataFrame retains its already-resolved schema. + fresh = spark.range(10) + assert fresh.columns == ["renamed"] + other_session = spark.newSession() + try: + assert other_session.range(10).columns == ["renamed"] + finally: + other_session.stop() + assert service.calls == 4 + assert df.columns == ["id"] + + +@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(mock_connect, failure): + spark, service = mock_connect + 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_non_struct_responses_are_not_cached(mock_connect): + spark, service = mock_connect + df = spark.range(10) + service.data_type = message_field(7, b"") + for count in (1, 2): + with pytest.raises(Exception, match="Schema is not a struct type"): + _ = df.schema + assert service.calls == count + 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(mock_connect): + spark, service = mock_connect + 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 + + +@pytest.mark.parametrize("accessor", ["schema", "columns", "dtypes"]) +def test_schema_cache_metadata_wait_releases_gil(mock_connect, accessor): + spark, service = mock_connect + df = spark.range(10) + service.block_response = True + + def allow_response_after_request(): + for _ in range(500): + if service.calls: + service.release_response() + return + 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 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