feat: 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

+            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

+        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

+        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

+
+        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

+        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

+    "# 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

+
+    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

+    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

+## 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

+        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

milenkovicm · 2026-03-29T10:27:52Z

hey @sandugood there are few tests failing here, if you wish to progress with this PR please have a look

sandugood · 2026-03-29T10:29:09Z

Converting to draft, going to fix it

milenkovicm · 2026-03-29T10:42:48Z

thanks

sandugood · 2026-03-30T18:19:07Z

Fixed the pytest functions and ran through the ruff

milenkovicm · 2026-04-01T18:35:04Z

thanks @sandugood will have a look asap

milenkovicm

looks good, thanks @sandugood

sandugood · 2026-04-08T20:15:36Z

thank you for your review @milenkovicm

sandugood · 2026-04-09T08:27:33Z

I think it's ready for merge. any additional comments? @martin-g

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

sandugood marked this pull request as draft March 29, 2026 10:28

sandugood added 3 commits March 30, 2026 21:08

Fixed pipeline errors

969bf9c

Merge remote-tracking branch 'upstream/main' into pr-1430

d785d80

Add ruff cache to .gitignore

9bb8def

sandugood marked this pull request as ready for review March 30, 2026 18:16

milenkovicm approved these changes Apr 8, 2026

View reviewed changes

milenkovicm mentioned this pull request Apr 8, 2026

feat(python): Improve Jupyter notebook support with SQL magic commands and examples #1430

Closed

Merge branch 'apache:main' into feature/jupyter-notebook-support

8d0d129

sandugood changed the title ~~Feature/jupyter notebook support~~ feat: jupyter notebook support Apr 8, 2026

milenkovicm merged commit 6091b52 into apache:main Apr 9, 2026
18 checks passed

	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!

milenkovicm commented Mar 29, 2026

Uh oh!

sandugood commented Mar 29, 2026

Uh oh!

milenkovicm commented Mar 29, 2026

Uh oh!

sandugood commented Mar 30, 2026

Uh oh!

milenkovicm commented Apr 1, 2026

Uh oh!

milenkovicm left a comment

Choose a reason for hiding this comment

Uh oh!

sandugood commented Apr 8, 2026

Uh oh!

sandugood commented Apr 9, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants