Skip to content

1866: fix: use DataFrame.__name__ instead of hardcoded string in metaclass method wrapping #73

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
martin-augment wants to merge 4 commits into
base: main
Choose a base branch
from
Open

1866: fix: use DataFrame.__name__ instead of hardcoded string in metaclass method wrapping #73

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
253c354
fix dataframe name - todo
djanand Jun 15, 2026
49dd6de
address review comments
djanand Jun 15, 2026
a24fd87
upstream itertools and tower-http deps update
djanand Jun 15, 2026
c9296a1
Add AI config files
martin-g Jun 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .cursor/rules.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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!

5 changes: 5 additions & 0 deletions .gemini/rules.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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!

5 changes: 5 additions & 0 deletions AGENTS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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!

5 changes: 5 additions & 0 deletions CLAUDE.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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!

2 changes: 2 additions & 0 deletions python/Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

6 changes: 2 additions & 4 deletions python/python/ballista/extension.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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__

@augmentcode augmentcode Bot Jun 15, 2026 •
edited
Loading

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

python/python/ballista/extension.py:79: This __annotations__["return"] == DataFrame.__name__ check only matches string return annotations equal to the class name; if DataFusion ever emits an actual DataFrame type (or a qualified string), these methods would silently stop being wrapped and return plain DataFrame instances.

Severity: medium

Other Locations
  • python/python/ballista/extension.py:118
  • python/python/tests/test_context.py:150

Fix This in Augment

🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.

@martin-augment martin-augment Jun 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

value:annoying; category:bug; feedback: The Augment AI reviewer is not correct! If the return type is not "DataFrame" then the unit tests will fail and the respective changes will be made to fix the problem. Until then there is no problem to be fixed.

):
Comment on lines 76 to 80

@gemini-code-assist gemini-code-assist Bot Jun 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Accessing __annotations__ directly on a callable can raise an AttributeError if the callable (such as a built-in or C-extension method from PyO3) does not define it or if it is None. Using getattr(base_value, "__annotations__", None) or {} is a safer, more defensive approach to prevent potential runtime crashes.

Suggested change
if (
callable(base_value)
and not base_name.startswith("__")
and base_value.__annotations__.get("return") == "DataFrame"
and base_value.__annotations__.get("return") == DataFrame.__name__
):
if (
callable(base_value)
and not base_name.startswith("__")
and (getattr(base_value, "__annotations__", None) or {}).get("return") == DataFrame.__name__
):

@martin-augment martin-augment Jun 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

value:annoying; category:bug; feedback: The Gemini AI reviewer is not correct! The DataFrame class uses type hints and all its members are static, i.e. there are no dynamically added ones. Raising an AttributeError and failing the test would be better than silently omit the assertion for a new field that is not typed.

#
# functions which return DataFrame are redefined
Expand Down Expand Up @@ -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__
):
Comment on lines 115 to 119

@gemini-code-assist gemini-code-assist Bot Jun 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Accessing __annotations__ directly on a callable can raise an AttributeError if the callable (such as a built-in or C-extension method from PyO3) does not define it or if it is None. Using getattr(base_value, "__annotations__", None) or {} is a safer, more defensive approach to prevent potential runtime crashes.

Suggested change
if (
callable(base_value)
and not base_name.startswith("__")
and base_value.__annotations__.get("return") == "DataFrame"
and base_value.__annotations__.get("return") == DataFrame.__name__
):
if (
callable(base_value)
and not base_name.startswith("__")
and (getattr(base_value, "__annotations__", None) or {}).get("return") == DataFrame.__name__
):

@martin-augment martin-augment Jun 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

value:annoying; category:bug; feedback: The Gemini AI reviewer is not correct! The DataFrame class uses type hints and all its members are static, i.e. there are no dynamically added ones. Raising an AttributeError and failing the test would be better than silently omit the assertion for a new field that is not typed.

#
# functions which return DataFrame are redefined
Expand Down
29 changes: 29 additions & 0 deletions python/python/tests/test_context.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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__

@augmentcode augmentcode Bot Jun 15, 2026 •
edited
Loading

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

python/python/tests/test_context.py:150: _assert_dataframe_returning_methods_wrapped assumes every callable(val) has a __annotations__ dict; if any DataFusion method is a C-extension/builtin descriptor without __annotations__, this will raise AttributeError and fail the test for reasons unrelated to wrapping.

Severity: low

Fix This in Augment

🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.

@martin-augment martin-augment Jun 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

value:annoying; category:bug; feedback: The Augment AI reviewer is not correct! The DataFrame class uses type hints and all its members are static, i.e. there are no dynamically added ones. Raising an AttributeError and failing the test would be better than silently omit the assertion for a new field that is not typed.

}
Comment on lines +145 to +151

@gemini-code-assist gemini-code-assist Bot Jun 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Accessing __annotations__ directly on a callable can raise an AttributeError if the callable does not define it or if it is None. Using getattr(val, "__annotations__", None) or {} is a safer, more defensive approach to prevent potential runtime crashes.

Suggested change
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__
}
should_be_wrapped = {
name
for name, val in base_cls.__dict__.items()
if callable(val)
and not name.startswith("__")
and (getattr(val, "__annotations__", None) or {}).get("return") == DataFrame.__name__
}

@martin-augment martin-augment Jun 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

value:annoying; category:bug; feedback: The Gemini AI reviewer is not correct! The DataFrame class uses type hints and all its members are static, i.e. there are no dynamically added ones. Raising an AttributeError and failing the test would be better than silently omit the assertion for a new field that is not typed.


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)
Loading