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

Author: kosiew

src/common/schema.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 use std::fmt::{self, Display, Formatter};
-use std::sync::Arc;
+use std::sync::{Arc, RwLock};
 use std::{any::Any, borrow::Cow};
 
 use arrow::datatypes::Schema;
@@ -25,6 +25,7 @@ use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::common::Constraints;
 use datafusion::datasource::TableType;
 use datafusion::logical_expr::{Expr, TableProviderFilterPushDown, TableSource};
+use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
 
 use datafusion::logical_expr::utils::split_conjunction;
@@ -33,17 +34,13 @@ use crate::sql::logical::PyLogicalPlan;
 
 use super::{data_type::DataTypeMap, function::SqlFunction};
 
-#[pyclass(name = "SqlSchema", module = "datafusion.common", subclass)]
+#[pyclass(name = "SqlSchema", module = "datafusion.common", subclass, frozen)]
 #[derive(Debug, Clone)]
 pub struct SqlSchema {
-    #[pyo3(get, set)]
-    pub name: String,
-    #[pyo3(get, set)]
-    pub tables: Vec<SqlTable>,
-    #[pyo3(get, set)]
-    pub views: Vec<SqlView>,
-    #[pyo3(get, set)]
-    pub functions: Vec<SqlFunction>,
+    name: Arc<RwLock<String>>,
+    tables: Arc<RwLock<Vec<SqlTable>>>,
+    views: Arc<RwLock<Vec<SqlView>>>,
+    functions: Arc<RwLock<Vec<SqlFunction>>>,
 }
 
 #[pyclass(name = "SqlTable", module = "datafusion.common", subclass)]
@@ -104,28 +101,98 @@ impl SqlSchema {
     #[new]
     pub fn new(schema_name: &str) -> Self {
         Self {
-            name: schema_name.to_owned(),
-            tables: Vec::new(),
-            views: Vec::new(),
-            functions: Vec::new(),
+            name: Arc::new(RwLock::new(schema_name.to_owned())),
+            tables: Arc::new(RwLock::new(Vec::new())),
+            views: Arc::new(RwLock::new(Vec::new())),
+            functions: Arc::new(RwLock::new(Vec::new())),
         }
     }
 
+    #[getter]
+    fn name(&self) -> PyResult<String> {
+        Ok(self
+            .name
+            .read()
+            .map_err(|_| PyRuntimeError::new_err("failed to read schema name"))?
+            .clone())
+    }
+
+    #[setter]
+    fn set_name(&self, value: String) -> PyResult<()> {
+        *self
+            .name
+            .write()
+            .map_err(|_| PyRuntimeError::new_err("failed to write schema name"))? = value;
+        Ok(())
+    }
+
+    #[getter]
+    fn tables(&self) -> PyResult<Vec<SqlTable>> {
+        Ok(self
+            .tables
+            .read()
+            .map_err(|_| PyRuntimeError::new_err("failed to read schema tables"))?
+            .clone())
+    }
+
+    #[setter]
+    fn set_tables(&self, tables: Vec<SqlTable>) -> PyResult<()> {
+        *self
+            .tables
+            .write()
+            .map_err(|_| PyRuntimeError::new_err("failed to write schema tables"))? = tables;
+        Ok(())
+    }
+
+    #[getter]
+    fn views(&self) -> PyResult<Vec<SqlView>> {
+        Ok(self
+            .views
+            .read()
+            .map_err(|_| PyRuntimeError::new_err("failed to read schema views"))?
+            .clone())
+    }
+
+    #[setter]
+    fn set_views(&self, views: Vec<SqlView>) -> PyResult<()> {
+        *self
+            .views
+            .write()
+            .map_err(|_| PyRuntimeError::new_err("failed to write schema views"))? = views;
+        Ok(())
+    }
+
+    #[getter]
+    fn functions(&self) -> PyResult<Vec<SqlFunction>> {
+        Ok(self
+            .functions
+            .read()
+            .map_err(|_| PyRuntimeError::new_err("failed to read schema functions"))?
+            .clone())
+    }
+
+    #[setter]
+    fn set_functions(&self, functions: Vec<SqlFunction>) -> PyResult<()> {
+        *self
+            .functions
+            .write()
+            .map_err(|_| PyRuntimeError::new_err("failed to write schema functions"))? = functions;
+        Ok(())
+    }
+
     pub fn table_by_name(&self, table_name: &str) -> Option<SqlTable> {
-        for tbl in &self.tables {
-            if tbl.name.eq(table_name) {
-                return Some(tbl.clone());
-            }
-        }
-        None
+        let tables = self.tables.read().expect("failed to read schema tables");
+        tables.iter().find(|tbl| tbl.name.eq(table_name)).cloned()
     }
 
-    pub fn add_table(&mut self, table: SqlTable) {
-        self.tables.push(table);
+    pub fn add_table(&self, table: SqlTable) {
+        let mut tables = self.tables.write().expect("failed to write schema tables");
+        tables.push(table);
     }
 
-    pub fn drop_table(&mut self, table_name: String) {
-        self.tables.retain(|x| !x.name.eq(&table_name));
+    pub fn drop_table(&self, table_name: String) {
+        let mut tables = self.tables.write().expect("failed to write schema tables");
+        tables.retain(|x| !x.name.eq(&table_name));
     }
 }
 

src/config.rs

@@ -15,40 +15,45 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::sync::{Arc, RwLock};
+
 use pyo3::prelude::*;
 use pyo3::types::*;
 
 use datafusion::config::ConfigOptions;
 
-use crate::errors::PyDataFusionResult;
+use crate::errors::{PyDataFusionError, PyDataFusionResult};
 use crate::utils::py_obj_to_scalar_value;
 
-#[pyclass(name = "Config", module = "datafusion", subclass)]
+#[pyclass(name = "Config", module = "datafusion", subclass, frozen)]
 #[derive(Clone)]
 pub(crate) struct PyConfig {
-    config: ConfigOptions,
+    config: Arc<RwLock<ConfigOptions>>,
 }
 
 #[pymethods]
 impl PyConfig {
     #[new]
     fn py_new() -> Self {
         Self {
-            config: ConfigOptions::new(),
+            config: Arc::new(RwLock::new(ConfigOptions::new())),
         }
     }
 
     /// Get configurations from environment variables
     #[staticmethod]
     pub fn from_env() -> PyDataFusionResult<Self> {
         Ok(Self {
-            config: ConfigOptions::from_env()?,
+            config: Arc::new(RwLock::new(ConfigOptions::from_env()?)),
         })
     }
 
     /// Get a configuration option
-    pub fn get<'py>(&mut self, key: &str, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        let options = self.config.to_owned();
+    pub fn get<'py>(&self, key: &str, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        let options = self
+            .config
+            .read()
+            .map_err(|_| PyDataFusionError::Common("failed to read configuration".to_string()))?;
         for entry in options.entries() {
             if entry.key == key {
                 return Ok(entry.value.into_pyobject(py)?);
@@ -58,25 +63,31 @@ 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 options = self
+            .config
+            .write()
+            .map_err(|_| PyDataFusionError::Common("failed to lock configuration".to_string()))?;
+        options.set(key, scalar_value.to_string().as_str())?;
         Ok(())
     }
 
     /// Get all configuration options
-    pub fn get_all(&mut self, py: Python) -> PyResult<PyObject> {
+    pub fn get_all(&self, py: Python) -> PyResult<PyObject> {
         let dict = PyDict::new(py);
-        let options = self.config.to_owned();
+        let options = self
+            .config
+            .read()
+            .map_err(|_| PyDataFusionError::Common("failed to read configuration".to_string()))?;
         for entry in options.entries() {
             dict.set_item(entry.key, entry.value.clone().into_pyobject(py)?)?;
         }
         Ok(dict.into())
     }
 
-    fn __repr__(&mut self, py: Python) -> PyResult<String> {
-        let dict = self.get_all(py);
-        match dict {
+    fn __repr__(&self, py: Python) -> PyResult<String> {
+        match self.get_all(py) {
             Ok(result) => Ok(format!("Config({result})")),
             Err(err) => Ok(format!("Error: {:?}", err.to_string())),
         }

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,30 +284,38 @@ 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.
-#[pyclass(name = "DataFrame", module = "datafusion", subclass)]
+#[pyclass(name = "DataFrame", module = "datafusion", subclass, frozen)]
 #[derive(Clone)]
 pub struct PyDataFrame {
     df: Arc<DataFrame>,
 
     // In IPython environment cache batches between __repr__ and _repr_html_ calls.
-    batches: Option<(Vec<RecordBatch>, bool)>,
+    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 (cached_batches, should_cache) = {
+            let mut cache = self.batches.lock().map_err(|_| {
+                PyDataFusionError::Common("failed to lock DataFrame display cache".to_string())
+            })?;
+            let should_cache = *is_ipython_env(py) && cache.is_none();
+            let batches = cache.take();
+            (batches, should_cache)
+        };
+
+        let (batches, has_more) = match cached_batches {
             Some(b) => b,
             None => wait_for_future(
                 py,
@@ -346,7 +354,10 @@ impl PyDataFrame {
         let html_str: String = html_result.extract()?;
 
         if should_cache {
-            self.batches = Some((batches, has_more));
+            let mut cache = self.batches.lock().map_err(|_| {
+                PyDataFusionError::Common("failed to lock DataFrame display cache".to_string())
+            })?;
+            *cache = Some((batches.clone(), has_more));
         }
 
         Ok(html_str)
@@ -376,7 +387,7 @@ impl PyDataFrame {
         }
     }
 
-    fn __repr__(&mut self, py: Python) -> PyDataFusionResult<String> {
+    fn __repr__(&self, py: Python) -> PyDataFusionResult<String> {
         self.prepare_repr_string(py, false)
     }
 
@@ -411,7 +422,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)
     }
 
@@ -874,7 +885,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>> {