Skip to content

Commit f08d5b0

Browse files
authored
Make Session Context pyclass frozen so interior mutability is only managed by rust (apache#1248)
* Remove unneeded mut self * Consider frozen and rely on underlying rust classes to manage mutability
1 parent 0ec0102 commit f08d5b0

File tree

2 files changed

+27
-35
lines changed

2 files changed

+27
-35
lines changed

src/context.rs

Lines changed: 26 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -296,7 +296,7 @@ impl PySQLOptions {
296296
/// `PySessionContext` is able to plan and execute DataFusion plans.
297297
/// It has a powerful optimizer, a physical planner for local execution, and a
298298
/// multi-threaded execution engine to perform the execution.
299-
#[pyclass(name = "SessionContext", module = "datafusion", subclass)]
299+
#[pyclass(frozen, name = "SessionContext", module = "datafusion", subclass)]
300300
#[derive(Clone)]
301301
pub struct PySessionContext {
302302
pub ctx: SessionContext,
@@ -348,7 +348,7 @@ impl PySessionContext {
348348
/// Register an object store with the given name
349349
#[pyo3(signature = (scheme, store, host=None))]
350350
pub fn register_object_store(
351-
&mut self,
351+
&self,
352352
scheme: &str,
353353
store: StorageContexts,
354354
host: Option<&str>,
@@ -380,7 +380,7 @@ impl PySessionContext {
380380
schema=None,
381381
file_sort_order=None))]
382382
pub fn register_listing_table(
383-
&mut self,
383+
&self,
384384
name: &str,
385385
path: &str,
386386
table_partition_cols: Vec<(String, PyArrowType<DataType>)>,
@@ -426,22 +426,22 @@ impl PySessionContext {
426426
Ok(())
427427
}
428428

429-
pub fn register_udtf(&mut self, func: PyTableFunction) {
429+
pub fn register_udtf(&self, func: PyTableFunction) {
430430
let name = func.name.clone();
431431
let func = Arc::new(func);
432432
self.ctx.register_udtf(&name, func);
433433
}
434434

435435
/// Returns a PyDataFrame whose plan corresponds to the SQL statement.
436-
pub fn sql(&mut self, query: &str, py: Python) -> PyDataFusionResult<PyDataFrame> {
436+
pub fn sql(&self, query: &str, py: Python) -> PyDataFusionResult<PyDataFrame> {
437437
let result = self.ctx.sql(query);
438438
let df = wait_for_future(py, result)??;
439439
Ok(PyDataFrame::new(df))
440440
}
441441

442442
#[pyo3(signature = (query, options=None))]
443443
pub fn sql_with_options(
444-
&mut self,
444+
&self,
445445
query: &str,
446446
options: Option<PySQLOptions>,
447447
py: Python,
@@ -458,7 +458,7 @@ impl PySessionContext {
458458

459459
#[pyo3(signature = (partitions, name=None, schema=None))]
460460
pub fn create_dataframe(
461-
&mut self,
461+
&self,
462462
partitions: PyArrowType<Vec<Vec<RecordBatch>>>,
463463
name: Option<&str>,
464464
schema: Option<PyArrowType<Schema>>,
@@ -493,14 +493,14 @@ impl PySessionContext {
493493
}
494494

495495
/// Create a DataFrame from an existing logical plan
496-
pub fn create_dataframe_from_logical_plan(&mut self, plan: PyLogicalPlan) -> PyDataFrame {
496+
pub fn create_dataframe_from_logical_plan(&self, plan: PyLogicalPlan) -> PyDataFrame {
497497
PyDataFrame::new(DataFrame::new(self.ctx.state(), plan.plan.as_ref().clone()))
498498
}
499499

500500
/// Construct datafusion dataframe from Python list
501501
#[pyo3(signature = (data, name=None))]
502502
pub fn from_pylist(
503-
&mut self,
503+
&self,
504504
data: Bound<'_, PyList>,
505505
name: Option<&str>,
506506
) -> PyResult<PyDataFrame> {
@@ -520,7 +520,7 @@ impl PySessionContext {
520520
/// Construct datafusion dataframe from Python dictionary
521521
#[pyo3(signature = (data, name=None))]
522522
pub fn from_pydict(
523-
&mut self,
523+
&self,
524524
data: Bound<'_, PyDict>,
525525
name: Option<&str>,
526526
) -> PyResult<PyDataFrame> {
@@ -540,7 +540,7 @@ impl PySessionContext {
540540
/// Construct datafusion dataframe from Arrow Table
541541
#[pyo3(signature = (data, name=None))]
542542
pub fn from_arrow(
543-
&mut self,
543+
&self,
544544
data: Bound<'_, PyAny>,
545545
name: Option<&str>,
546546
py: Python,
@@ -574,11 +574,7 @@ impl PySessionContext {
574574
/// Construct datafusion dataframe from pandas
575575
#[allow(clippy::wrong_self_convention)]
576576
#[pyo3(signature = (data, name=None))]
577-
pub fn from_pandas(
578-
&mut self,
579-
data: Bound<'_, PyAny>,
580-
name: Option<&str>,
581-
) -> PyResult<PyDataFrame> {
577+
pub fn from_pandas(&self, data: Bound<'_, PyAny>, name: Option<&str>) -> PyResult<PyDataFrame> {
582578
// Obtain GIL token
583579
let py = data.py();
584580

@@ -594,11 +590,7 @@ impl PySessionContext {
594590

595591
/// Construct datafusion dataframe from polars
596592
#[pyo3(signature = (data, name=None))]
597-
pub fn from_polars(
598-
&mut self,
599-
data: Bound<'_, PyAny>,
600-
name: Option<&str>,
601-
) -> PyResult<PyDataFrame> {
593+
pub fn from_polars(&self, data: Bound<'_, PyAny>, name: Option<&str>) -> PyResult<PyDataFrame> {
602594
// Convert Polars dataframe to Arrow Table
603595
let table = data.call_method0("to_arrow")?;
604596

@@ -607,18 +599,18 @@ impl PySessionContext {
607599
Ok(df)
608600
}
609601

610-
pub fn register_table(&mut self, name: &str, table: &PyTable) -> PyDataFusionResult<()> {
602+
pub fn register_table(&self, name: &str, table: &PyTable) -> PyDataFusionResult<()> {
611603
self.ctx.register_table(name, table.table())?;
612604
Ok(())
613605
}
614606

615-
pub fn deregister_table(&mut self, name: &str) -> PyDataFusionResult<()> {
607+
pub fn deregister_table(&self, name: &str) -> PyDataFusionResult<()> {
616608
self.ctx.deregister_table(name)?;
617609
Ok(())
618610
}
619611

620612
pub fn register_catalog_provider(
621-
&mut self,
613+
&self,
622614
name: &str,
623615
provider: Bound<'_, PyAny>,
624616
) -> PyDataFusionResult<()> {
@@ -647,7 +639,7 @@ impl PySessionContext {
647639

648640
/// Construct datafusion dataframe from Arrow Table
649641
pub fn register_table_provider(
650-
&mut self,
642+
&self,
651643
name: &str,
652644
provider: Bound<'_, PyAny>,
653645
) -> PyDataFusionResult<()> {
@@ -671,7 +663,7 @@ impl PySessionContext {
671663
}
672664

673665
pub fn register_record_batches(
674-
&mut self,
666+
&self,
675667
name: &str,
676668
partitions: PyArrowType<Vec<Vec<RecordBatch>>>,
677669
) -> PyDataFusionResult<()> {
@@ -689,7 +681,7 @@ impl PySessionContext {
689681
schema=None,
690682
file_sort_order=None))]
691683
pub fn register_parquet(
692-
&mut self,
684+
&self,
693685
name: &str,
694686
path: &str,
695687
table_partition_cols: Vec<(String, PyArrowType<DataType>)>,
@@ -732,7 +724,7 @@ impl PySessionContext {
732724
file_extension=".csv",
733725
file_compression_type=None))]
734726
pub fn register_csv(
735-
&mut self,
727+
&self,
736728
name: &str,
737729
path: &Bound<'_, PyAny>,
738730
schema: Option<PyArrowType<Schema>>,
@@ -780,7 +772,7 @@ impl PySessionContext {
780772
table_partition_cols=vec![],
781773
file_compression_type=None))]
782774
pub fn register_json(
783-
&mut self,
775+
&self,
784776
name: &str,
785777
path: PathBuf,
786778
schema: Option<PyArrowType<Schema>>,
@@ -819,7 +811,7 @@ impl PySessionContext {
819811
file_extension=".avro",
820812
table_partition_cols=vec![]))]
821813
pub fn register_avro(
822-
&mut self,
814+
&self,
823815
name: &str,
824816
path: PathBuf,
825817
schema: Option<PyArrowType<Schema>>,
@@ -860,17 +852,17 @@ impl PySessionContext {
860852
Ok(())
861853
}
862854

863-
pub fn register_udf(&mut self, udf: PyScalarUDF) -> PyResult<()> {
855+
pub fn register_udf(&self, udf: PyScalarUDF) -> PyResult<()> {
864856
self.ctx.register_udf(udf.function);
865857
Ok(())
866858
}
867859

868-
pub fn register_udaf(&mut self, udaf: PyAggregateUDF) -> PyResult<()> {
860+
pub fn register_udaf(&self, udaf: PyAggregateUDF) -> PyResult<()> {
869861
self.ctx.register_udaf(udaf.function);
870862
Ok(())
871863
}
872864

873-
pub fn register_udwf(&mut self, udwf: PyWindowUDF) -> PyResult<()> {
865+
pub fn register_udwf(&self, udwf: PyWindowUDF) -> PyResult<()> {
874866
self.ctx.register_udwf(udwf.function);
875867
Ok(())
876868
}
@@ -942,7 +934,7 @@ impl PySessionContext {
942934
#[allow(clippy::too_many_arguments)]
943935
#[pyo3(signature = (path, schema=None, schema_infer_max_records=1000, file_extension=".json", table_partition_cols=vec![], file_compression_type=None))]
944936
pub fn read_json(
945-
&mut self,
937+
&self,
946938
path: PathBuf,
947939
schema: Option<PyArrowType<Schema>>,
948940
schema_infer_max_records: usize,

src/substrait.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,7 @@ impl PySubstraitConsumer {
138138
/// Convert Substrait Plan to DataFusion DataFrame
139139
#[staticmethod]
140140
pub fn from_substrait_plan(
141-
ctx: &mut PySessionContext,
141+
ctx: &PySessionContext,
142142
plan: PyPlan,
143143
py: Python,
144144
) -> PyDataFusionResult<PyLogicalPlan> {

0 commit comments

Comments
 (0)