feat: enhance PyDataFrame and PySessionContext to maintain session state across operations

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))
+        Ok(PyDataFrame::new(df, self.ctx.state().into()))
     }
 
     #[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))
+        Ok(PyDataFrame::new(df, self.ctx.state().into()))
     }
 
     #[pyo3(signature = (partitions, name=None, schema=None))]
@@ -486,13 +486,16 @@ impl PySessionContext {
 
         let table = wait_for_future(py, self._table(&table_name))??;
 
-        let df = PyDataFrame::new(table);
+        let df = PyDataFrame::new(table, self.ctx.state().into());
         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()))
+        PyDataFrame::new(
+            DataFrame::new(self.ctx.state(), plan.plan.as_ref().clone()),
+            self.ctx.state().into(),
+        )
     }
 
     /// Construct datafusion dataframe from Python list
@@ -913,7 +916,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)),
+            Ok(df) => Ok(PyDataFrame::new(df, self.ctx.state().into())),
             Err(e) => {
                 if let datafusion::error::DataFusionError::Plan(msg) = &e {
                     if msg.contains("No table named") {
@@ -930,7 +933,10 @@ impl PySessionContext {
     }
 
     pub fn empty_table(&self) -> PyDataFusionResult<PyDataFrame> {
-        Ok(PyDataFrame::new(self.ctx.read_empty()?))
+        Ok(PyDataFrame::new(
+            self.ctx.read_empty()?,
+            self.ctx.state().into(),
+        ))
     }
 
     pub fn session_id(&self) -> String {
@@ -970,7 +976,7 @@ impl PySessionContext {
             let result = self.ctx.read_json(path, options);
             wait_for_future(py, result)??
         };
-        Ok(PyDataFrame::new(df))
+        Ok(PyDataFrame::new(df, self.ctx.state().into()))
     }
 
     #[allow(clippy::too_many_arguments)]
@@ -1020,12 +1026,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)??);
+            let df = PyDataFrame::new(wait_for_future(py, result)??, self.ctx.state().into());
             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)??);
+            let df = PyDataFrame::new(wait_for_future(py, result)??, self.ctx.state().into());
             Ok(df)
         }
     }
@@ -1068,7 +1074,7 @@ impl PySessionContext {
             .collect();
 
         let result = self.ctx.read_parquet(path, options);
-        let df = PyDataFrame::new(wait_for_future(py, result)??);
+        let df = PyDataFrame::new(wait_for_future(py, result)??, self.ctx.state().into());
         Ok(df)
     }
 
@@ -1097,12 +1103,12 @@ impl PySessionContext {
             let read_future = self.ctx.read_avro(path, options);
             wait_for_future(py, read_future)??
         };
-        Ok(PyDataFrame::new(df))
+        Ok(PyDataFrame::new(df, self.ctx.state().into()))
     }
 
     pub fn read_table(&self, table: &PyTable) -> PyDataFusionResult<PyDataFrame> {
         let df = self.ctx.read_table(table.table())?;
-        Ok(PyDataFrame::new(df))
+        Ok(PyDataFrame::new(df, self.ctx.state().into()))
     }
 
     fn __repr__(&self) -> PyResult<String> {
@@ -1133,7 +1139,7 @@ impl PySessionContext {
         let 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))
+        Ok(PyRecordBatchStream::new(stream, state.into()))
     }
 }
 

src/dataframe.rs

@@ -292,15 +292,20 @@ impl PyParquetColumnOptions {
 pub struct PyDataFrame {
     df: Arc<DataFrame>,
 
+    // Hold the session state so streams/readers can keep the
+    // underlying SessionContext alive while Python iterates.
+    session_state: Arc<SessionState>,
+
     // In IPython environment cache batches between __repr__ and _repr_html_ calls.
     batches: Option<(Vec<RecordBatch>, bool)>,
 }
 
 impl PyDataFrame {
     /// creates a new PyDataFrame
-    pub fn new(df: DataFrame) -> Self {
+    pub fn new(df: DataFrame, session_state: Arc<SessionState>) -> Self {
         Self {
             df: Arc::new(df),
+            session_state,
             batches: None,
         }
     }
@@ -481,7 +486,7 @@ impl PyDataFrame {
     fn describe(&self, py: Python) -> PyDataFusionResult<Self> {
         let df = self.df.as_ref().clone();
         let stat_df = spawn_future(py, async move { df.describe().await })?;
-        Ok(Self::new(stat_df))
+        Ok(Self::new(stat_df, self.session_state.clone()))
     }
 
     /// Returns the schema from the logical plan
@@ -511,31 +516,31 @@ impl PyDataFrame {
     fn select_columns(&self, args: Vec<PyBackedStr>) -> PyDataFusionResult<Self> {
         let args = args.iter().map(|s| s.as_ref()).collect::<Vec<&str>>();
         let df = self.df.as_ref().clone().select_columns(&args)?;
-        Ok(Self::new(df))
+        Ok(Self::new(df, self.session_state.clone()))
     }
 
     #[pyo3(signature = (*args))]
     fn select(&self, args: Vec<PyExpr>) -> PyDataFusionResult<Self> {
         let expr: Vec<Expr> = args.into_iter().map(|e| e.into()).collect();
         let df = self.df.as_ref().clone().select(expr)?;
-        Ok(Self::new(df))
+        Ok(Self::new(df, self.session_state.clone()))
     }
 
     #[pyo3(signature = (*args))]
     fn drop(&self, args: Vec<PyBackedStr>) -> PyDataFusionResult<Self> {
         let cols = args.iter().map(|s| s.as_ref()).collect::<Vec<&str>>();
         let df = self.df.as_ref().clone().drop_columns(&cols)?;
-        Ok(Self::new(df))
+        Ok(Self::new(df, self.session_state.clone()))
     }
 
     fn filter(&self, predicate: PyExpr) -> PyDataFusionResult<Self> {
         let df = self.df.as_ref().clone().filter(predicate.into())?;
-        Ok(Self::new(df))
+        Ok(Self::new(df, self.session_state.clone()))
     }
 
     fn with_column(&self, name: &str, expr: PyExpr) -> PyDataFusionResult<Self> {
         let df = self.df.as_ref().clone().with_column(name, expr.into())?;
-        Ok(Self::new(df))
+        Ok(Self::new(df, self.session_state.clone()))
     }
 
     fn with_columns(&self, exprs: Vec<PyExpr>) -> PyDataFusionResult<Self> {
@@ -545,7 +550,7 @@ impl PyDataFrame {
             let name = format!("{}", expr.schema_name());
             df = df.with_column(name.as_str(), expr)?
         }
-        Ok(Self::new(df))
+        Ok(Self::new(df, self.session_state.clone()))
     }
 
     /// Rename one column by applying a new projection. This is a no-op if the column to be
@@ -556,27 +561,27 @@ impl PyDataFrame {
             .as_ref()
             .clone()
             .with_column_renamed(old_name, new_name)?;
-        Ok(Self::new(df))
+        Ok(Self::new(df, self.session_state.clone()))
     }
 
     fn aggregate(&self, group_by: Vec<PyExpr>, aggs: Vec<PyExpr>) -> PyDataFusionResult<Self> {
         let group_by = group_by.into_iter().map(|e| e.into()).collect();
         let aggs = aggs.into_iter().map(|e| e.into()).collect();
         let df = self.df.as_ref().clone().aggregate(group_by, aggs)?;
-        Ok(Self::new(df))
+        Ok(Self::new(df, self.session_state.clone()))
     }
 
     #[pyo3(signature = (*exprs))]
     fn sort(&self, exprs: Vec<PySortExpr>) -> PyDataFusionResult<Self> {
         let exprs = to_sort_expressions(exprs);
         let df = self.df.as_ref().clone().sort(exprs)?;
-        Ok(Self::new(df))
+        Ok(Self::new(df, self.session_state.clone()))
     }
 
     #[pyo3(signature = (count, offset=0))]
     fn limit(&self, count: usize, offset: usize) -> PyDataFusionResult<Self> {
         let df = self.df.as_ref().clone().limit(offset, Some(count))?;
-        Ok(Self::new(df))
+        Ok(Self::new(df, self.session_state.clone()))
     }
 
     /// Executes the plan, returning a list of `RecordBatch`es.
@@ -593,7 +598,7 @@ impl PyDataFrame {
     /// Cache DataFrame.
     fn cache(&self, py: Python) -> PyDataFusionResult<Self> {
         let df = wait_for_future(py, self.df.as_ref().clone().cache())??;
-        Ok(Self::new(df))
+        Ok(Self::new(df, self.session_state.clone()))
     }
 
     /// Executes this DataFrame and collects all results into a vector of vector of RecordBatch
@@ -618,7 +623,7 @@ impl PyDataFrame {
     /// Filter out duplicate rows
     fn distinct(&self) -> PyDataFusionResult<Self> {
         let df = self.df.as_ref().clone().distinct()?;
-        Ok(Self::new(df))
+        Ok(Self::new(df, self.session_state.clone()))
     }
 
     fn join(
@@ -652,7 +657,7 @@ impl PyDataFrame {
             &right_keys,
             None,
         )?;
-        Ok(Self::new(df))
+        Ok(Self::new(df, self.session_state.clone()))
     }
 
     fn join_on(
@@ -681,7 +686,7 @@ impl PyDataFrame {
             .as_ref()
             .clone()
             .join_on(right.df.as_ref().clone(), join_type, exprs)?;
-        Ok(Self::new(df))
+        Ok(Self::new(df, self.session_state.clone()))
     }
 
     /// Print the query plan
@@ -714,7 +719,7 @@ impl PyDataFrame {
             .as_ref()
             .clone()
             .repartition(Partitioning::RoundRobinBatch(num))?;
-        Ok(Self::new(new_df))
+        Ok(Self::new(new_df, self.session_state.clone()))
     }
 
     /// Repartition a `DataFrame` based on a logical partitioning scheme.
@@ -726,7 +731,7 @@ impl PyDataFrame {
             .as_ref()
             .clone()
             .repartition(Partitioning::Hash(expr, num))?;
-        Ok(Self::new(new_df))
+        Ok(Self::new(new_df, self.session_state.clone()))
     }
 
     /// Calculate the union of two `DataFrame`s, preserving duplicate rows.The
@@ -742,7 +747,7 @@ impl PyDataFrame {
             self.df.as_ref().clone().union(py_df.df.as_ref().clone())?
         };
 
-        Ok(Self::new(new_df))
+        Ok(Self::new(new_df, self.session_state.clone()))
     }
 
     /// Calculate the distinct union of two `DataFrame`s.  The
@@ -753,7 +758,7 @@ impl PyDataFrame {
             .as_ref()
             .clone()
             .union_distinct(py_df.df.as_ref().clone())?;
-        Ok(Self::new(new_df))
+        Ok(Self::new(new_df, self.session_state.clone()))
     }
 
     #[pyo3(signature = (column, preserve_nulls=true))]
@@ -766,7 +771,7 @@ impl PyDataFrame {
             .as_ref()
             .clone()
             .unnest_columns_with_options(&[column], unnest_options)?;
-        Ok(Self::new(df))
+        Ok(Self::new(df, self.session_state.clone()))
     }
 
     #[pyo3(signature = (columns, preserve_nulls=true))]
@@ -784,7 +789,7 @@ impl PyDataFrame {
             .as_ref()
             .clone()
             .unnest_columns_with_options(&cols, unnest_options)?;
-        Ok(Self::new(df))
+        Ok(Self::new(df, self.session_state.clone()))
     }
 
     /// Calculate the intersection of two `DataFrame`s.  The two `DataFrame`s must have exactly the same schema
@@ -794,13 +799,13 @@ impl PyDataFrame {
             .as_ref()
             .clone()
             .intersect(py_df.df.as_ref().clone())?;
-        Ok(Self::new(new_df))
+        Ok(Self::new(new_df, self.session_state.clone()))
     }
 
     /// Calculate the exception of two `DataFrame`s.  The two `DataFrame`s must have exactly the same schema
     fn except_all(&self, py_df: PyDataFrame) -> PyDataFusionResult<Self> {
         let new_df = self.df.as_ref().clone().except(py_df.df.as_ref().clone())?;
-        Ok(Self::new(new_df))
+        Ok(Self::new(new_df, self.session_state.clone()))
     }
 
     /// Write a `DataFrame` to a CSV file.
@@ -957,7 +962,7 @@ impl PyDataFrame {
         requested_schema: Option<Bound<'py, PyCapsule>>,
     ) -> PyDataFusionResult<Bound<'py, PyCapsule>> {
         let df = self.df.as_ref().clone();
-        let state = df.session_state().clone();
+        let state = self.session_state.clone();
         let streams = spawn_future(py, async move { df.execute_stream_partitioned().await })?;
         let streams = streams
             .into_iter()
@@ -997,14 +1002,14 @@ impl PyDataFrame {
 
     fn execute_stream(&self, py: Python) -> PyDataFusionResult<PyRecordBatchStream> {
         let df = self.df.as_ref().clone();
-        let state = df.session_state().clone();
+        let state = self.session_state.clone();
         let stream = spawn_future(py, async move { df.execute_stream().await })?;
         Ok(PyRecordBatchStream::new(stream, state))
     }
 
     fn execute_stream_partitioned(&self, py: Python) -> PyResult<Vec<PyRecordBatchStream>> {
         let df = self.df.as_ref().clone();
-        let state = df.session_state().clone();
+        let state = self.session_state.clone();
         let streams = spawn_future(py, async move { df.execute_stream_partitioned().await })?;
         Ok(streams
             .into_iter()
@@ -1073,7 +1078,7 @@ impl PyDataFrame {
         };
 
         let df = self.df.as_ref().clone().fill_null(scalar_value, cols)?;
-        Ok(Self::new(df))
+        Ok(Self::new(df, self.session_state.clone()))
     }
 }