feat/making global context accessible for users

python/datafusion/context.py

@@ -468,6 +468,8 @@ class SessionContext:
     See :ref:`user_guide_concepts` in the online documentation for more information.
     """
 
+    _global_instance = None
+
     def __init__(
         self,
         config: SessionConfig | None = None,
@@ -498,6 +500,18 @@ def __init__(
 
         self.ctx = SessionContextInternal(config, runtime)
 
+    @classmethod
+    def global_ctx(cls) -> "SessionContextInternal":
+        """Retrieve the global context.
+
+        Returns:
+            A `SessionContextInternal` object that corresponds to the global context
+        """
+        if cls._global_instance is None:
+            internal_ctx = SessionContextInternal.global_ctx()
+            cls._global_instance = internal_ctx
+        return cls._global_instance
+
     def enable_url_table(self) -> "SessionContext":
         """Control if local files can be queried as tables.
 

python/datafusion/io.py

@@ -26,7 +26,7 @@
 from datafusion.dataframe import DataFrame
 from datafusion.expr import Expr
 
-from ._internal import SessionContext as SessionContextInternal
+from datafusion.context import SessionContext
 
 
 def read_parquet(
@@ -65,7 +65,7 @@ def read_parquet(
     if table_partition_cols is None:
         table_partition_cols = []
     return DataFrame(
-        SessionContextInternal._global_ctx().read_parquet(
+        SessionContext.global_ctx().read_parquet(
             str(path),
             table_partition_cols,
             parquet_pruning,
@@ -107,7 +107,7 @@ def read_json(
     if table_partition_cols is None:
         table_partition_cols = []
     return DataFrame(
-        SessionContextInternal._global_ctx().read_json(
+        SessionContext.global_ctx().read_json(
             str(path),
             schema,
             schema_infer_max_records,
@@ -158,7 +158,7 @@ def read_csv(
     path = [str(p) for p in path] if isinstance(path, list) else str(path)
 
     return DataFrame(
-        SessionContextInternal._global_ctx().read_csv(
+        SessionContext.global_ctx().read_csv(
             path,
             schema,
             has_header,
@@ -195,7 +195,7 @@ def read_avro(
     if file_partition_cols is None:
         file_partition_cols = []
     return DataFrame(
-        SessionContextInternal._global_ctx().read_avro(
+        SessionContext.global_ctx().read_avro(
             str(path), schema, file_partition_cols, file_extension
         )
     )

python/tests/test_context.py

@@ -32,6 +32,8 @@
     literal,
 )
 
+from datafusion._internal import SessionContext as SessionContextInternal
+
 
 def test_create_context_no_args():
     SessionContext()
@@ -629,3 +631,32 @@ def test_sql_with_options_no_statements(ctx):
     options = SQLOptions().with_allow_statements(False)
     with pytest.raises(Exception, match="SetVariable"):
         ctx.sql_with_options(sql, options=options)
+
+
+def test_global_context_type():
+    ctx = SessionContext.global_ctx()
+    assert isinstance(ctx, SessionContextInternal)
+
+
+def test_global_context_is_singleton():
+    ctx1 = SessionContext.global_ctx()
+    ctx2 = SessionContext.global_ctx()
+    assert ctx1 is ctx2
+
+
+@pytest.fixture
+def batch():
+    return pa.RecordBatch.from_arrays(
+        [pa.array([4, 5, 6])],
+        names=["a"],
+    )
+
+
+def test_create_dataframe_with_global_ctx(batch):
+    ctx = SessionContext.global_ctx()
+
+    df = ctx.create_dataframe([[batch]])
+
+    result = df.collect()[0].column(0)
+
+    assert result == pa.array([4, 5, 6])

src/context.rs

@@ -308,7 +308,7 @@ impl PySessionContext {
 
     #[classmethod]
     #[pyo3(signature = ())]
-    fn _global_ctx(_cls: &Bound<'_, PyType>) -> PyResult<Self> {
+    fn global_ctx(_cls: &Bound<'_, PyType>) -> PyResult<Self> {
         Ok(Self {
             ctx: get_global_ctx().clone(),
         })