[dbtsai]

What changes were proposed in this pull request?

Cache resolved schemas across metadata reads and exact clones; failures remain retryable. Propagate snapshots through selected schema-preserving transformations, following Spark 4.2; set operations require exactly equal schemas from the same session. to(schema) and stateless partition/group pandas/Arrow UDFs use declared struct schemas.

Propagation never analyzes parents or waits for in-flight analysis. Other plans resolve independently. Python metadata waits release the GIL, and the shared cache recovers defensively after Rust panics.

Why are the changes needed?

Avoid redundant analysis RPCs while allowing Python threads to progress during metadata waits.

Does this PR introduce any user-facing change?

Yes. Metadata reads reuse resolved schemas. Concurrent readers share one analysis request; a stalled request blocks those readers. External table/view changes can leave cached metadata stale; recreate the DataFrame from its source to refresh it.

How was this patch tested?

• 398 Rust library tests and 318 CI-selected offline Python tests passed. Regression checks reject the prior build and a blocking propagation mutation.
• 19 live cases matched reference PySpark on Spark 4.2, comparing schemas, metadata, columns, dtypes, and collected rows.
• Release/WASM integration compilation, formatting, and wheel build passed. Advisory Clippy reports the same 923 diagnostics as the base.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: OpenAI Codex (GPT-6)

[HyukjinKwon]

Review summary

The production change holds together. DataFrame::schema() keeps the cache lock from the cache check through the analyze RPC and the store, and stores only a successfully decoded schema, so failures stay retryable. Poison recovery is safe because the slot only ever holds None or a complete value. zip now goes through DataFrame::new, so every derived plan starts with an empty cache while exact clones share one. The schema, columns, and dtypes getters all release the GIL around the core call. The one finding is in the new offline tests: the concurrent-access test does not hold the first AnalyzePlan open, so it checks the single-flight lock only probabilistically. The existing response gate in the fixture would make it deterministic.

Findings

1 total: 0 P0, 0 P1, 0 P2, 1 P3.

Nit (P3)

• Concurrent schema-cache test does not force the contended interleaving — python/tests/test_dropin_offline.py:3000 — see inline.

[HyukjinKwon]

Nit (P3): The mock's AnalyzePlan handler replies immediately here because service.block_response is never set. Whether threads 2..8 reach DataFrame::schema() while the first RPC is still in flight is therefore left to the scheduler. Suppose a future change dropped the guard before block_on(analyze_plan) and re-checked the cache afterwards. This test could still see service.calls == 1 and pass, so its signal for the single-flight guarantee, the main reason for holding the mutex across the RPC, is only probabilistic. The fixture already has the gate that would make it deterministic, as used by test_schema_cache_metadata_wait_releases_gil. Set block_response = True, start the 8 reads, wait until service.calls is non-zero and give the other readers a moment to block on the cache, then call release_response() and assert one call.

Verification:

• Regression: With the first analysis response held until several readers have entered the schema accessor, concurrent schema reads on one DataFrame produce exactly one analysis request on the server and equal, distinct schema objects.
• Behavior: The gated test stays bounded: if the response is never released, or a reader keeps the GIL, the test fails within its timeouts and does not hang.

[dbtsai]

Fixed in d0183ef. The test now starts nine mixed schema/columns/dtypes readers behind a barrier, holds the first response open, and releases the gate in finally. It checks one analysis RPC and equal, distinct schema objects, with bounded waits.

I also tested its sensitivity against a build that deliberately drops the cache lock before analysis: the revised test failed in all three runs, observing nine RPCs each time. The fixed candidate passes all 290 offline Python tests and three additional runs of the contention test.

[viirya]

The core change looks right to me. schema_cache is private, and zip was the only struct-literal construction outside DataFrame::new; it now goes through new. Nothing mutates plan or session after construction. So one cache always corresponds to one (session, plan) pair. The clones that share it (persist/unpersist, GroupedData, asTable) all keep the same schema. Every pyspark-rs caller of schema()/columns()/input_columns() now runs inside py.detach, so no thread can block on the mutex while holding the GIL. The struct check matches the assert isinstance(schema, StructType) in the reference client.schema().

A few things I'd like to see addressed, mostly on the scope of the description and the tests:

• The description says DataFrames reuse resolved schemas "as reference PySpark does", but this only partly matches. Since 4.0 the reference also passes _cached_schema on to schema-preserving derivations: alias, filter, limit, sort, sample, hint, distinct, and union* when the schemas match (via _merge_cached_schema). createDataFrame(..., schema) and to(schema) also prefill the cache. Resolving every new plan independently is a reasonable, more conservative first step. Could the description say so, and could we leave propagation as a follow-up? One of the Rust tests currently asserts the opposite of that follow-up (see inline).
• Holding the lock across the RPC gives schema() single-flight semantics the reference doesn't have. If one AnalyzePlan stalls, every other reader of that DataFrame waits behind it instead of sending its own request. I think that trade-off is fine, but it's worth one sentence in the doc comment.
• The offline mock server is useful well beyond this feature. It's about 120 lines: hand-rolled wire encoders with hard-coded field numbers, plus a gRPC server in a subprocess. It's also inlined into an already large test file under schema_cache_* names. Could it move to its own helper module (e.g. tests/_mock_connect_server.py) with neutral names, so the next transport-level offline test can reuse it?
• Some tests check properties the types already guarantee. Inline notes below cover which ones I think can go.

[viirya]

Since the lock is held across analyze_plan, a slow or stalled analysis also blocks every other reader of this DataFrame and its clones, including GroupedData.input_columns(). Reference PySpark doesn't synchronize here at all, so this is new behavior on our side. Could the doc comment mention it?

[dbtsai]

Addressed in ed7efa6.

Added this trade-off to schema() rustdoc: concurrent readers of the DataFrame and its exact clones wait for one analysis RPC, so a stalled analysis blocks those metadata readers. Propagation itself uses a nonblocking snapshot and skips an in-flight lookup.

[viirya]

Every other mutex in the crate uses .lock().unwrap(). A short note on why this one needs special handling would help the next reader: it's the only lock held across a blocking, fallible call. Otherwise it looks like an inconsistency. Also, I don't think a panic while holding this lock is reachable from the Python side. The test can only trigger one by calling the sync API inside a tokio runtime. If keeping it is mainly defensive, that's fine, but the comment should say so.

[dbtsai]

Addressed in ed7efa6.

Clarified that this is defensive recovery for Rust callers: the lock spans a blocking RPC, and calling the sync API inside a Tokio runtime can panic there. The slot contains only None or a complete schema. The test also verifies recovery clears the poisoned flag.

[viirya]

With this check, schema() only ever returns DataType::Struct, so the _ => arms in columns() (line 842), dtypes() (line 1483), and pyspark-rs dtypes (crates/pyspark-rs/src/dataframe.rs:1387) can't be reached anymore. Could we drop them (e.g. let DataType::Struct { fields } = self.schema()? else { unreachable!() }) so the struct invariant is enforced in one place?

[dbtsai]

Addressed in ed7efa6.

Removed the duplicate non-struct error branches from Rust columns/dtypes and the Python dtypes binding. These now destructure the Struct invariant enforced by schema(); declared output schemas are also filtered to Struct before seeding.

[viirya]

Asserting that limit/select start with an empty cache ties the test to the current choice, not to a correctness requirement. Reference PySpark does carry the cached schema through limit (and filter, sort, alias, ...). For select and zip the assertion makes sense, since the schema can change. For limit, I'd drop it so propagating schemas like the reference doesn't need to start by deleting a test.

[dbtsai]

Addressed in ed7efa6.

Removed the limit-empty assertion and implemented propagation here. Warm limit/filter/sort/alias and the other selected Spark 4.2 transformations snapshot the resolved schema. Cold children keep independent caches. Tests cover both cases and ensure construction never waits for an in-flight parent analysis.

[viirya]

DataType and its fields own their data (String, HashMap, serde_json::Value), so clone() is always a deep copy. The test checks Clone, not the cache. I think it can be removed.

[dbtsai]

Addressed in ed7efa6.

Removed the Rust deep-clone test. The remaining tests focus on cache sharing, construction-time snapshots, exact schema equality for set operations, declared outputs, and panic recovery.

[viirya]

Same idea on the Python side: data_type_to_py builds new Python objects on every getter call, so returned values are independent by construction. This test would only start catching something if we cached Python objects, which this PR doesn't. Could it be trimmed or removed?

[dbtsai]

Addressed in ed7efa6.

Trimmed this to a small public behavior check that mutating a returned schema or metadata list does not affect later reads. Removed the nested mutation assertions and helper encoding for that case.

[viirya]

The struct check lives only in DataFrame::schema(), and columns/dtypes both go through it. So running this once per accessor (3 cases), on top of failures_are_not_cached, mostly repeats the same path. One case would cover it.

[dbtsai]

Addressed in ed7efa6.

Reduced non-struct response coverage to one schema accessor case. It still verifies that repeated invalid responses are not cached and that a later valid response succeeds.

[viirya]

This asserts ReleaseSession behavior from a fixture teardown, which isn't part of the schema-cache change. When it fails, the failure gets reported against whichever schema-cache test was running. Could we drop it, or move it into a test of its own?

[dbtsai]

Addressed in ed7efa6.

Removed the ReleaseSession count and teardown assertion. Teardown still calls spark.stop() and always joins or terminates the mock-server process.

[dbtsai]

Implemented propagation in ed7efa6, along with the documentation and test cleanup from this review.

Selected unary transformations snapshot already-resolved schemas. Set operations reuse them only when both cached schemas match exactly in the same session. to(schema) and stateless partition/group pandas/Arrow UDFs use declared struct outputs. Propagation does not resolve parents or wait on an in-flight analysis. The mock service now lives in tests/_mock_connect_server.py with neutral names.

One reference detail: pinned Spark 4.2's createDataFrame constructs DataFrame(plan, self) without seeding _cached_schema, including with an explicit StructType. I kept that behavior.

Validation: 398 Rust library tests and 318 CI-selected offline Python tests pass. Nineteen live cases match reference PySpark on Spark 4.2 for schemas, metadata, columns, dtypes, and rows. The prior build fails all 16 selected propagation cases; a deliberately blocking snapshot fails the bounded lock-wait test. Release/WASM integration targets compile; Clippy has the same 923 diagnostics as the base. The PR description now reflects this scope and validation.

[viirya]

Thanks for taking the propagation on in this PR. I checked it against every _cached_schema assignment in v4.2.0 connect/dataframe.py and connect/group.py, and the set matches:

• The unary operations are alias, coalesce, repartition*, dropDuplicates, distinct, filter, limit, sort*, sample, hint, randomSplit, and observe.
• The set operations go through the exact-equality merge.
• to and the stateless map/group/cogroup pandas/Arrow UDFs are seeded from the declared struct.
• dropDuplicatesWithinWatermark, offset, and the stateful group APIs stay unpropagated, which also matches.

The declared UDF schema comes from python_udf.output_type, which is the type actually sent to the server, so the seeded value agrees with what analysis would return. Using try_lock for the snapshot is the right call, since transformations run with the GIL held. with_merged_schema drops each guard before taking the next, so self-unions are safe.

You're right about createDataFrame. I was reading master rather than v4.2.0, where it doesn't seed _cached_schema. Keeping that behavior is correct.

All the earlier inline points are resolved. Two optional nits remain:

• On a cache hit, schema() deep-clones the DataType while still holding the lock. A child built on another thread during that window fails try_lock, starts with an empty cache, and later issues an analysis it didn't need. The result is still correct, but propagation becomes timing-dependent. Cloning the Arc under the guard and doing the deep clone after the guard is dropped would shrink that window to almost nothing (inline).
• The schema() rustdoc now covers caching, propagation rules, blocking, and staleness. The propagation rules might read better next to with_cached_schema/with_merged_schema, or as a one-line "mirrors Spark 4.2 _cached_schema propagation". Not blocking.

LGTM otherwise.

[viirya]

Small one: the deep clone here happens while the guard is held. A concurrent with_cached_schema on another thread then fails try_lock and builds a child with an empty cache. Something like this keeps the critical section to a refcount bump:

if let Some(schema) = cached.clone() {
    drop(cached);
    return Ok(schema.as_ref().clone());
}

[dbtsai]

Thanks, fixed in 1f9e525. Both cache hits and successful analysis now drop the guard before deep-cloning the schema. I also moved the propagation notes beside the cache helpers. Merged as 617fa148.