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

Author: kosiew

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

@@ -16,7 +16,7 @@
 // under the License.
 
 use std::fmt::{self, Display, Formatter};
-use std::sync::{Arc, RwLock};
+use std::sync::Arc;
 use std::{any::Any, borrow::Cow};
 
 use arrow::datatypes::Schema;
@@ -25,7 +25,6 @@ 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;
@@ -34,6 +33,8 @@ use crate::sql::logical::PyLogicalPlan;
 
 use super::{data_type::DataTypeMap, function::SqlFunction};
 
+use parking_lot::RwLock;
+
 #[pyclass(name = "SqlSchema", module = "datafusion.common", subclass, frozen)]
 #[derive(Debug, Clone)]
 pub struct SqlSchema {
@@ -110,88 +111,60 @@ impl SqlSchema {
 
     #[getter]
     fn name(&self) -> PyResult<String> {
-        Ok(self
-            .name
-            .read()
-            .map_err(|_| PyRuntimeError::new_err("failed to read schema name"))?
-            .clone())
+        Ok(self.name.read().clone())
     }
 
     #[setter]
     fn set_name(&self, value: String) -> PyResult<()> {
-        *self
-            .name
-            .write()
-            .map_err(|_| PyRuntimeError::new_err("failed to write schema name"))? = value;
+        *self.name.write() = value;
         Ok(())
     }
 
     #[getter]
     fn tables(&self) -> PyResult<Vec<SqlTable>> {
-        Ok(self
-            .tables
-            .read()
-            .map_err(|_| PyRuntimeError::new_err("failed to read schema tables"))?
-            .clone())
+        Ok(self.tables.read().clone())
     }
 
     #[setter]
     fn set_tables(&self, tables: Vec<SqlTable>) -> PyResult<()> {
-        *self
-            .tables
-            .write()
-            .map_err(|_| PyRuntimeError::new_err("failed to write schema tables"))? = tables;
+        *self.tables.write() = tables;
         Ok(())
     }
 
     #[getter]
     fn views(&self) -> PyResult<Vec<SqlView>> {
-        Ok(self
-            .views
-            .read()
-            .map_err(|_| PyRuntimeError::new_err("failed to read schema views"))?
-            .clone())
+        Ok(self.views.read().clone())
     }
 
     #[setter]
     fn set_views(&self, views: Vec<SqlView>) -> PyResult<()> {
-        *self
-            .views
-            .write()
-            .map_err(|_| PyRuntimeError::new_err("failed to write schema views"))? = views;
+        *self.views.write() = views;
         Ok(())
     }
 
     #[getter]
     fn functions(&self) -> PyResult<Vec<SqlFunction>> {
-        Ok(self
-            .functions
-            .read()
-            .map_err(|_| PyRuntimeError::new_err("failed to read schema functions"))?
-            .clone())
+        Ok(self.functions.read().clone())
     }
 
     #[setter]
     fn set_functions(&self, functions: Vec<SqlFunction>) -> PyResult<()> {
-        *self
-            .functions
-            .write()
-            .map_err(|_| PyRuntimeError::new_err("failed to write schema functions"))? = functions;
+        *self.functions.write() = functions;
         Ok(())
     }
 
     pub fn table_by_name(&self, table_name: &str) -> Option<SqlTable> {
-        let tables = self.tables.read().expect("failed to read schema tables");
+        let tables = self.tables.read();
         tables.iter().find(|tbl| tbl.name.eq(table_name)).cloned()
     }
 
     pub fn add_table(&self, table: SqlTable) {
-        let mut tables = self.tables.write().expect("failed to write schema tables");
+        let mut tables = self.tables.write();
         tables.push(table);
     }
 
     pub fn drop_table(&self, table_name: String) {
-        let mut tables = self.tables.write().expect("failed to write schema tables");
+        let mut tables = self.tables.write();
         tables.retain(|x| !x.name.eq(&table_name));
     }
 }

src/common/schema.rs

@@ -15,15 +15,16 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::sync::{Arc, RwLock};
+use std::sync::Arc;
 
 use pyo3::prelude::*;
 use pyo3::types::*;
 
 use datafusion::config::ConfigOptions;
 
-use crate::errors::{PyDataFusionError, PyDataFusionResult};
+use crate::errors::PyDataFusionResult;
 use crate::utils::py_obj_to_scalar_value;
+use parking_lot::RwLock;
 
 #[pyclass(name = "Config", module = "datafusion", subclass, frozen)]
 #[derive(Clone)]
@@ -50,10 +51,7 @@ impl PyConfig {
 
     /// Get a configuration option
     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()))?;
+        let options = self.config.read();
         for entry in options.entries() {
             if entry.key == key {
                 return Ok(entry.value.into_pyobject(py)?);
@@ -65,21 +63,15 @@ impl PyConfig {
     /// Set a configuration option
     pub fn set(&self, key: &str, value: PyObject, py: Python) -> PyDataFusionResult<()> {
         let scalar_value = py_obj_to_scalar_value(py, value)?;
-        let mut options = self
-            .config
-            .write()
-            .map_err(|_| PyDataFusionError::Common("failed to lock configuration".to_string()))?;
+        let mut options = self.config.write();
         options.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
-            .read()
-            .map_err(|_| PyDataFusionError::Common("failed to read configuration".to_string()))?;
+        let options = self.config.read();
         for entry in options.entries() {
             dict.set_item(entry.key, entry.value.clone().into_pyobject(py)?)?;
         }

src/config.rs

@@ -17,7 +17,7 @@
 
 use std::collections::HashMap;
 use std::ffi::CString;
-use std::sync::{Arc, Mutex};
+use std::sync::Arc;
 
 use arrow::array::{new_null_array, RecordBatch, RecordBatchIterator, RecordBatchReader};
 use arrow::compute::can_cast_types;
@@ -58,6 +58,8 @@ use crate::{
     expr::{sort_expr::PySortExpr, PyExpr},
 };
 
+use parking_lot::Mutex;
+
 // https://github.com/apache/datafusion-python/pull/1016#discussion_r1983239116
 // - we have not decided on the table_provider approach yet
 // this is an interim implementation
@@ -307,9 +309,7 @@ impl PyDataFrame {
         let PythonFormatter { formatter, config } = get_python_formatter_with_config(py)?;
 
         let (cached_batches, should_cache) = {
-            let mut cache = self.batches.lock().map_err(|_| {
-                PyDataFusionError::Common("failed to lock DataFrame display cache".to_string())
-            })?;
+            let mut cache = self.batches.lock();
             let should_cache = *is_ipython_env(py) && cache.is_none();
             let batches = cache.take();
             (batches, should_cache)
@@ -354,9 +354,7 @@ impl PyDataFrame {
         let html_str: String = html_result.extract()?;
 
         if should_cache {
-            let mut cache = self.batches.lock().map_err(|_| {
-                PyDataFusionError::Common("failed to lock DataFrame display cache".to_string())
-            })?;
+            let mut cache = self.batches.lock();
             *cache = Some((batches.clone(), has_more));
         }
 

src/dataframe.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::sync::{Arc, Mutex};
+use std::sync::Arc;
 
 use crate::{
     errors::{PyDataFusionError, PyDataFusionResult},
@@ -24,23 +24,14 @@ use crate::{
 use datafusion::logical_expr::conditional_expressions::CaseBuilder;
 use pyo3::prelude::*;
 
+use parking_lot::{Mutex, MutexGuard};
+
 #[pyclass(name = "CaseBuilder", module = "datafusion.expr", subclass, frozen)]
 #[derive(Clone)]
 pub struct PyCaseBuilder {
     case_builder: Arc<Mutex<Option<CaseBuilder>>>,
 }
 
-impl From<PyCaseBuilder> for CaseBuilder {
-    fn from(case_builder: PyCaseBuilder) -> Self {
-        case_builder
-            .case_builder
-            .lock()
-            .expect("Case builder mutex poisoned")
-            .take()
-            .expect("CaseBuilder has already been consumed")
-    }
-}
-
 impl From<CaseBuilder> for PyCaseBuilder {
     fn from(case_builder: CaseBuilder) -> PyCaseBuilder {
         PyCaseBuilder {
@@ -50,25 +41,27 @@ impl From<CaseBuilder> for PyCaseBuilder {
 }
 
 impl PyCaseBuilder {
-    fn lock_case_builder(
-        &self,
-    ) -> PyDataFusionResult<std::sync::MutexGuard<'_, Option<CaseBuilder>>> {
-        self.case_builder
-            .lock()
-            .map_err(|_| PyDataFusionError::Common("failed to lock CaseBuilder".to_string()))
+    fn lock_case_builder(&self) -> MutexGuard<'_, Option<CaseBuilder>> {
+        self.case_builder.lock()
     }
 
     fn take_case_builder(&self) -> PyDataFusionResult<CaseBuilder> {
-        let mut guard = self.lock_case_builder()?;
+        let mut guard = self.lock_case_builder();
         guard.take().ok_or_else(|| {
             PyDataFusionError::Common("CaseBuilder has already been consumed".to_string())
         })
     }
 
-    fn store_case_builder(&self, builder: CaseBuilder) -> PyDataFusionResult<()> {
-        let mut guard = self.lock_case_builder()?;
+    fn store_case_builder(&self, builder: CaseBuilder) {
+        let mut guard = self.lock_case_builder();
         *guard = Some(builder);
-        Ok(())
+    }
+
+    pub fn into_case_builder(self) -> PyDataFusionResult<CaseBuilder> {
+        let mut guard = self.case_builder.lock();
+        guard.take().ok_or_else(|| {
+            PyDataFusionError::Common("CaseBuilder has already been consumed".to_string())
+        })
     }
 }
 
@@ -77,7 +70,7 @@ impl PyCaseBuilder {
     fn when(&self, when: PyExpr, then: PyExpr) -> PyDataFusionResult<PyCaseBuilder> {
         let mut builder = self.take_case_builder()?;
         let next_builder = builder.when(when.expr, then.expr);
-        self.store_case_builder(next_builder)?;
+        self.store_case_builder(next_builder);
         Ok(self.clone())
     }
 
@@ -86,7 +79,7 @@ impl PyCaseBuilder {
         match builder.otherwise(else_expr.expr) {
             Ok(expr) => Ok(expr.clone().into()),
             Err(err) => {
-                self.store_case_builder(builder)?;
+                self.store_case_builder(builder);
                 Err(err.into())
             }
         }
@@ -97,7 +90,7 @@ impl PyCaseBuilder {
         match builder.end() {
             Ok(expr) => Ok(expr.clone().into()),
             Err(err) => {
-                self.store_case_builder(builder)?;
+                self.store_case_builder(builder);
                 Err(err.into())
             }
         }