Freeze PyO3 wrappers & introduce interior mutability to avoid PyO3 borrow errors (#1253)

* Refactor schema, config, dataframe, and expression classes to use RwLock and Mutex for interior mutability

* Add error handling to CaseBuilder methods to preserve builder state

* Refactor to use parking_lot for interior mutability in schema, config, dataframe, and conditional expression modules

* Add concurrency tests for SqlSchema, Config, and DataFrame

* Add tests for CaseBuilder to ensure builder state is preserved on success

* Add test for independent handles in CaseBuilder to verify behavior

* Fix CaseBuilder to preserve state correctly in when() method

* Refactor to use named constant for boolean literals in test_expr.py

* fix ruff errors

* Refactor to introduce type aliases for cached batches in dataframe.rs

* Cherry pick from #1252

* Add most expr - cherry pick from #1252

* Add source root - cherry pick #1252

* Fix license comment formatting in config.rs

* Refactor caching logic to use a local variable for IPython environment check

* Add test for ensuring exposed pyclasses default to frozen

* Add PyO3 class mutability guidelines reference to contributor guide

* Mark boolean expression classes as frozen for immutability

* Refactor PyCaseBuilder methods to eliminate redundant take/store logic

* Refactor PyConfig methods to improve readability by encapsulating configuration reads

* Resolve patch apply conflicts for CaseBuilder concurrency improvements

- Added CaseBuilderHandle guard that keeps the underlying CaseBuilder alive while holding the mutex and restores it on drop
- Updated when, otherwise, and end methods to operate through the guard and consume the builder explicitly
- This prevents transient None states during concurrent access and improves thread safety

* Resolve Config optimization conflicts for improved read/write concurrency

- Released Config read guard before converting values to Python objects in get and get_all
- Ensures locks are held only while collecting scalar entries, not during expensive Python object conversion
- Added regression test that runs Config.get_all and Config.set concurrently to guard against read/write contention regressions
- Improves overall performance by reducing lock contention in multi-threaded scenarios

* Refactor PyConfig get methods for improved readability and performance

* Refactor test_expr.py to replace positional boolean literals with named constants for improved linting compliance

* fix ruff errors

* Add license header to test_pyclass_frozen.py for compliance

* Alternate approach to case expression

* Replace case builter with keeping the expressions and then applying as required

* Update unit tests

* Refactor case and when functions to utilize PyCaseBuilder for improved clarity and functionality

* Update src/expr/conditional_expr.rs

---------

Co-authored-by: ntjohnson1 <24689722+ntjohnson1@users.noreply.github.com>
Co-authored-by: Tim Saucer <timsaucer@gmail.com>

Author: 3 people

Cargo.lock

@@ -51,6 +51,7 @@ futures = "0.3"
 object_store = { version = "0.12.3", features = ["aws", "gcp", "azure", "http"] }
 url = "2"
 log = "0.4.27"
+parking_lot = "0.12"
 
 [build-dependencies]
 prost-types = "0.13.1" # keep in line with `datafusion-substrait`

Cargo.toml

@@ -137,6 +137,67 @@ and you want to create a sharable FFI counterpart, you could write:
     let my_provider = MyTableProvider::default();
     let ffi_provider = FFI_TableProvider::new(Arc::new(my_provider), false, None);
 
+.. _ffi_pyclass_mutability:
+
+PyO3 class mutability guidelines
+--------------------------------
+
+PyO3 bindings should present immutable wrappers whenever a struct stores shared or
+interior-mutable state. In practice this means that any ``#[pyclass]`` containing an
+``Arc<RwLock<_>>`` or similar synchronized primitive must opt into ``#[pyclass(frozen)]``
+unless there is a compelling reason not to.
+
+The :mod:`datafusion` configuration helpers illustrate the preferred pattern. The
+``PyConfig`` class in :file:`src/config.rs` stores an ``Arc<RwLock<ConfigOptions>>`` and is
+explicitly frozen so callers interact with configuration state through provided methods
+instead of mutating the container directly:
+
+.. code-block:: rust
+
+    #[pyclass(name = "Config", module = "datafusion", subclass, frozen)]
+    #[derive(Clone)]
+    pub(crate) struct PyConfig {
+        config: Arc<RwLock<ConfigOptions>>,
+    }
+
+The same approach applies to execution contexts. ``PySessionContext`` in
+:file:`src/context.rs` stays frozen even though it shares mutable state internally via
+``SessionContext``. This ensures PyO3 tracks borrows correctly while Python-facing APIs
+clone the inner ``SessionContext`` or return new wrappers instead of mutating the
+existing instance in place:
+
+.. code-block:: rust
+
+    #[pyclass(frozen, name = "SessionContext", module = "datafusion", subclass)]
+    #[derive(Clone)]
+    pub struct PySessionContext {
+        pub ctx: SessionContext,
+    }
+
+Occasionally a type must remain mutable—for example when PyO3 attribute setters need to
+update fields directly. In these rare cases add an inline justification so reviewers and
+future contributors understand why ``frozen`` is unsafe to enable. ``DataTypeMap`` in
+:file:`src/common/data_type.rs` includes such a comment because PyO3 still needs to track
+field updates:
+
+.. code-block:: rust
+
+    // TODO: This looks like this needs pyo3 tracking so leaving unfrozen for now
+    #[derive(Debug, Clone)]
+    #[pyclass(name = "DataTypeMap", module = "datafusion.common", subclass)]
+    pub struct DataTypeMap {
+        #[pyo3(get, set)]
+        pub arrow_type: PyDataType,
+        #[pyo3(get, set)]
+        pub python_type: PythonType,
+        #[pyo3(get, set)]
+        pub sql_type: SqlType,
+    }
+
+When reviewers encounter a mutable ``#[pyclass]`` without a comment, they should request
+an explanation or ask that ``frozen`` be added. Keeping these wrappers frozen by default
+helps avoid subtle bugs stemming from PyO3's interior mutability tracking.
+
 If you were interfacing with a library that provided the above ``FFI_TableProvider`` and
 you needed to turn it back into an ``TableProvider``, you can turn it into a
 ``ForeignTableProvider`` with implements the ``TableProvider`` trait.

docs/source/contributor-guide/ffi.rst

@@ -26,6 +26,10 @@ We welcome and encourage contributions of all kinds, such as:
 In addition to submitting new PRs, we have a healthy tradition of community members reviewing each other’s PRs.
 Doing so is a great way to help the community as well as get more familiar with Rust and the relevant codebases.
 
+Before opening a pull request that touches PyO3 bindings, please review the
+:ref:`PyO3 class mutability guidelines <ffi_pyclass_mutability>` so you can flag missing
+``#[pyclass(frozen)]`` annotations during development and review.
+
 How to develop
 --------------
 

docs/source/contributor-guide/introduction.rst

@@ -107,6 +107,9 @@ convention = "google"
 [tool.ruff.lint.pycodestyle]
 max-doc-length = 88
 
+[tool.ruff.lint.flake8-boolean-trap]
+extend-allowed-calls = ["lit", "datafusion.lit"]
+
 # Disable docstring checking for these directories
 [tool.ruff.lint.per-file-ignores]
 "python/tests/*" = [

pyproject.toml

@@ -0,0 +1,126 @@
+# 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.
+
+from __future__ import annotations
+
+from concurrent.futures import ThreadPoolExecutor
+
+import pyarrow as pa
+from datafusion import Config, SessionContext, col, lit
+from datafusion import functions as f
+from datafusion.common import SqlSchema
+
+
+def _run_in_threads(fn, count: int = 8) -> None:
+    with ThreadPoolExecutor(max_workers=count) as executor:
+        futures = [executor.submit(fn, i) for i in range(count)]
+        for future in futures:
+            # Propagate any exception raised in the worker thread.
+            future.result()
+
+
+def test_concurrent_access_to_shared_structures() -> None:
+    """Exercise SqlSchema, Config, and DataFrame concurrently."""
+
+    schema = SqlSchema("concurrency")
+    config = Config()
+    ctx = SessionContext()
+
+    batch = pa.record_batch([pa.array([1, 2, 3], type=pa.int32())], names=["value"])
+    df = ctx.create_dataframe([[batch]])
+
+    config_key = "datafusion.execution.batch_size"
+    expected_rows = batch.num_rows
+
+    def worker(index: int) -> None:
+        schema.name = f"concurrency-{index}"
+        assert schema.name.startswith("concurrency-")
+        # Exercise getters that use internal locks.
+        assert isinstance(schema.tables, list)
+        assert isinstance(schema.views, list)
+        assert isinstance(schema.functions, list)
+
+        config.set(config_key, str(1024 + index))
+        assert config.get(config_key) is not None
+        # Access the full config map to stress lock usage.
+        assert config_key in config.get_all()
+
+        batches = df.collect()
+        assert sum(batch.num_rows for batch in batches) == expected_rows
+
+    _run_in_threads(worker, count=12)
+
+
+def test_config_set_during_get_all() -> None:
+    """Ensure config writes proceed while another thread reads all entries."""
+
+    config = Config()
+    key = "datafusion.execution.batch_size"
+
+    def reader() -> None:
+        for _ in range(200):
+            # get_all should not hold the lock while converting to Python objects
+            config.get_all()
+
+    def writer() -> None:
+        for index in range(200):
+            config.set(key, str(1024 + index))
+
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        reader_future = executor.submit(reader)
+        writer_future = executor.submit(writer)
+        reader_future.result(timeout=10)
+        writer_future.result(timeout=10)
+
+    assert config.get(key) is not None
+
+
+def test_case_builder_reuse_from_multiple_threads() -> None:
+    """Ensure the case builder can be safely reused across threads."""
+
+    ctx = SessionContext()
+    values = pa.array([0, 1, 2, 3, 4], type=pa.int32())
+    df = ctx.create_dataframe([[pa.record_batch([values], names=["value"])]])
+
+    base_builder = f.case(col("value"))
+
+    def add_case(i: int) -> None:
+        nonlocal base_builder
+        base_builder = base_builder.when(lit(i), lit(f"value-{i}"))
+
+    _run_in_threads(add_case, count=8)
+
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        otherwise_future = executor.submit(base_builder.otherwise, lit("default"))
+        case_expr = otherwise_future.result()
+
+    result = df.select(case_expr.alias("label")).collect()
+    assert sum(batch.num_rows for batch in result) == len(values)
+
+    predicate_builder = f.when(col("value") == lit(0), lit("zero"))
+
+    def add_predicate(i: int) -> None:
+        predicate_builder.when(col("value") == lit(i + 1), lit(f"value-{i + 1}"))
+
+    _run_in_threads(add_predicate, count=4)
+
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        end_future = executor.submit(predicate_builder.end)
+        predicate_expr = end_future.result()
+
+    result = df.select(predicate_expr.alias("label")).collect()
+    assert sum(batch.num_rows for batch in result) == len(values)

python/tests/test_concurrency.py

@@ -16,6 +16,7 @@
 # under the License.
 
 import re
+from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, timezone
 
 import pyarrow as pa
@@ -200,6 +201,98 @@ def traverse_logical_plan(plan):
     assert not variant.negated()
 
 
+def test_case_builder_error_preserves_builder_state():
+    case_builder = functions.when(lit(True), lit(1))
+
+    with pytest.raises(Exception) as exc_info:
+        _ = case_builder.otherwise(lit("bad"))
+
+    err_msg = str(exc_info.value)
+    assert "multiple data types" in err_msg
+    assert "CaseBuilder has already been consumed" not in err_msg
+
+    _ = case_builder.end()
+
+    err_msg = str(exc_info.value)
+    assert "multiple data types" in err_msg
+    assert "CaseBuilder has already been consumed" not in err_msg
+
+
+def test_case_builder_success_preserves_builder_state():
+    ctx = SessionContext()
+    df = ctx.from_pydict({"flag": [False]}, name="tbl")
+
+    case_builder = functions.when(col("flag"), lit("true"))
+
+    expr_default_one = case_builder.otherwise(lit("default-1")).alias("result")
+    result_one = df.select(expr_default_one).collect()
+    assert result_one[0].column(0).to_pylist() == ["default-1"]
+
+    expr_default_two = case_builder.otherwise(lit("default-2")).alias("result")
+    result_two = df.select(expr_default_two).collect()
+    assert result_two[0].column(0).to_pylist() == ["default-2"]
+
+    expr_end_one = case_builder.end().alias("result")
+    end_one = df.select(expr_end_one).collect()
+    assert end_one[0].column(0).to_pylist() == [None]
+
+
+def test_case_builder_when_handles_are_independent():
+    ctx = SessionContext()
+    df = ctx.from_pydict(
+        {
+            "flag": [True, False, False, False],
+            "value": [1, 15, 25, 5],
+        },
+        name="tbl",
+    )
+
+    base_builder = functions.when(col("flag"), lit("flag-true"))
+
+    first_builder = base_builder.when(col("value") > lit(10), lit("gt10"))
+    second_builder = base_builder.when(col("value") > lit(20), lit("gt20"))
+
+    first_builder = first_builder.when(lit(True), lit("final-one"))
+
+    expr_first = first_builder.otherwise(lit("fallback-one")).alias("first")
+    expr_second = second_builder.otherwise(lit("fallback-two")).alias("second")
+
+    result = df.select(expr_first, expr_second).collect()[0]
+
+    assert result.column(0).to_pylist() == [
+        "flag-true",
+        "gt10",
+        "gt10",
+        "final-one",
+    ]
+    assert result.column(1).to_pylist() == [
+        "flag-true",
+        "fallback-two",
+        "gt20",
+        "fallback-two",
+    ]
+
+
+def test_case_builder_when_thread_safe():
+    case_builder = functions.when(lit(True), lit(1))
+
+    def build_expr(value: int) -> bool:
+        builder = case_builder.when(lit(True), lit(value))
+        builder.otherwise(lit(value))
+        return True
+
+    with ThreadPoolExecutor(max_workers=8) as executor:
+        futures = [executor.submit(build_expr, idx) for idx in range(16)]
+        results = [future.result() for future in futures]
+
+    assert all(results)
+
+    # Ensure the shared builder remains usable after concurrent `when` calls.
+    follow_up_builder = case_builder.when(lit(True), lit(42))
+    assert isinstance(follow_up_builder, type(case_builder))
+    follow_up_builder.otherwise(lit(7))
+
+
 def test_expr_getitem() -> None:
     ctx = SessionContext()
     data = {