refactor: replace SessionState with SessionContext in PyDataFrame and related structures

Author: kosiew

src/context.rs

@@ -434,7 +434,7 @@ impl PySessionContext {
     pub fn sql(&mut self, query: &str, py: Python) -> PyDataFusionResult<PyDataFrame> {
         let result = self.ctx.sql(query);
         let df = wait_for_future(py, result)??;
-        Ok(PyDataFrame::new(df, self.ctx.state().into()))
+        Ok(PyDataFrame::new(df, Arc::new(self.ctx.clone())))
     }
 
     #[pyo3(signature = (query, options=None))]
@@ -451,7 +451,7 @@ impl PySessionContext {
         };
         let result = self.ctx.sql_with_options(query, options);
         let df = wait_for_future(py, result)??;
-        Ok(PyDataFrame::new(df, self.ctx.state().into()))
+        Ok(PyDataFrame::new(df, Arc::new(self.ctx.clone())))
     }
 
     #[pyo3(signature = (partitions, name=None, schema=None))]
@@ -486,16 +486,14 @@ impl PySessionContext {
 
         let table = wait_for_future(py, self._table(&table_name))??;
 
-        let df = PyDataFrame::new(table, self.ctx.state().into());
+        let df = PyDataFrame::new(table, Arc::new(self.ctx.clone()));
         Ok(df)
     }
 
     /// Create a DataFrame from an existing logical plan
     pub fn create_dataframe_from_logical_plan(&mut self, plan: PyLogicalPlan) -> PyDataFrame {
-        PyDataFrame::new(
-            DataFrame::new(self.ctx.state(), plan.plan.as_ref().clone()),
-            self.ctx.state().into(),
-        )
+        let ctx = Arc::new(self.ctx.clone());
+        PyDataFrame::new(DataFrame::new(ctx.state(), plan.plan.as_ref().clone()), ctx)
     }
 
     /// Construct datafusion dataframe from Python list
@@ -916,7 +914,7 @@ impl PySessionContext {
         let res = wait_for_future(py, self.ctx.table(name))
             .map_err(|e| PyKeyError::new_err(e.to_string()))?;
         match res {
-            Ok(df) => Ok(PyDataFrame::new(df, self.ctx.state().into())),
+            Ok(df) => Ok(PyDataFrame::new(df, Arc::new(self.ctx.clone()))),
             Err(e) => {
                 if let datafusion::error::DataFusionError::Plan(msg) = &e {
                     if msg.contains("No table named") {
@@ -935,7 +933,7 @@ impl PySessionContext {
     pub fn empty_table(&self) -> PyDataFusionResult<PyDataFrame> {
         Ok(PyDataFrame::new(
             self.ctx.read_empty()?,
-            self.ctx.state().into(),
+            Arc::new(self.ctx.clone()),
         ))
     }
 
@@ -976,7 +974,7 @@ impl PySessionContext {
             let result = self.ctx.read_json(path, options);
             wait_for_future(py, result)??
         };
-        Ok(PyDataFrame::new(df, self.ctx.state().into()))
+        Ok(PyDataFrame::new(df, Arc::new(self.ctx.clone())))
     }
 
     #[allow(clippy::too_many_arguments)]
@@ -1026,12 +1024,12 @@ impl PySessionContext {
             let paths = path.extract::<Vec<String>>()?;
             let paths = paths.iter().map(|p| p as &str).collect::<Vec<&str>>();
             let result = self.ctx.read_csv(paths, options);
-            let df = PyDataFrame::new(wait_for_future(py, result)??, self.ctx.state().into());
+            let df = PyDataFrame::new(wait_for_future(py, result)??, Arc::new(self.ctx.clone()));
             Ok(df)
         } else {
             let path = path.extract::<String>()?;
             let result = self.ctx.read_csv(path, options);
-            let df = PyDataFrame::new(wait_for_future(py, result)??, self.ctx.state().into());
+            let df = PyDataFrame::new(wait_for_future(py, result)??, Arc::new(self.ctx.clone()));
             Ok(df)
         }
     }
@@ -1074,7 +1072,7 @@ impl PySessionContext {
             .collect();
 
         let result = self.ctx.read_parquet(path, options);
-        let df = PyDataFrame::new(wait_for_future(py, result)??, self.ctx.state().into());
+        let df = PyDataFrame::new(wait_for_future(py, result)??, Arc::new(self.ctx.clone()));
         Ok(df)
     }
 
@@ -1103,12 +1101,12 @@ impl PySessionContext {
             let read_future = self.ctx.read_avro(path, options);
             wait_for_future(py, read_future)??
         };
-        Ok(PyDataFrame::new(df, self.ctx.state().into()))
+        Ok(PyDataFrame::new(df, Arc::new(self.ctx.clone())))
     }
 
     pub fn read_table(&self, table: &PyTable) -> PyDataFusionResult<PyDataFrame> {
         let df = self.ctx.read_table(table.table())?;
-        Ok(PyDataFrame::new(df, self.ctx.state().into()))
+        Ok(PyDataFrame::new(df, Arc::new(self.ctx.clone())))
     }
 
     fn __repr__(&self) -> PyResult<String> {
@@ -1135,11 +1133,12 @@ impl PySessionContext {
         part: usize,
         py: Python,
     ) -> PyDataFusionResult<PyRecordBatchStream> {
-        let state = self.ctx.state();
-        let ctx: TaskContext = TaskContext::from(&state);
+        let ctx = Arc::new(self.ctx.clone());
+        let state = ctx.state();
+        let task_ctx: TaskContext = TaskContext::from(&state);
         let plan = plan.plan.clone();
-        let stream = spawn_future(py, async move { plan.execute(part, Arc::new(ctx)) })?;
-        Ok(PyRecordBatchStream::new(stream, state.into()))
+        let stream = spawn_future(py, async move { plan.execute(part, Arc::new(task_ctx)) })?;
+        Ok(PyRecordBatchStream::new(stream, ctx))
     }
 }