(datafusion-ballista) branch main updated: fix: use DataFrame.__name__ instead of hardcoded string in metaclass method wrapping (#1866)

[mg...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

martin-g pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git


The following commit(s) were added to refs/heads/main by this push:
     new 493dbfb81 fix: use DataFrame.__name__ instead of hardcoded string in metaclass method wrapping (#1866)
493dbfb81 is described below

commit 493dbfb81e0c95cce4a775d2bbb45daa2cc78dc6
Author: Jogesh Dinavahi <jo...@gmail.com>
AuthorDate: Mon Jun 15 23:34:24 2026 -0700

    fix: use DataFrame.__name__ instead of hardcoded string in metaclass method wrapping (#1866)
    
    * fix dataframe name - todo
    
    * address review comments
    
    * upstream itertools and tower-http deps update
    
    * address review comment - remove whole comment
---
 python/Cargo.lock                   |  2 ++
 python/python/ballista/extension.py | 10 ++--------
 python/python/tests/test_context.py | 29 +++++++++++++++++++++++++++++
 3 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/python/Cargo.lock b/python/Cargo.lock
index 8ccc4358b..dc1c986dc 100644
--- a/python/Cargo.lock
+++ b/python/Cargo.lock
@@ -642,6 +642,7 @@ dependencies = [
  "futures",
  "http",
  "insta",
+ "itertools",
  "log",
  "object_store",
  "parking_lot",
@@ -652,6 +653,7 @@ dependencies = [
  "tokio",
  "tokio-stream",
  "tonic",
+ "tower-http",
  "uuid",
 ]
 
diff --git a/python/python/ballista/extension.py b/python/python/ballista/extension.py
index 90a3651bc..645c99518 100644
--- a/python/python/ballista/extension.py
+++ b/python/python/ballista/extension.py
@@ -71,13 +71,10 @@ class RedefiningDataFrameMeta(type):
             return method_wrapper
 
         for base_name, base_value in bases[0].__dict__.items():
-            #
-            # TODO: could we not use 'DataFrame' as a string here?
-            #
             if (
                 callable(base_value)
                 and not base_name.startswith("__")
-                and base_value.__annotations__.get("return") == "DataFrame"
+                and base_value.__annotations__.get("return") == DataFrame.__name__
             ):
                 #
                 # functions which return DataFrame are redefined
@@ -111,13 +108,10 @@ class RedefiningSessionContextMeta(type):
             return method_wrapper
 
         for base_name, base_value in bases[0].__dict__.items():
-            #
-            # could we not use 'DataFrame' as a string here?
-            #
             if (
                 callable(base_value)
                 and not base_name.startswith("__")
-                and base_value.__annotations__.get("return") == "DataFrame"
+                and base_value.__annotations__.get("return") == DataFrame.__name__
             ):
                 #
                 # functions which return DataFrame are redefined
diff --git a/python/python/tests/test_context.py b/python/python/tests/test_context.py
index bb92eb24c..047c8986e 100644
--- a/python/python/tests/test_context.py
+++ b/python/python/tests/test_context.py
@@ -16,6 +16,7 @@
 # under the License.
 
 from ballista import BallistaSessionContext, setup_test_cluster
+from ballista.extension import DataFrame, DistributedDataFrame, SessionContext
 from datafusion import col, lit
 import pytest
 import pyarrow as pa
@@ -138,3 +139,31 @@ def test_write_json(ctx, tmp_path):
     df.write_json(out_dir)
     json_files = list((tmp_path / "out").glob("*.json"))
     assert len(json_files) > 0
+
+
+def _assert_dataframe_returning_methods_wrapped(base_cls, sub_cls):
+    should_be_wrapped = {
+        name
+        for name, val in base_cls.__dict__.items()
+        if callable(val)
+        and not name.startswith("__")
+        and val.__annotations__.get("return") == DataFrame.__name__
+    }
+
+    assert should_be_wrapped
+    for name in should_be_wrapped:
+        assert name in sub_cls.__dict__, f"{name} not found in {sub_cls.__name__}"
+        assert callable(sub_cls.__dict__[name]), (
+            f"{name} is not callable in {sub_cls.__name__}"
+        )
+        assert sub_cls.__dict__[name] is not base_cls.__dict__[name], (
+            f"{name} was not replaced in {sub_cls.__name__}"
+        )
+
+
+def test_distributed_dataframe_wraps_dataframe_returning_methods():
+    _assert_dataframe_returning_methods_wrapped(DataFrame, DistributedDataFrame)
+
+
+def test_ballista_session_context_wraps_dataframe_returning_methods():
+    _assert_dataframe_returning_methods_wrapped(SessionContext, BallistaSessionContext)


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org