From 253c354039ac78949795c9e7b276a17516047eef Mon Sep 17 00:00:00 2001 From: Jogesh Anand Date: Mon, 15 Jun 2026 00:12:29 -0700 Subject: [PATCH 1/4] fix dataframe name - todo --- python/python/ballista/extension.py | 6 ++---- python/python/tests/test_context.py | 30 +++++++++++++++++++++++++++++ 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/python/python/ballista/extension.py b/python/python/ballista/extension.py index 90a3651bc6..5dbbf2a85e 100644 --- a/python/python/ballista/extension.py +++ b/python/python/ballista/extension.py @@ -72,12 +72,11 @@ def method_wrapper(*args, **kwargs): 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 @@ -112,12 +111,11 @@ def method_wrapper(*args, **kwargs): 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 bb92eb24c5..c2a53f09c3 100644 --- a/python/python/tests/test_context.py +++ b/python/python/tests/test_context.py @@ -138,3 +138,33 @@ 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 test_distributed_dataframe_wraps_dataframe_returning_methods(): + from ballista.extension import DistributedDataFrame, DataFrame + + should_be_wrapped = { + name + for name, val in DataFrame.__dict__.items() + if callable(val) + and not name.startswith("__") + and val.__annotations__.get("return") == DataFrame.__name__ + } + + assert should_be_wrapped + assert should_be_wrapped.issubset(DistributedDataFrame.__dict__) + + +def test_ballista_session_context_wraps_dataframe_returning_methods(): + from ballista.extension import BallistaSessionContext, SessionContext, DataFrame + + should_be_wrapped = { + name + for name, val in SessionContext.__dict__.items() + if callable(val) + and not name.startswith("__") + and val.__annotations__.get("return") == DataFrame.__name__ + } + + assert should_be_wrapped + assert should_be_wrapped.issubset(BallistaSessionContext.__dict__) From 49dd6de1ef54ff655ce27293d693f4aeb7742f21 Mon Sep 17 00:00:00 2001 From: Jogesh Anand Date: Mon, 15 Jun 2026 00:36:08 -0700 Subject: [PATCH 2/4] address review comments --- python/python/tests/test_context.py | 31 ++++++++++++++--------------- 1 file changed, 15 insertions(+), 16 deletions(-) diff --git a/python/python/tests/test_context.py b/python/python/tests/test_context.py index c2a53f09c3..047c8986e7 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 @@ -140,31 +141,29 @@ def test_write_json(ctx, tmp_path): assert len(json_files) > 0 -def test_distributed_dataframe_wraps_dataframe_returning_methods(): - from ballista.extension import DistributedDataFrame, DataFrame - +def _assert_dataframe_returning_methods_wrapped(base_cls, sub_cls): should_be_wrapped = { name - for name, val in DataFrame.__dict__.items() + 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 - assert should_be_wrapped.issubset(DistributedDataFrame.__dict__) + 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_ballista_session_context_wraps_dataframe_returning_methods(): - from ballista.extension import BallistaSessionContext, SessionContext, DataFrame +def test_distributed_dataframe_wraps_dataframe_returning_methods(): + _assert_dataframe_returning_methods_wrapped(DataFrame, DistributedDataFrame) - should_be_wrapped = { - name - for name, val in SessionContext.__dict__.items() - if callable(val) - and not name.startswith("__") - and val.__annotations__.get("return") == DataFrame.__name__ - } - assert should_be_wrapped - assert should_be_wrapped.issubset(BallistaSessionContext.__dict__) +def test_ballista_session_context_wraps_dataframe_returning_methods(): + _assert_dataframe_returning_methods_wrapped(SessionContext, BallistaSessionContext) From a24fd87373f40637efad93dcf673964d933de548 Mon Sep 17 00:00:00 2001 From: Jogesh Anand Date: Mon, 15 Jun 2026 00:43:26 -0700 Subject: [PATCH 3/4] upstream itertools and tower-http deps update --- python/Cargo.lock | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/Cargo.lock b/python/Cargo.lock index 8ccc4358b7..dc1c986dc8 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", ] From c9296a1a2772b07957b8eeff47880293d794f302 Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Mon, 15 Jun 2026 15:53:47 +0300 Subject: [PATCH 4/4] Add AI config files --- .cursor/rules.md | 5 +++++ .gemini/rules.md | 5 +++++ AGENTS.md | 5 +++++ CLAUDE.md | 5 +++++ 4 files changed, 20 insertions(+) create mode 100644 .cursor/rules.md create mode 100644 .gemini/rules.md create mode 100644 AGENTS.md create mode 100644 CLAUDE.md diff --git a/.cursor/rules.md b/.cursor/rules.md new file mode 100644 index 0000000000..00385e6f5f --- /dev/null +++ b/.cursor/rules.md @@ -0,0 +1,5 @@ +The AI agents should never link to any issue or a pull request +in any GitHub repository in the code reviews! + +The AI agents should not review AI agents' config files like CLAUDE.md or AGENTS.md! + diff --git a/.gemini/rules.md b/.gemini/rules.md new file mode 100644 index 0000000000..00385e6f5f --- /dev/null +++ b/.gemini/rules.md @@ -0,0 +1,5 @@ +The AI agents should never link to any issue or a pull request +in any GitHub repository in the code reviews! + +The AI agents should not review AI agents' config files like CLAUDE.md or AGENTS.md! + diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000000..00385e6f5f --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,5 @@ +The AI agents should never link to any issue or a pull request +in any GitHub repository in the code reviews! + +The AI agents should not review AI agents' config files like CLAUDE.md or AGENTS.md! + diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000000..00385e6f5f --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,5 @@ +The AI agents should never link to any issue or a pull request +in any GitHub repository in the code reviews! + +The AI agents should not review AI agents' config files like CLAUDE.md or AGENTS.md! +