Feature/jupyter notebook support by sandugood · Pull Request #1513 · apache/datafusion-ballista

sandugood · 2026-03-16T20:10:10Z

Which issue does this PR close?

Refactoring of PR #1430

Rationale for this change

Latest commit to the mentioned PR was in January. Refactoring felt like a nice thing.

What changes are included in this PR?

Fix of both @line_magic and @cell_magic decorators. Combined them under the @line_cell_magic decorator
Nicely HTML-formatted output of cells. Using display() func and also printing in a loop where applicable

Are there any user-facing changes?

No.

…s and examples This PR implements all items from the checklist in issue apache#1398: ## Implementation Checklist - [x] Add example .ipynb notebooks to python/examples/ - getting_started.ipynb - Basic connection and queries - dataframe_api.ipynb - DataFrame transformations - distributed_queries.ipynb - Multi-stage query examples - [x] Document notebook support in Python README - Added comprehensive Jupyter section with examples - [x] Create ballista.jupyter module with magic commands - Full implementation with BallistaMagics class - [x] Add %ballista connect/status/tables/schema line magics - connect: Connect to Ballista cluster - status: Show connection status - tables: List registered tables - schema: Show table schema - disconnect: Disconnect from cluster - history: Show query history - [x] Add %%sql cell magic - Line magic for single-line queries - Cell magic for multi-line queries - Variable assignment support - --no-display and --limit options - [x] Add explain_visual() method for query plan rendering - Generates DOT/SVG visualization - Supports Jupyter _repr_html_ - Fallback when graphviz not installed - [x] Add progress indicator support for long-running queries - collect_with_progress() method - Callback support for custom progress handling - Jupyter-aware display - [x] Consider JupySQL integration - Documented as alternative in README ## Additional Features - ExecutionPlanVisualization class for plan rendering - tables() method on BallistaSessionContext - Optional jupyter dependency in pyproject.toml - Comprehensive test coverage (45 tests passing) Closes apache#1398

martin-g · 2026-03-16T21:15:08Z

python/python/ballista/jupyter.py

+            if not tables:
+                return "No tables registered.\n\nUse ctx.register_parquet() or ctx.register_csv() to register tables."
+            schema_count = len(tables.keys())
+            table_count = len(tables.values())


Suggested change

table_count = len(tables.values())

table_count = sum(len(v) for v in tables.values())

martin-g · 2026-03-16T21:18:02Z

python/python/ballista/jupyter.py

+        for row in lines:
+            print(row)
+
+    def _show_help(self) -> Optional[str]:


Is this function supposed to return something ?
Currently it just prints and implicitly returns None but https://github.com/apache/datafusion-ballista/pull/1513/changes#diff-db55df3b39944095b30c0a940790394d25c009baf49be871acd19d9eab2f30b8R90-R93 seems to expect Some

Instead of return of str in functions changed the return type to Optional[str]. Remade it to display() if the ipython is available and defaults to print otherwise.
Also changed the testcases to account for that.

martin-g · 2026-03-16T21:22:46Z

python/python/ballista/jupyter.py

+        if IPYTHON_AVAILABLE:
+            display(HTML(html))
+        else:
+            print("\n".join(status_lines))


Is this function supposed to return Some ?
https://github.com/apache/datafusion-ballista/pull/1513/changes#diff-db55df3b39944095b30c0a940790394d25c009baf49be871acd19d9eab2f30b8R68 tries to call .lower() on the implicit None

Same as previous comment

martin-g · 2026-03-16T21:24:08Z

python/python/ballista/jupyter.py

+
+        try:
+            # Query the table with LIMIT 0 to get schema without data
+            df = self.ctx.sql(f"SELECT * FROM {table_name} LIMIT 0")


Should table_name be sanitized ? Currently it could lead to SQL injection

We are considering a local development, I guess?
It is the same as:
%sql some SQL code

so it might be potentially any SQL code, executed by the data-scientist/analyst.

martin-g · 2026-03-16T21:25:38Z

python/python/ballista/jupyter.py

+        else:
+            def decorator(func):
+                return func
+            return decorator


No fallback for line_cell_magic. @line_cell_magic will fail if IPython is not available

martin-g · 2026-03-16T21:44:55Z

python/examples/getting_started.ipynb

+    "# ctx = BallistaSessionContext(\"df://your-scheduler:50050\")\n",
+    "\n",
+    "# host, port = setup_test_cluster()\n",
+    "host, port = \"localhost\", \"39431\"\n",


Is this a debug leftover ? IMO line 79 should be used

Yes, debug leftover. Fixed

martin-g · 2026-03-16T21:45:22Z

python/python/ballista/extension.py

+
+    def collect_with_progress(
+        self,
+        callback: Optional[callable] = None,


Suggested change

callback: Optional[callable] = None,

callback: Optional[Callable] = None,

Fixed. Thank you

martin-g · 2026-03-16T21:47:27Z

python/python/ballista/jupyter.py

+    LIMIT 10
+        """
+        for row in help_info.split("\n"):
+            print(row)


Suggested change

print(row)

print(row)

return help_info

martin-g · 2026-03-16T21:50:11Z

python/README.md

+## Jupyter Notebook Support
+
+PyBallista provides first-class Jupyter notebook support with SQL magic commands and rich HTML rendering.
+


Suggested change

Install Jupyter extras first:

```bash

pip install "ballista[jupyter]"

martin-g · 2026-03-16T21:52:23Z

python/python/ballista/extension.py

+        try:
+            catalog = self.catalog()
+            schema_names = list(catalog.schema_names())
+            if schema_names is not None:


A list is never None.

Suggested change

if schema_names is not None:

if schema_names:

Fixed. Thank you

sandugood · 2026-03-17T11:44:21Z

My main question was regarding the SQL-sanitizing. It seems to me that it is not needed (adding another layer of logic) in the current implementation. But I'm opened to suggestions and thank you for your review.

milenkovicm · 2026-03-18T10:02:03Z

My main question was regarding the SQL-sanitizing. It seems to me that it is not needed (adding another layer of logic) in the current implementation. But I'm opened to suggestions and thank you for your review.

hey @sandugood thanks for the contribution, i would agree with you, I don't think this is security issue as user can do whatever it wants anyway.

will have a look at the pr later in more details

littleKitchen and others added 4 commits January 31, 2026 00:45

Refactored magic commands

40df656

Changed .gitignore and added docstring

3592b2a

Removed previous ver of function

a188629

github-actions bot added the python label Mar 16, 2026

martin-g reviewed Mar 16, 2026

View reviewed changes

sandugood added 2 commits March 17, 2026 14:16

Changed the jupyter, extension and fixed the README.md

3117377

Added test cases

0b750c2

	table_count = len(tables.values())
	table_count = sum(len(v) for v in tables.values())

	callback: Optional[callable] = None,
	callback: Optional[Callable] = None,

		## Jupyter Notebook Support

		PyBallista provides first-class Jupyter notebook support with SQL magic commands and rich HTML rendering.

+Install Jupyter extras first:
+```bash
+pip install "ballista[jupyter]"

Conversation

sandugood commented Mar 16, 2026

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sandugood commented Mar 17, 2026

Uh oh!

milenkovicm commented Mar 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants