Boring fixes for frozen

Author: ntjohnson1

src/config.rs

@@ -18,38 +18,39 @@
 use pyo3::prelude::*;
 use pyo3::types::*;
 
+use std::sync::{Arc, Mutex};
+
 use datafusion::config::ConfigOptions;
 
 use crate::errors::PyDataFusionResult;
 use crate::utils::py_obj_to_scalar_value;
 
-// TODO: Not frozen because set needs access
-#[pyclass(name = "Config", module = "datafusion", subclass)]
+#[pyclass(frozen, name = "Config", module = "datafusion", subclass)]
 #[derive(Clone)]
 pub(crate) struct PyConfig {
-    config: ConfigOptions,
+    config: Arc<Mutex<ConfigOptions>>,
 }
 
 #[pymethods]
 impl PyConfig {
     #[new]
     fn py_new() -> Self {
         Self {
-            config: ConfigOptions::new(),
+            config: Arc::new(Mutex::new(ConfigOptions::new())),
         }
     }
 
     /// Get configurations from environment variables
     #[staticmethod]
     pub fn from_env() -> PyDataFusionResult<Self> {
         Ok(Self {
-            config: ConfigOptions::from_env()?,
+            config: Arc::new(Mutex::new(ConfigOptions::from_env()?)),
         })
     }
 
     /// Get a configuration option
     pub fn get<'py>(&self, key: &str, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        let options = self.config.to_owned();
+        let options = self.config.lock().unwrap();
         for entry in options.entries() {
             if entry.key == key {
                 return Ok(entry.value.into_pyobject(py)?);
@@ -59,16 +60,17 @@ impl PyConfig {
     }
 
     /// Set a configuration option
-    pub fn set(&mut self, key: &str, value: PyObject, py: Python) -> PyDataFusionResult<()> {
+    pub fn set(&self, key: &str, value: PyObject, py: Python) -> PyDataFusionResult<()> {
         let scalar_value = py_obj_to_scalar_value(py, value)?;
-        self.config.set(key, scalar_value.to_string().as_str())?;
+        let mut config = self.config.lock().unwrap();
+        config.set(key, scalar_value.to_string().as_str())?;
         Ok(())
     }
 
     /// Get all configuration options
     pub fn get_all(&self, py: Python) -> PyResult<PyObject> {
         let dict = PyDict::new(py);
-        let options = self.config.to_owned();
+        let options = self.config.lock().unwrap();
         for entry in options.entries() {
             dict.set_item(entry.key, entry.value.clone().into_pyobject(py)?)?;
         }

src/dataframe.rs

@@ -17,7 +17,7 @@
 
 use std::collections::HashMap;
 use std::ffi::CString;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 
 use arrow::array::{new_null_array, RecordBatch, RecordBatchIterator, RecordBatchReader};
 use arrow::compute::can_cast_types;
@@ -284,31 +284,31 @@ impl PyParquetColumnOptions {
 /// A PyDataFrame is a representation of a logical plan and an API to compose statements.
 /// Use it to build a plan and `.collect()` to execute the plan and collect the result.
 /// The actual execution of a plan runs natively on Rust and Arrow on a multi-threaded environment.
-// TODO: Not frozen because batches don't currently handle interior mutability
-#[pyclass(name = "DataFrame", module = "datafusion", subclass)]
+#[pyclass(frozen, name = "DataFrame", module = "datafusion", subclass)]
 #[derive(Clone)]
 pub struct PyDataFrame {
     df: Arc<DataFrame>,
 
     // In IPython environment cache batches between __repr__ and _repr_html_ calls.
-    batches: Option<(Vec<RecordBatch>, bool)>,
+    #[allow(clippy::type_complexity)] // Currently only used once
+    batches: Arc<Mutex<Option<(Vec<RecordBatch>, bool)>>>,
 }
 
 impl PyDataFrame {
     /// creates a new PyDataFrame
     pub fn new(df: DataFrame) -> Self {
         Self {
             df: Arc::new(df),
-            batches: None,
+            batches: Arc::new(Mutex::new(None)),
         }
     }
 
-    fn prepare_repr_string(&mut self, py: Python, as_html: bool) -> PyDataFusionResult<String> {
+    fn prepare_repr_string(&self, py: Python, as_html: bool) -> PyDataFusionResult<String> {
         // Get the Python formatter and config
         let PythonFormatter { formatter, config } = get_python_formatter_with_config(py)?;
 
-        let should_cache = *is_ipython_env(py) && self.batches.is_none();
-        let (batches, has_more) = match self.batches.take() {
+        let should_cache = *is_ipython_env(py) && self.batches.lock().unwrap().is_none();
+        let (batches, has_more) = match self.batches.lock().unwrap().take() {
             Some(b) => b,
             None => wait_for_future(
                 py,
@@ -347,7 +347,7 @@ impl PyDataFrame {
         let html_str: String = html_result.extract()?;
 
         if should_cache {
-            self.batches = Some((batches, has_more));
+            *self.batches.lock().unwrap() = Some((batches, has_more));
         }
 
         Ok(html_str)
@@ -377,7 +377,7 @@ impl PyDataFrame {
         }
     }
 
-    fn __repr__(&mut self, py: Python) -> PyDataFusionResult<String> {
+    fn __repr__(&self, py: Python) -> PyDataFusionResult<String> {
         self.prepare_repr_string(py, false)
     }
 
@@ -412,7 +412,7 @@ impl PyDataFrame {
         Ok(format!("DataFrame()\n{batches_as_displ}{additional_str}"))
     }
 
-    fn _repr_html_(&mut self, py: Python) -> PyDataFusionResult<String> {
+    fn _repr_html_(&self, py: Python) -> PyDataFusionResult<String> {
         self.prepare_repr_string(py, true)
     }
 
@@ -875,7 +875,7 @@ impl PyDataFrame {
 
     #[pyo3(signature = (requested_schema=None))]
     fn __arrow_c_stream__<'py>(
-        &'py mut self,
+        &'py self,
         py: Python<'py>,
         requested_schema: Option<Bound<'py, PyCapsule>>,
     ) -> PyDataFusionResult<Bound<'py, PyCapsule>> {