You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "sandugood (via GitHub)" <gi...@apache.org> on 2026/03/16 20:10:13 UTC
[PR] Feature/jupyter notebook support [datafusion-ballista]
sandugood opened a new pull request, #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513
# Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
-->
Refactoring of PR #1430
# Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.
-->
Latest commit to the mentioned PR was in January. Refactoring felt like a nice thing.
# What changes are included in this PR?
<!--
There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
-->
1. Fix of both @line_magic and @cell_magic decorators. Combined them under the @line_cell_magic decorator
2. Nicely HTML-formatted output of cells. Using display() func and also printing in a loop where applicable
# Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
-->
No.
<!--
If there are any breaking changes to public APIs, please add the `api change` label.
-->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feature/jupyter notebook support [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on code in PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513#discussion_r2946190395
##########
python/python/ballista/jupyter.py:
##########
@@ -0,0 +1,493 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Ballista Jupyter Magic Commands.
+
+This module provides IPython magic commands for interacting with Ballista
+clusters directly from Jupyter notebooks.
+
+Usage:
+ %load_ext ballista.jupyter
+
+ # Connect to a Ballista cluster
+ %ballista connect df://localhost:50050
+
+ # Check connection status
+ %ballista status
+
+ # List registered tables
+ %ballista tables
+
+ # Show schema for a table
+ %ballista schema my_table
+
+ # Execute a simple SQL query (line magic)
+ %sql SELECT COUNT(*) FROM my_table
+
+ # Execute a complex SQL query (cell magic)
+ %%sql
+ SELECT
+ customer_id,
+ SUM(amount) as total
+ FROM orders
+ GROUP BY customer_id
+ ORDER BY total DESC
+ LIMIT 10
+"""
+
+from typing import Optional, List, Dict, Any
+import time
+import threading
+
+try:
+ from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic
+ from IPython.display import display, HTML
+
+ IPYTHON_AVAILABLE = True
+except ImportError:
+ IPYTHON_AVAILABLE = False
+
+ # Provide stub classes for when IPython is not available
+ class Magics:
+ def __init__(self, shell=None):
+ self.shell = shell
+
+ def magics_class(cls):
+ return cls
+
+ def line_magic(name_or_func):
+ """Stub line_magic decorator for when IPython is not available."""
+ # Handle both @line_magic and @line_magic("name") usage
+ if callable(name_or_func):
+ # Used as @line_magic without arguments
+ return name_or_func
+ else:
+ # Used as @line_magic("name") with arguments
+ def decorator(func):
+ return func
+ return decorator
+
+ def cell_magic(name_or_func):
+ """Stub cell_magic decorator for when IPython is not available."""
+ # Handle both @cell_magic and @cell_magic("name") usage
+ if callable(name_or_func):
+ return name_or_func
+ else:
+ def decorator(func):
+ return func
+ return decorator
+
+
+from .extension import BallistaSessionContext, DistributedDataFrame
+
+
+class BallistaConnectionError(Exception):
+ """Raised when not connected to a Ballista cluster."""
+
+ pass
+
+
+@magics_class
+class BallistaMagics(Magics):
+ """
+ IPython magic commands for Ballista.
+
+ Provides convenient commands for connecting to Ballista clusters,
+ executing SQL queries, and exploring table schemas.
+ """
+
+ def __init__(self, shell=None):
+ super().__init__(shell)
+ self._ctx: Optional[BallistaSessionContext] = None
+ self._address: Optional[str] = None
+ self._last_result: Optional[DistributedDataFrame] = None
+ self._query_history: List[Dict[str, Any]] = []
+
+ @property
+ def ctx(self) -> BallistaSessionContext:
+ """Get the current context, raising an error if not connected."""
+ if self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ return self._ctx
+
+ @property
+ def is_connected(self) -> bool:
+ """Check if connected to a Ballista cluster."""
+ return self._ctx is not None
+
+ @line_magic
+ def ballista(self, line: str) -> Optional[str]:
+ """
+ Ballista management commands.
+
+ Usage:
+ %ballista connect df://localhost:50050 - Connect to cluster
+ %ballista status - Show connection status
+ %ballista tables - List registered tables
+ %ballista schema <table> - Show table schema
+ %ballista disconnect - Disconnect from cluster
+ %ballista history - Show query history
+
+ Examples:
+ %ballista connect df://localhost:50050
+ %ballista tables
+ %ballista schema orders
+ """
+ parts = line.strip().split(maxsplit=1)
+ if not parts:
+ return self._show_help()
+
+ cmd = parts[0].lower()
+ args = parts[1] if len(parts) > 1 else ""
+
+ if cmd == "connect":
+ return self._connect(args)
+ elif cmd == "status":
+ return self._status()
+ elif cmd == "tables":
+ return self._tables()
+ elif cmd == "schema":
+ return self._schema(args)
+ elif cmd == "disconnect":
+ return self._disconnect()
+ elif cmd == "history":
+ return self._show_history()
+ elif cmd == "help":
+ return self._show_help()
+ else:
+ return f"Unknown command: {cmd}. Use '%ballista help' for available commands."
+
+ @line_magic
+ def register(self, line: str) -> Optional[str]:
+ """Register a new table"""
+ if not line:
+ return "You should provide file extension and table name to register"
+ elif self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ else:
+ args = line.strip().split()
+ file_type = args[0]
+ if len(args) < 2:
+ return f"You should provide table name for this .{file_type} file"
+ table_name = args[1]
+ if len(args) < 3:
+ return f"You should provide path to your file"
+ file_name = args[2]
+
+ if file_type == "parquet":
+ self._ctx.register_parquet(table_name, file_name)
+ elif file_type == "csv":
+ self._ctx.register_csv(table_name, file_name)
+ else:
+ raise NotImplemented("Currently not supporting the inserted file format")
+
+ @line_cell_magic
+ def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]:
+ """
+ Execute a SQL query (both line and cell magic).
+
+ Two cases possible: with cell or without cell
+
+ Examples:
+ 1. Without a cell (line_magic)
+ %sql SELECT * FROM test_table
+ 2. With a cell (cell_magic)
+ %%sql --no-display
+ SELECT
+ id,
+ bool_col,
+ tinyint_col
+ FROM test_data_v1
+ WHERE id > 2
+ ORDER BY id
+ LIMIT 5
+ """
+ if not cell:
+ return self._execute_sql(line.strip()) if line.strip() else None
+ else:
+ args = line.strip().split()
+ display_results = True
+ limit = 50
+ var_name = None
+
+ i = 0
+ while i < len(args):
+ if args[i] == "--no-display":
+ display_results = False
+ elif args[i] == "--limit" and i + 1 < len(args):
+ try:
+ limit = int(args[i + 1])
+ i += 1
+ except ValueError:
+ pass
+ elif not args[i].startswith("--"):
+ var_name = args[i]
+ i += 1
+
+ query = cell.strip()
+ if not query:
+ return None
+
+ result = self._execute_sql(query, display_results=display_results, limit=limit)
+
+ # Store in user namespace if variable name provided
+ if var_name and self.shell is not None:
+ self.shell.user_ns[var_name] = result
+
+ return result
+
+ def _connect(self, address: str) -> Optional[str]:
+ """Connect to a Ballista cluster."""
+ if not address:
+ return "Usage: %ballista connect df://host:port"
+
+ # Normalize address
+ if not address.startswith("df://"):
+ address = f"df://{address}"
+
+ try:
+ self._ctx = BallistaSessionContext(address)
+ self._address = address
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Connected to Ballista cluster at {address}"))
+ else:
+ return f"✓ Connected to Ballista cluster at {address}"
+ except Exception as e:
+ self._ctx = None
+ self._address = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✗ Failed to connect to {address}: {e}"))
+ else:
+ return f"✗ Failed to connect to {address}: {e}"
+
+ def _disconnect(self) -> Optional[str]:
+ """Disconnect from the Ballista cluster."""
+ if not self.is_connected:
+ return "Not connected to any cluster."
+
+ address = self._address
+ self._ctx = None
+ self._address = None
+ self._last_result = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Disconnected from {address}"))
+ else:
+ return f"✓ Disconnected from {address}"
+
+ def _status(self) -> Optional[str]:
+ """Show connection status."""
+ if not self.is_connected:
+ return "Status: Not connected\n\nUse '%ballista connect df://host:port' to connect."
+
+ status_lines = [
+ f"Status: Connected",
+ f"Address: {self._address}",
+ f"Session ID: {self._ctx.session_id}",
+ f"Queries executed: {len(self._query_history)}",
+ ]
+
+ if self._last_result is not None:
+ status_lines.append("Last result: Available (access via '_' or '_last_result')")
+
+ def _format_html_status_output(line: str) -> str:
+ name, value = line.split(":", 1)
+ return f"<p><strong>{name}:</strong> {value.strip()}</p>"
+
+ html = "".join(_format_html_status_output(line) for line in status_lines)
+ if IPYTHON_AVAILABLE:
+ display(HTML(html))
+ else:
+ print("\n".join(status_lines))
+
+ def _tables(self) -> Optional[str]:
+ """List all registered tables."""
+ try:
+ # Get table names from the catalog and their respective schemas
+ tables = self.ctx.get_tables()
+ 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())
+ # Build a nice table display (HTML-formatted if applicable)
+ lines = [
+ {"content": f"Total: {table_count} table(s) in {schema_count} schema(s)", "is_info": True},
+ {"content": "Registered tables:", "is_info": True},
+ *[{"content": f"Schema: {schema_name}. Tables: {", ".join(list(table_names))}", "is_info": False}
Review Comment:
Fixed
##########
python/python/ballista/jupyter.py:
##########
@@ -0,0 +1,493 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Ballista Jupyter Magic Commands.
+
+This module provides IPython magic commands for interacting with Ballista
+clusters directly from Jupyter notebooks.
+
+Usage:
+ %load_ext ballista.jupyter
+
+ # Connect to a Ballista cluster
+ %ballista connect df://localhost:50050
+
+ # Check connection status
+ %ballista status
+
+ # List registered tables
+ %ballista tables
+
+ # Show schema for a table
+ %ballista schema my_table
+
+ # Execute a simple SQL query (line magic)
+ %sql SELECT COUNT(*) FROM my_table
+
+ # Execute a complex SQL query (cell magic)
+ %%sql
+ SELECT
+ customer_id,
+ SUM(amount) as total
+ FROM orders
+ GROUP BY customer_id
+ ORDER BY total DESC
+ LIMIT 10
+"""
+
+from typing import Optional, List, Dict, Any
+import time
+import threading
+
+try:
+ from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic
+ from IPython.display import display, HTML
+
+ IPYTHON_AVAILABLE = True
+except ImportError:
+ IPYTHON_AVAILABLE = False
+
+ # Provide stub classes for when IPython is not available
+ class Magics:
+ def __init__(self, shell=None):
+ self.shell = shell
+
+ def magics_class(cls):
+ return cls
+
+ def line_magic(name_or_func):
+ """Stub line_magic decorator for when IPython is not available."""
+ # Handle both @line_magic and @line_magic("name") usage
+ if callable(name_or_func):
+ # Used as @line_magic without arguments
+ return name_or_func
+ else:
+ # Used as @line_magic("name") with arguments
+ def decorator(func):
+ return func
+ return decorator
+
+ def cell_magic(name_or_func):
+ """Stub cell_magic decorator for when IPython is not available."""
+ # Handle both @cell_magic and @cell_magic("name") usage
+ if callable(name_or_func):
+ return name_or_func
+ else:
+ def decorator(func):
+ return func
+ return decorator
+
+
+from .extension import BallistaSessionContext, DistributedDataFrame
+
+
+class BallistaConnectionError(Exception):
+ """Raised when not connected to a Ballista cluster."""
+
+ pass
+
+
+@magics_class
+class BallistaMagics(Magics):
+ """
+ IPython magic commands for Ballista.
+
+ Provides convenient commands for connecting to Ballista clusters,
+ executing SQL queries, and exploring table schemas.
+ """
+
+ def __init__(self, shell=None):
+ super().__init__(shell)
+ self._ctx: Optional[BallistaSessionContext] = None
+ self._address: Optional[str] = None
+ self._last_result: Optional[DistributedDataFrame] = None
+ self._query_history: List[Dict[str, Any]] = []
+
+ @property
+ def ctx(self) -> BallistaSessionContext:
+ """Get the current context, raising an error if not connected."""
+ if self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ return self._ctx
+
+ @property
+ def is_connected(self) -> bool:
+ """Check if connected to a Ballista cluster."""
+ return self._ctx is not None
+
+ @line_magic
+ def ballista(self, line: str) -> Optional[str]:
+ """
+ Ballista management commands.
+
+ Usage:
+ %ballista connect df://localhost:50050 - Connect to cluster
+ %ballista status - Show connection status
+ %ballista tables - List registered tables
+ %ballista schema <table> - Show table schema
+ %ballista disconnect - Disconnect from cluster
+ %ballista history - Show query history
+
+ Examples:
+ %ballista connect df://localhost:50050
+ %ballista tables
+ %ballista schema orders
+ """
+ parts = line.strip().split(maxsplit=1)
+ if not parts:
+ return self._show_help()
+
+ cmd = parts[0].lower()
+ args = parts[1] if len(parts) > 1 else ""
+
+ if cmd == "connect":
+ return self._connect(args)
+ elif cmd == "status":
+ return self._status()
+ elif cmd == "tables":
+ return self._tables()
+ elif cmd == "schema":
+ return self._schema(args)
+ elif cmd == "disconnect":
+ return self._disconnect()
+ elif cmd == "history":
+ return self._show_history()
+ elif cmd == "help":
+ return self._show_help()
+ else:
+ return f"Unknown command: {cmd}. Use '%ballista help' for available commands."
+
+ @line_magic
+ def register(self, line: str) -> Optional[str]:
+ """Register a new table"""
+ if not line:
+ return "You should provide file extension and table name to register"
+ elif self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ else:
+ args = line.strip().split()
+ file_type = args[0]
+ if len(args) < 2:
+ return f"You should provide table name for this .{file_type} file"
+ table_name = args[1]
+ if len(args) < 3:
+ return f"You should provide path to your file"
+ file_name = args[2]
+
+ if file_type == "parquet":
+ self._ctx.register_parquet(table_name, file_name)
+ elif file_type == "csv":
+ self._ctx.register_csv(table_name, file_name)
+ else:
+ raise NotImplemented("Currently not supporting the inserted file format")
+
+ @line_cell_magic
+ def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]:
+ """
+ Execute a SQL query (both line and cell magic).
+
+ Two cases possible: with cell or without cell
+
+ Examples:
+ 1. Without a cell (line_magic)
+ %sql SELECT * FROM test_table
+ 2. With a cell (cell_magic)
+ %%sql --no-display
+ SELECT
+ id,
+ bool_col,
+ tinyint_col
+ FROM test_data_v1
+ WHERE id > 2
+ ORDER BY id
+ LIMIT 5
+ """
+ if not cell:
+ return self._execute_sql(line.strip()) if line.strip() else None
+ else:
+ args = line.strip().split()
+ display_results = True
+ limit = 50
+ var_name = None
+
+ i = 0
+ while i < len(args):
+ if args[i] == "--no-display":
+ display_results = False
+ elif args[i] == "--limit" and i + 1 < len(args):
+ try:
+ limit = int(args[i + 1])
+ i += 1
+ except ValueError:
+ pass
+ elif not args[i].startswith("--"):
+ var_name = args[i]
+ i += 1
+
+ query = cell.strip()
+ if not query:
+ return None
+
+ result = self._execute_sql(query, display_results=display_results, limit=limit)
+
+ # Store in user namespace if variable name provided
+ if var_name and self.shell is not None:
+ self.shell.user_ns[var_name] = result
+
+ return result
+
+ def _connect(self, address: str) -> Optional[str]:
+ """Connect to a Ballista cluster."""
+ if not address:
+ return "Usage: %ballista connect df://host:port"
+
+ # Normalize address
+ if not address.startswith("df://"):
+ address = f"df://{address}"
+
+ try:
+ self._ctx = BallistaSessionContext(address)
+ self._address = address
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Connected to Ballista cluster at {address}"))
+ else:
+ return f"✓ Connected to Ballista cluster at {address}"
+ except Exception as e:
+ self._ctx = None
+ self._address = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✗ Failed to connect to {address}: {e}"))
+ else:
+ return f"✗ Failed to connect to {address}: {e}"
+
+ def _disconnect(self) -> Optional[str]:
+ """Disconnect from the Ballista cluster."""
+ if not self.is_connected:
+ return "Not connected to any cluster."
+
+ address = self._address
+ self._ctx = None
+ self._address = None
+ self._last_result = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Disconnected from {address}"))
+ else:
+ return f"✓ Disconnected from {address}"
+
+ def _status(self) -> Optional[str]:
+ """Show connection status."""
+ if not self.is_connected:
+ return "Status: Not connected\n\nUse '%ballista connect df://host:port' to connect."
+
+ status_lines = [
+ f"Status: Connected",
+ f"Address: {self._address}",
+ f"Session ID: {self._ctx.session_id}",
+ f"Queries executed: {len(self._query_history)}",
+ ]
+
+ if self._last_result is not None:
+ status_lines.append("Last result: Available (access via '_' or '_last_result')")
+
+ def _format_html_status_output(line: str) -> str:
+ name, value = line.split(":", 1)
+ return f"<p><strong>{name}:</strong> {value.strip()}</p>"
+
+ html = "".join(_format_html_status_output(line) for line in status_lines)
+ if IPYTHON_AVAILABLE:
+ display(HTML(html))
+ else:
+ print("\n".join(status_lines))
+
+ def _tables(self) -> Optional[str]:
+ """List all registered tables."""
+ try:
+ # Get table names from the catalog and their respective schemas
+ tables = self.ctx.get_tables()
+ 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())
+ # Build a nice table display (HTML-formatted if applicable)
+ lines = [
+ {"content": f"Total: {table_count} table(s) in {schema_count} schema(s)", "is_info": True},
+ {"content": "Registered tables:", "is_info": True},
+ *[{"content": f"Schema: {schema_name}. Tables: {", ".join(list(table_names))}", "is_info": False}
+ for schema_name, table_names in tables.items()]
+ ]
+
+ def _format_html_tables_output(line: str, is_info: bool = False) -> str:
+ if is_info:
+ return f"<p><strong>{line}<strong></p>"
+ else:
+ return f"<p><pre><i>{line}</i></pre></p>"
+
+ if IPYTHON_AVAILABLE:
+ display(
+ HTML("".join(_format_html_tables_output(val["content"], val["is_info"]) for val in lines))
+ )
+ else:
+ return "".join(val["content"] for val in lines)
+ except Exception as e:
+ return f"Error listing tables: {e}"
+
+ def _schema(self, table_name: str) -> Optional[str]:
+ """Show schema for a table."""
+ if not table_name:
+ return "Usage: %ballista schema <table_name>"
+
+ try:
+ # Query the table with LIMIT 0 to get schema without data
+ df = self.ctx.sql(f"SELECT * FROM {table_name} LIMIT 0")
+ schema = df.schema()
+
+ lines = [f"Schema for '{table_name}':", "-" * 50]
+ for field in schema:
+ nullable = "NULL" if field.nullable else "NOT NULL"
+ lines.append(f" {field.name:20} {str(field.type):15} {nullable}")
+ lines.append("-" * 50)
+ lines.append(f"Total: {len(schema)} column(s)")
+
+ return "\n".join(lines)
+ except Exception as e:
+ return f"Error getting schema for '{table_name}': {e}"
+
+ def _execute_sql(
+ self,
+ query: str,
+ display_results: bool = True,
+ limit: int = 50,
Review Comment:
Deleted from function args
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feature/jupyter notebook support [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on code in PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513#discussion_r2946175063
##########
python/python/ballista/jupyter.py:
##########
@@ -0,0 +1,493 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Ballista Jupyter Magic Commands.
+
+This module provides IPython magic commands for interacting with Ballista
+clusters directly from Jupyter notebooks.
+
+Usage:
+ %load_ext ballista.jupyter
+
+ # Connect to a Ballista cluster
+ %ballista connect df://localhost:50050
+
+ # Check connection status
+ %ballista status
+
+ # List registered tables
+ %ballista tables
+
+ # Show schema for a table
+ %ballista schema my_table
+
+ # Execute a simple SQL query (line magic)
+ %sql SELECT COUNT(*) FROM my_table
+
+ # Execute a complex SQL query (cell magic)
+ %%sql
+ SELECT
+ customer_id,
+ SUM(amount) as total
+ FROM orders
+ GROUP BY customer_id
+ ORDER BY total DESC
+ LIMIT 10
+"""
+
+from typing import Optional, List, Dict, Any
+import time
+import threading
+
+try:
+ from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic
+ from IPython.display import display, HTML
+
+ IPYTHON_AVAILABLE = True
+except ImportError:
+ IPYTHON_AVAILABLE = False
+
+ # Provide stub classes for when IPython is not available
+ class Magics:
+ def __init__(self, shell=None):
+ self.shell = shell
+
+ def magics_class(cls):
+ return cls
+
+ def line_magic(name_or_func):
+ """Stub line_magic decorator for when IPython is not available."""
+ # Handle both @line_magic and @line_magic("name") usage
+ if callable(name_or_func):
+ # Used as @line_magic without arguments
+ return name_or_func
+ else:
+ # Used as @line_magic("name") with arguments
+ def decorator(func):
+ return func
+ return decorator
+
+ def cell_magic(name_or_func):
+ """Stub cell_magic decorator for when IPython is not available."""
+ # Handle both @cell_magic and @cell_magic("name") usage
+ if callable(name_or_func):
+ return name_or_func
+ else:
+ def decorator(func):
+ return func
+ return decorator
+
+
+from .extension import BallistaSessionContext, DistributedDataFrame
+
+
+class BallistaConnectionError(Exception):
+ """Raised when not connected to a Ballista cluster."""
+
+ pass
+
+
+@magics_class
+class BallistaMagics(Magics):
+ """
+ IPython magic commands for Ballista.
+
+ Provides convenient commands for connecting to Ballista clusters,
+ executing SQL queries, and exploring table schemas.
+ """
+
+ def __init__(self, shell=None):
+ super().__init__(shell)
+ self._ctx: Optional[BallistaSessionContext] = None
+ self._address: Optional[str] = None
+ self._last_result: Optional[DistributedDataFrame] = None
+ self._query_history: List[Dict[str, Any]] = []
+
+ @property
+ def ctx(self) -> BallistaSessionContext:
+ """Get the current context, raising an error if not connected."""
+ if self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ return self._ctx
+
+ @property
+ def is_connected(self) -> bool:
+ """Check if connected to a Ballista cluster."""
+ return self._ctx is not None
+
+ @line_magic
+ def ballista(self, line: str) -> Optional[str]:
+ """
+ Ballista management commands.
+
+ Usage:
+ %ballista connect df://localhost:50050 - Connect to cluster
+ %ballista status - Show connection status
+ %ballista tables - List registered tables
+ %ballista schema <table> - Show table schema
+ %ballista disconnect - Disconnect from cluster
+ %ballista history - Show query history
+
+ Examples:
+ %ballista connect df://localhost:50050
+ %ballista tables
+ %ballista schema orders
+ """
+ parts = line.strip().split(maxsplit=1)
+ if not parts:
+ return self._show_help()
+
+ cmd = parts[0].lower()
+ args = parts[1] if len(parts) > 1 else ""
+
+ if cmd == "connect":
+ return self._connect(args)
+ elif cmd == "status":
+ return self._status()
+ elif cmd == "tables":
+ return self._tables()
+ elif cmd == "schema":
+ return self._schema(args)
+ elif cmd == "disconnect":
+ return self._disconnect()
+ elif cmd == "history":
+ return self._show_history()
+ elif cmd == "help":
+ return self._show_help()
+ else:
+ return f"Unknown command: {cmd}. Use '%ballista help' for available commands."
+
+ @line_magic
+ def register(self, line: str) -> Optional[str]:
+ """Register a new table"""
+ if not line:
+ return "You should provide file extension and table name to register"
+ elif self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ else:
+ args = line.strip().split()
+ file_type = args[0]
+ if len(args) < 2:
+ return f"You should provide table name for this .{file_type} file"
+ table_name = args[1]
+ if len(args) < 3:
+ return f"You should provide path to your file"
+ file_name = args[2]
+
+ if file_type == "parquet":
+ self._ctx.register_parquet(table_name, file_name)
+ elif file_type == "csv":
+ self._ctx.register_csv(table_name, file_name)
+ else:
+ raise NotImplemented("Currently not supporting the inserted file format")
+
+ @line_cell_magic
+ def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]:
+ """
+ Execute a SQL query (both line and cell magic).
+
+ Two cases possible: with cell or without cell
+
+ Examples:
+ 1. Without a cell (line_magic)
+ %sql SELECT * FROM test_table
+ 2. With a cell (cell_magic)
+ %%sql --no-display
+ SELECT
+ id,
+ bool_col,
+ tinyint_col
+ FROM test_data_v1
+ WHERE id > 2
+ ORDER BY id
+ LIMIT 5
+ """
+ if not cell:
+ return self._execute_sql(line.strip()) if line.strip() else None
+ else:
+ args = line.strip().split()
+ display_results = True
+ limit = 50
+ var_name = None
+
+ i = 0
+ while i < len(args):
+ if args[i] == "--no-display":
+ display_results = False
+ elif args[i] == "--limit" and i + 1 < len(args):
+ try:
+ limit = int(args[i + 1])
+ i += 1
+ except ValueError:
+ pass
+ elif not args[i].startswith("--"):
+ var_name = args[i]
+ i += 1
+
+ query = cell.strip()
+ if not query:
+ return None
+
+ result = self._execute_sql(query, display_results=display_results, limit=limit)
+
+ # Store in user namespace if variable name provided
+ if var_name and self.shell is not None:
+ self.shell.user_ns[var_name] = result
+
+ return result
+
+ def _connect(self, address: str) -> Optional[str]:
+ """Connect to a Ballista cluster."""
+ if not address:
+ return "Usage: %ballista connect df://host:port"
+
+ # Normalize address
+ if not address.startswith("df://"):
+ address = f"df://{address}"
+
+ try:
+ self._ctx = BallistaSessionContext(address)
+ self._address = address
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Connected to Ballista cluster at {address}"))
+ else:
+ return f"✓ Connected to Ballista cluster at {address}"
+ except Exception as e:
+ self._ctx = None
+ self._address = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✗ Failed to connect to {address}: {e}"))
+ else:
+ return f"✗ Failed to connect to {address}: {e}"
+
+ def _disconnect(self) -> Optional[str]:
+ """Disconnect from the Ballista cluster."""
+ if not self.is_connected:
+ return "Not connected to any cluster."
+
+ address = self._address
+ self._ctx = None
+ self._address = None
+ self._last_result = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Disconnected from {address}"))
+ else:
+ return f"✓ Disconnected from {address}"
+
+ def _status(self) -> Optional[str]:
+ """Show connection status."""
+ if not self.is_connected:
+ return "Status: Not connected\n\nUse '%ballista connect df://host:port' to connect."
+
+ status_lines = [
+ f"Status: Connected",
+ f"Address: {self._address}",
+ f"Session ID: {self._ctx.session_id}",
+ f"Queries executed: {len(self._query_history)}",
+ ]
+
+ if self._last_result is not None:
+ status_lines.append("Last result: Available (access via '_' or '_last_result')")
+
+ def _format_html_status_output(line: str) -> str:
+ name, value = line.split(":", 1)
+ return f"<p><strong>{name}:</strong> {value.strip()}</p>"
+
+ html = "".join(_format_html_status_output(line) for line in status_lines)
+ if IPYTHON_AVAILABLE:
+ display(HTML(html))
+ else:
+ print("\n".join(status_lines))
+
+ def _tables(self) -> Optional[str]:
+ """List all registered tables."""
+ try:
+ # Get table names from the catalog and their respective schemas
+ tables = self.ctx.get_tables()
+ 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())
+ # Build a nice table display (HTML-formatted if applicable)
+ lines = [
+ {"content": f"Total: {table_count} table(s) in {schema_count} schema(s)", "is_info": True},
+ {"content": "Registered tables:", "is_info": True},
+ *[{"content": f"Schema: {schema_name}. Tables: {", ".join(list(table_names))}", "is_info": False}
+ for schema_name, table_names in tables.items()]
+ ]
+
+ def _format_html_tables_output(line: str, is_info: bool = False) -> str:
+ if is_info:
+ return f"<p><strong>{line}<strong></p>"
+ else:
+ return f"<p><pre><i>{line}</i></pre></p>"
+
+ if IPYTHON_AVAILABLE:
+ display(
+ HTML("".join(_format_html_tables_output(val["content"], val["is_info"]) for val in lines))
+ )
+ else:
+ return "".join(val["content"] for val in lines)
+ except Exception as e:
+ return f"Error listing tables: {e}"
+
+ def _schema(self, table_name: str) -> Optional[str]:
+ """Show schema for a table."""
+ if not table_name:
+ return "Usage: %ballista schema <table_name>"
+
+ try:
+ # Query the table with LIMIT 0 to get schema without data
+ df = self.ctx.sql(f"SELECT * FROM {table_name} LIMIT 0")
+ schema = df.schema()
+
+ lines = [f"Schema for '{table_name}':", "-" * 50]
+ for field in schema:
+ nullable = "NULL" if field.nullable else "NOT NULL"
+ lines.append(f" {field.name:20} {str(field.type):15} {nullable}")
+ lines.append("-" * 50)
+ lines.append(f"Total: {len(schema)} column(s)")
+
+ return "\n".join(lines)
+ except Exception as e:
+ return f"Error getting schema for '{table_name}': {e}"
+
+ def _execute_sql(
+ self,
+ query: str,
+ display_results: bool = True,
+ limit: int = 50,
+ ) -> Optional[DistributedDataFrame]:
+ """Execute a SQL query and return the result."""
+ start_time = time.time()
+
+ try:
+ result = self.ctx.sql(query)
+ elapsed = time.time() - start_time
+
+ # Store result
+ self._last_result = result
+ if self.shell is not None and hasattr(self.shell, 'user_ns'):
+ self.shell.user_ns["_last_result"] = result
+
+ # Record in history
+ self._query_history.append(
+ {
+ "query": query,
+ "elapsed_seconds": elapsed,
+ "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
+ }
+ )
+
+ # Display if requested and in notebook environment
+ return result
+
+ except Exception as e:
+ error_msg = f"Query failed: {e}"
+ if IPYTHON_AVAILABLE:
+ try:
+ display(HTML(f'<div style="color: red; font-weight: bold;">{error_msg}</div>'))
+ except Exception:
+ print(error_msg)
+ else:
+ print(error_msg)
+ return None
+
+ def _show_history(self) -> Optional[str]:
+ """Show query history."""
+ if not self._query_history:
+ return "No queries executed yet."
+
+ lines = ["Query History:", "-" * 60]
+ for i, entry in enumerate(self._query_history[-10:], 1): # Last 10 queries
+ query_preview = entry["query"][:50] + "..." if len(entry["query"]) > 50 else entry["query"]
+ query_preview = query_preview.replace("\n", " ")
+ lines.append(f"{i}. [{entry['timestamp']}] ({entry['elapsed_seconds']:.2f}s)")
+ lines.append(f" {query_preview}")
+ lines.append("-" * 60)
+
+ for row in lines:
+ print(row)
+
+ def _show_help(self) -> Optional[str]:
Review Comment:
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.
##########
python/python/ballista/jupyter.py:
##########
@@ -0,0 +1,493 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Ballista Jupyter Magic Commands.
+
+This module provides IPython magic commands for interacting with Ballista
+clusters directly from Jupyter notebooks.
+
+Usage:
+ %load_ext ballista.jupyter
+
+ # Connect to a Ballista cluster
+ %ballista connect df://localhost:50050
+
+ # Check connection status
+ %ballista status
+
+ # List registered tables
+ %ballista tables
+
+ # Show schema for a table
+ %ballista schema my_table
+
+ # Execute a simple SQL query (line magic)
+ %sql SELECT COUNT(*) FROM my_table
+
+ # Execute a complex SQL query (cell magic)
+ %%sql
+ SELECT
+ customer_id,
+ SUM(amount) as total
+ FROM orders
+ GROUP BY customer_id
+ ORDER BY total DESC
+ LIMIT 10
+"""
+
+from typing import Optional, List, Dict, Any
+import time
+import threading
+
+try:
+ from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic
+ from IPython.display import display, HTML
+
+ IPYTHON_AVAILABLE = True
+except ImportError:
+ IPYTHON_AVAILABLE = False
+
+ # Provide stub classes for when IPython is not available
+ class Magics:
+ def __init__(self, shell=None):
+ self.shell = shell
+
+ def magics_class(cls):
+ return cls
+
+ def line_magic(name_or_func):
+ """Stub line_magic decorator for when IPython is not available."""
+ # Handle both @line_magic and @line_magic("name") usage
+ if callable(name_or_func):
+ # Used as @line_magic without arguments
+ return name_or_func
+ else:
+ # Used as @line_magic("name") with arguments
+ def decorator(func):
+ return func
+ return decorator
+
+ def cell_magic(name_or_func):
+ """Stub cell_magic decorator for when IPython is not available."""
+ # Handle both @cell_magic and @cell_magic("name") usage
+ if callable(name_or_func):
+ return name_or_func
+ else:
+ def decorator(func):
+ return func
+ return decorator
+
+
+from .extension import BallistaSessionContext, DistributedDataFrame
+
+
+class BallistaConnectionError(Exception):
+ """Raised when not connected to a Ballista cluster."""
+
+ pass
+
+
+@magics_class
+class BallistaMagics(Magics):
+ """
+ IPython magic commands for Ballista.
+
+ Provides convenient commands for connecting to Ballista clusters,
+ executing SQL queries, and exploring table schemas.
+ """
+
+ def __init__(self, shell=None):
+ super().__init__(shell)
+ self._ctx: Optional[BallistaSessionContext] = None
+ self._address: Optional[str] = None
+ self._last_result: Optional[DistributedDataFrame] = None
+ self._query_history: List[Dict[str, Any]] = []
+
+ @property
+ def ctx(self) -> BallistaSessionContext:
+ """Get the current context, raising an error if not connected."""
+ if self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ return self._ctx
+
+ @property
+ def is_connected(self) -> bool:
+ """Check if connected to a Ballista cluster."""
+ return self._ctx is not None
+
+ @line_magic
+ def ballista(self, line: str) -> Optional[str]:
+ """
+ Ballista management commands.
+
+ Usage:
+ %ballista connect df://localhost:50050 - Connect to cluster
+ %ballista status - Show connection status
+ %ballista tables - List registered tables
+ %ballista schema <table> - Show table schema
+ %ballista disconnect - Disconnect from cluster
+ %ballista history - Show query history
+
+ Examples:
+ %ballista connect df://localhost:50050
+ %ballista tables
+ %ballista schema orders
+ """
+ parts = line.strip().split(maxsplit=1)
+ if not parts:
+ return self._show_help()
+
+ cmd = parts[0].lower()
+ args = parts[1] if len(parts) > 1 else ""
+
+ if cmd == "connect":
+ return self._connect(args)
+ elif cmd == "status":
+ return self._status()
+ elif cmd == "tables":
+ return self._tables()
+ elif cmd == "schema":
+ return self._schema(args)
+ elif cmd == "disconnect":
+ return self._disconnect()
+ elif cmd == "history":
+ return self._show_history()
+ elif cmd == "help":
+ return self._show_help()
+ else:
+ return f"Unknown command: {cmd}. Use '%ballista help' for available commands."
+
+ @line_magic
+ def register(self, line: str) -> Optional[str]:
+ """Register a new table"""
+ if not line:
+ return "You should provide file extension and table name to register"
+ elif self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ else:
+ args = line.strip().split()
+ file_type = args[0]
+ if len(args) < 2:
+ return f"You should provide table name for this .{file_type} file"
+ table_name = args[1]
+ if len(args) < 3:
+ return f"You should provide path to your file"
+ file_name = args[2]
+
+ if file_type == "parquet":
+ self._ctx.register_parquet(table_name, file_name)
+ elif file_type == "csv":
+ self._ctx.register_csv(table_name, file_name)
+ else:
+ raise NotImplemented("Currently not supporting the inserted file format")
+
+ @line_cell_magic
+ def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]:
+ """
+ Execute a SQL query (both line and cell magic).
+
+ Two cases possible: with cell or without cell
+
+ Examples:
+ 1. Without a cell (line_magic)
+ %sql SELECT * FROM test_table
+ 2. With a cell (cell_magic)
+ %%sql --no-display
+ SELECT
+ id,
+ bool_col,
+ tinyint_col
+ FROM test_data_v1
+ WHERE id > 2
+ ORDER BY id
+ LIMIT 5
+ """
+ if not cell:
+ return self._execute_sql(line.strip()) if line.strip() else None
+ else:
+ args = line.strip().split()
+ display_results = True
+ limit = 50
+ var_name = None
+
+ i = 0
+ while i < len(args):
+ if args[i] == "--no-display":
+ display_results = False
+ elif args[i] == "--limit" and i + 1 < len(args):
+ try:
+ limit = int(args[i + 1])
+ i += 1
+ except ValueError:
+ pass
+ elif not args[i].startswith("--"):
+ var_name = args[i]
+ i += 1
+
+ query = cell.strip()
+ if not query:
+ return None
+
+ result = self._execute_sql(query, display_results=display_results, limit=limit)
+
+ # Store in user namespace if variable name provided
+ if var_name and self.shell is not None:
+ self.shell.user_ns[var_name] = result
+
+ return result
+
+ def _connect(self, address: str) -> Optional[str]:
+ """Connect to a Ballista cluster."""
+ if not address:
+ return "Usage: %ballista connect df://host:port"
+
+ # Normalize address
+ if not address.startswith("df://"):
+ address = f"df://{address}"
+
+ try:
+ self._ctx = BallistaSessionContext(address)
+ self._address = address
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Connected to Ballista cluster at {address}"))
+ else:
+ return f"✓ Connected to Ballista cluster at {address}"
+ except Exception as e:
+ self._ctx = None
+ self._address = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✗ Failed to connect to {address}: {e}"))
+ else:
+ return f"✗ Failed to connect to {address}: {e}"
+
+ def _disconnect(self) -> Optional[str]:
+ """Disconnect from the Ballista cluster."""
+ if not self.is_connected:
+ return "Not connected to any cluster."
+
+ address = self._address
+ self._ctx = None
+ self._address = None
+ self._last_result = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Disconnected from {address}"))
+ else:
+ return f"✓ Disconnected from {address}"
+
+ def _status(self) -> Optional[str]:
+ """Show connection status."""
+ if not self.is_connected:
+ return "Status: Not connected\n\nUse '%ballista connect df://host:port' to connect."
+
+ status_lines = [
+ f"Status: Connected",
+ f"Address: {self._address}",
+ f"Session ID: {self._ctx.session_id}",
+ f"Queries executed: {len(self._query_history)}",
+ ]
+
+ if self._last_result is not None:
+ status_lines.append("Last result: Available (access via '_' or '_last_result')")
+
+ def _format_html_status_output(line: str) -> str:
+ name, value = line.split(":", 1)
+ return f"<p><strong>{name}:</strong> {value.strip()}</p>"
+
+ html = "".join(_format_html_status_output(line) for line in status_lines)
+ if IPYTHON_AVAILABLE:
+ display(HTML(html))
+ else:
+ print("\n".join(status_lines))
Review Comment:
Same as previous comment
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feature/jupyter notebook support [datafusion-ballista]
Posted by "martin-g (via GitHub)" <gi...@apache.org>.
martin-g commented on code in PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513#discussion_r2943046549
##########
python/python/ballista/jupyter.py:
##########
@@ -0,0 +1,493 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Ballista Jupyter Magic Commands.
+
+This module provides IPython magic commands for interacting with Ballista
+clusters directly from Jupyter notebooks.
+
+Usage:
+ %load_ext ballista.jupyter
+
+ # Connect to a Ballista cluster
+ %ballista connect df://localhost:50050
+
+ # Check connection status
+ %ballista status
+
+ # List registered tables
+ %ballista tables
+
+ # Show schema for a table
+ %ballista schema my_table
+
+ # Execute a simple SQL query (line magic)
+ %sql SELECT COUNT(*) FROM my_table
+
+ # Execute a complex SQL query (cell magic)
+ %%sql
+ SELECT
+ customer_id,
+ SUM(amount) as total
+ FROM orders
+ GROUP BY customer_id
+ ORDER BY total DESC
+ LIMIT 10
+"""
+
+from typing import Optional, List, Dict, Any
+import time
+import threading
+
+try:
+ from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic
+ from IPython.display import display, HTML
+
+ IPYTHON_AVAILABLE = True
+except ImportError:
+ IPYTHON_AVAILABLE = False
+
+ # Provide stub classes for when IPython is not available
+ class Magics:
+ def __init__(self, shell=None):
+ self.shell = shell
+
+ def magics_class(cls):
+ return cls
+
+ def line_magic(name_or_func):
+ """Stub line_magic decorator for when IPython is not available."""
+ # Handle both @line_magic and @line_magic("name") usage
+ if callable(name_or_func):
+ # Used as @line_magic without arguments
+ return name_or_func
+ else:
+ # Used as @line_magic("name") with arguments
+ def decorator(func):
+ return func
+ return decorator
+
+ def cell_magic(name_or_func):
+ """Stub cell_magic decorator for when IPython is not available."""
+ # Handle both @cell_magic and @cell_magic("name") usage
+ if callable(name_or_func):
+ return name_or_func
+ else:
+ def decorator(func):
+ return func
+ return decorator
+
+
+from .extension import BallistaSessionContext, DistributedDataFrame
+
+
+class BallistaConnectionError(Exception):
+ """Raised when not connected to a Ballista cluster."""
+
+ pass
+
+
+@magics_class
+class BallistaMagics(Magics):
+ """
+ IPython magic commands for Ballista.
+
+ Provides convenient commands for connecting to Ballista clusters,
+ executing SQL queries, and exploring table schemas.
+ """
+
+ def __init__(self, shell=None):
+ super().__init__(shell)
+ self._ctx: Optional[BallistaSessionContext] = None
+ self._address: Optional[str] = None
+ self._last_result: Optional[DistributedDataFrame] = None
+ self._query_history: List[Dict[str, Any]] = []
+
+ @property
+ def ctx(self) -> BallistaSessionContext:
+ """Get the current context, raising an error if not connected."""
+ if self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ return self._ctx
+
+ @property
+ def is_connected(self) -> bool:
+ """Check if connected to a Ballista cluster."""
+ return self._ctx is not None
+
+ @line_magic
+ def ballista(self, line: str) -> Optional[str]:
+ """
+ Ballista management commands.
+
+ Usage:
+ %ballista connect df://localhost:50050 - Connect to cluster
+ %ballista status - Show connection status
+ %ballista tables - List registered tables
+ %ballista schema <table> - Show table schema
+ %ballista disconnect - Disconnect from cluster
+ %ballista history - Show query history
+
+ Examples:
+ %ballista connect df://localhost:50050
+ %ballista tables
+ %ballista schema orders
+ """
+ parts = line.strip().split(maxsplit=1)
+ if not parts:
+ return self._show_help()
+
+ cmd = parts[0].lower()
+ args = parts[1] if len(parts) > 1 else ""
+
+ if cmd == "connect":
+ return self._connect(args)
+ elif cmd == "status":
+ return self._status()
+ elif cmd == "tables":
+ return self._tables()
+ elif cmd == "schema":
+ return self._schema(args)
+ elif cmd == "disconnect":
+ return self._disconnect()
+ elif cmd == "history":
+ return self._show_history()
+ elif cmd == "help":
+ return self._show_help()
+ else:
+ return f"Unknown command: {cmd}. Use '%ballista help' for available commands."
+
+ @line_magic
+ def register(self, line: str) -> Optional[str]:
+ """Register a new table"""
+ if not line:
+ return "You should provide file extension and table name to register"
+ elif self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ else:
+ args = line.strip().split()
+ file_type = args[0]
+ if len(args) < 2:
+ return f"You should provide table name for this .{file_type} file"
+ table_name = args[1]
+ if len(args) < 3:
+ return f"You should provide path to your file"
+ file_name = args[2]
+
+ if file_type == "parquet":
+ self._ctx.register_parquet(table_name, file_name)
+ elif file_type == "csv":
+ self._ctx.register_csv(table_name, file_name)
+ else:
+ raise NotImplemented("Currently not supporting the inserted file format")
+
+ @line_cell_magic
+ def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]:
+ """
+ Execute a SQL query (both line and cell magic).
+
+ Two cases possible: with cell or without cell
+
+ Examples:
+ 1. Without a cell (line_magic)
+ %sql SELECT * FROM test_table
+ 2. With a cell (cell_magic)
+ %%sql --no-display
+ SELECT
+ id,
+ bool_col,
+ tinyint_col
+ FROM test_data_v1
+ WHERE id > 2
+ ORDER BY id
+ LIMIT 5
+ """
+ if not cell:
+ return self._execute_sql(line.strip()) if line.strip() else None
+ else:
+ args = line.strip().split()
+ display_results = True
+ limit = 50
+ var_name = None
+
+ i = 0
+ while i < len(args):
+ if args[i] == "--no-display":
+ display_results = False
+ elif args[i] == "--limit" and i + 1 < len(args):
+ try:
+ limit = int(args[i + 1])
+ i += 1
+ except ValueError:
+ pass
+ elif not args[i].startswith("--"):
+ var_name = args[i]
+ i += 1
+
+ query = cell.strip()
+ if not query:
+ return None
+
+ result = self._execute_sql(query, display_results=display_results, limit=limit)
+
+ # Store in user namespace if variable name provided
+ if var_name and self.shell is not None:
+ self.shell.user_ns[var_name] = result
+
+ return result
+
+ def _connect(self, address: str) -> Optional[str]:
+ """Connect to a Ballista cluster."""
+ if not address:
+ return "Usage: %ballista connect df://host:port"
+
+ # Normalize address
+ if not address.startswith("df://"):
+ address = f"df://{address}"
+
+ try:
+ self._ctx = BallistaSessionContext(address)
+ self._address = address
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Connected to Ballista cluster at {address}"))
+ else:
+ return f"✓ Connected to Ballista cluster at {address}"
+ except Exception as e:
+ self._ctx = None
+ self._address = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✗ Failed to connect to {address}: {e}"))
+ else:
+ return f"✗ Failed to connect to {address}: {e}"
+
+ def _disconnect(self) -> Optional[str]:
+ """Disconnect from the Ballista cluster."""
+ if not self.is_connected:
+ return "Not connected to any cluster."
+
+ address = self._address
+ self._ctx = None
+ self._address = None
+ self._last_result = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Disconnected from {address}"))
+ else:
+ return f"✓ Disconnected from {address}"
+
+ def _status(self) -> Optional[str]:
+ """Show connection status."""
+ if not self.is_connected:
+ return "Status: Not connected\n\nUse '%ballista connect df://host:port' to connect."
+
+ status_lines = [
+ f"Status: Connected",
+ f"Address: {self._address}",
+ f"Session ID: {self._ctx.session_id}",
+ f"Queries executed: {len(self._query_history)}",
+ ]
+
+ if self._last_result is not None:
+ status_lines.append("Last result: Available (access via '_' or '_last_result')")
+
+ def _format_html_status_output(line: str) -> str:
+ name, value = line.split(":", 1)
+ return f"<p><strong>{name}:</strong> {value.strip()}</p>"
+
+ html = "".join(_format_html_status_output(line) for line in status_lines)
+ if IPYTHON_AVAILABLE:
+ display(HTML(html))
+ else:
+ print("\n".join(status_lines))
Review Comment:
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
##########
python/python/ballista/extension.py:
##########
@@ -229,9 +231,319 @@ def write_parquet(
df = self._to_internal_df()
df.write_parquet(str(path), compression.value, compression_level)
+ def explain_visual(self, analyze: bool = False) -> "ExecutionPlanVisualization":
+ """
+ Generate a visual representation of the execution plan.
+
+ This method creates an SVG visualization of the query execution plan,
+ which can be displayed directly in Jupyter notebooks.
+
+ Args:
+ analyze: If True, includes runtime statistics from actual execution.
+
+ Returns:
+ ExecutionPlanVisualization: An object that renders as SVG in Jupyter.
+
+ Example:
+ >>> df = ctx.sql("SELECT * FROM orders WHERE amount > 100")
+ >>> df.explain_visual() # Displays SVG in notebook
+ >>> viz = df.explain_visual(analyze=True)
+ >>> viz.save("plan.svg") # Save to file
+ """
+ # Get the execution plan as a string representation
+ # Note: explain() prints but doesn't return a string, so we use logical_plan()
+ try:
+ plan = self.logical_plan()
+ plan_str = plan.display_indent()
+ except Exception:
+ # Fallback if logical_plan() fails
+ plan_str = "Unable to retrieve execution plan"
+ return ExecutionPlanVisualization(plan_str, analyze=analyze)
+
+ def collect_with_progress(
+ self,
+ callback: Optional[callable] = None,
+ poll_interval: float = 0.5,
+ ):
+ """
+ Collect results with progress indication.
+
+ For long-running queries, this method provides progress updates
+ through a callback function or displays a progress bar in Jupyter.
+
+ Args:
+ callback: Optional function to call with progress updates.
+ Signature: callback(status: str, progress: float)
+ poll_interval: How often to check progress (seconds).
+
+ Returns:
+ The collected result batches.
+
+ Example:
+ >>> def my_callback(status, progress):
+ ... print(f"{status}: {progress:.1%}")
+ >>> batches = df.collect_with_progress(callback=my_callback)
+ """
+ import threading
+ import time
+
+ result = [None]
+ error = [None]
+ done = threading.Event()
+
+ def execute():
+ try:
+ result[0] = self.collect()
+ except Exception as e:
+ error[0] = e
+ finally:
+ done.set()
+
+ thread = threading.Thread(target=execute)
+ thread.start()
+
+ # Check if we're in a Jupyter environment
+ try:
+ from IPython.display import display, clear_output
+ from IPython import get_ipython
+
+ in_jupyter = get_ipython() is not None
+ except (ImportError, AttributeError):
+ in_jupyter = False
+
+ start_time = time.time()
+
+ if in_jupyter and callback is None:
+ # Display a simple progress indicator
+ try:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ clear_output(wait=True)
+ print(f"⏳ Query executing... ({elapsed:.1f}s elapsed)")
+
+ clear_output(wait=True)
+ elapsed = time.time() - start_time
+ print(f"✓ Query completed in {elapsed:.1f}s")
+ except Exception:
+ pass # Ignore display errors
+ elif callback is not None:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ callback(f"Executing ({elapsed:.1f}s)", -1.0) # -1 means indeterminate
+
+ elapsed = time.time() - start_time
+ callback(f"Completed in {elapsed:.1f}s", 1.0)
+ else:
+ done.wait()
+
+ thread.join()
+
+ if error[0] is not None:
+ raise error[0]
+
+ return result[0]
+
+
+class ExecutionPlanVisualization:
+ """
+ A wrapper for execution plan visualizations that can render as SVG in Jupyter.
+
+ This class takes the text representation of an execution plan and converts
+ it to a Graphviz DOT format, which is then rendered as SVG.
+ """
+
+ def __init__(self, plan_str: str, analyze: bool = False):
+ self.plan_str = plan_str
+ self.analyze = analyze
+ self._svg_cache: Optional[str] = None
+
+ def _parse_plan_to_dot(self) -> str:
+ """Convert the plan string to DOT format for Graphviz."""
+ lines = self.plan_str.strip().split("\n")
+
+ dot_lines = [
+ "digraph ExecutionPlan {",
+ ' rankdir=TB;',
+ ' node [shape=box, style="rounded,filled", fontname="Helvetica"];',
+ ' edge [fontname="Helvetica"];',
+ "",
+ ]
+
+ nodes = []
+ edges = []
+ node_id = 0
+ stack = [] # (indent_level, node_id)
+
+ for line in lines:
+ if not line.strip():
+ continue
+
+ # Calculate indent level
+ indent = len(line) - len(line.lstrip())
+ content = line.strip()
+
+ # Skip non-plan lines
+ if content.startswith("physical_plan") or content.startswith("logical_plan"):
+ continue
+
+ # Create a node for this plan element
+ current_id = node_id
+ node_id += 1
+
+ # Determine node color based on operation type
+ color = "#E3F2FD" # Default light blue
+ if "Scan" in content or "TableScan" in content:
+ color = "#E8F5E9" # Light green for scans
+ elif "Filter" in content:
+ color = "#FFF3E0" # Light orange for filters
+ elif "Aggregate" in content or "HashAggregate" in content:
+ color = "#F3E5F5" # Light purple for aggregations
+ elif "Join" in content:
+ color = "#FFEBEE" # Light red for joins
+ elif "Sort" in content:
+ color = "#E0F7FA" # Light cyan for sorts
+ elif "Projection" in content:
+ color = "#FFF8E1" # Light amber for projections
+
+ # Escape special characters for DOT format
+ label = content.replace('"', '\\"').replace("\n", "\\n")
+ if len(label) > 60:
+ # Wrap long labels
+ label = label[:57] + "..."
+
+ nodes.append(f' node{current_id} [label="{label}", fillcolor="{color}"];')
+
+ # Connect to parent based on indentation
+ while stack and stack[-1][0] >= indent:
+ stack.pop()
+
+ if stack:
+ parent_id = stack[-1][1]
+ edges.append(f" node{parent_id} -> node{current_id};")
+
+ stack.append((indent, current_id))
+
+ dot_lines.extend(nodes)
+ dot_lines.append("")
+ dot_lines.extend(edges)
+ dot_lines.append("}")
+
+ return "\n".join(dot_lines)
+
+ def to_dot(self) -> str:
+ """Get the DOT representation of the execution plan."""
+ return self._parse_plan_to_dot()
+
+ def to_svg(self) -> str:
+ """
+ Convert the execution plan to SVG format.
+
+ Requires graphviz to be installed. If graphviz is not available,
+ returns a simple HTML representation instead.
+ """
+ if self._svg_cache is not None:
+ return self._svg_cache
+
+ dot_source = self._parse_plan_to_dot()
+
+ try:
+ import subprocess
+
+ # Try to use graphviz's dot command
+ process = subprocess.run(
+ ["dot", "-Tsvg"],
+ input=dot_source.encode(),
+ capture_output=True,
+ timeout=30,
+ )
+
+ if process.returncode == 0:
+ self._svg_cache = process.stdout.decode()
+ return self._svg_cache
+ except (subprocess.SubprocessError, FileNotFoundError, subprocess.TimeoutExpired):
+ pass
+
+ # Fallback: return a pre-formatted HTML representation
+ escaped_plan = (
+ self.plan_str.replace("&", "&")
+ .replace("<", "<")
+ .replace(">", ">")
+ )
+ self._svg_cache = f"""
+ <div style="font-family: monospace; background: #f5f5f5; padding: 10px;
+ border-radius: 5px; overflow-x: auto;">
+ <div style="color: #666; margin-bottom: 5px;">
+ Execution Plan {'(with statistics)' if self.analyze else ''}
+ <br><small>Install graphviz for visual diagram: brew install graphviz</small>
+ </div>
+ <pre style="margin: 0;">{escaped_plan}</pre>
+ </div>
+ """
+ return self._svg_cache
+
+ def save(self, path: str) -> None:
+ """Save the visualization to a file (SVG or DOT format)."""
+ if path.endswith(".dot"):
+ content = self.to_dot()
+ else:
+ content = self.to_svg()
+
+ with open(path, "w") as f:
+ f.write(content)
+
+ def _repr_html_(self) -> str:
+ """HTML representation for Jupyter notebooks."""
+ return self.to_svg()
+
+ def _repr_svg_(self) -> str:
+ """SVG representation for Jupyter notebooks."""
+ svg = self.to_svg()
+ # Only return if it's actual SVG content
+ if svg.strip().startswith("<svg") or svg.strip().startswith("<?xml"):
+ return svg
+ return ""
+
+ def __repr__(self) -> str:
+ """String representation."""
+ return f"ExecutionPlanVisualization(analyze={self.analyze})\n{self.plan_str}"
+
class BallistaSessionContext(SessionContext, metaclass=RedefiningSessionContextMeta):
+ """
+ A session context for connecting to and querying a Ballista cluster.
+
+ This class extends DataFusion's SessionContext to work with distributed
+ Ballista clusters, automatically routing query execution to the cluster
+ while maintaining API compatibility with local DataFusion usage.
+
+ Example:
+ >>> from ballista import BallistaSessionContext
+ >>> ctx = BallistaSessionContext("df://localhost:50050")
+ >>> df = ctx.sql("SELECT * FROM my_table LIMIT 10")
+ >>> df.show()
+
+ For Jupyter notebook users:
+ >>> %load_ext ballista.jupyter
+ >>> %ballista connect df://localhost:50050
+ >>> %sql SELECT * FROM my_table
+ """
+
def __init__(self, address: str, config=None, runtime=None):
super().__init__(config, runtime)
self.address = address
self.session_id = self.session_id()
+
+ def get_tables(self) -> Optional[dict[str, str]]:
Review Comment:
The dict values are List of strings, not just a String.
```suggestion
def get_tables(self) -> Dict[str, List[str]]:
```
##########
python/python/ballista/extension.py:
##########
@@ -229,9 +231,319 @@ def write_parquet(
df = self._to_internal_df()
df.write_parquet(str(path), compression.value, compression_level)
+ def explain_visual(self, analyze: bool = False) -> "ExecutionPlanVisualization":
+ """
+ Generate a visual representation of the execution plan.
+
+ This method creates an SVG visualization of the query execution plan,
+ which can be displayed directly in Jupyter notebooks.
+
+ Args:
+ analyze: If True, includes runtime statistics from actual execution.
+
+ Returns:
+ ExecutionPlanVisualization: An object that renders as SVG in Jupyter.
+
+ Example:
+ >>> df = ctx.sql("SELECT * FROM orders WHERE amount > 100")
+ >>> df.explain_visual() # Displays SVG in notebook
+ >>> viz = df.explain_visual(analyze=True)
+ >>> viz.save("plan.svg") # Save to file
+ """
+ # Get the execution plan as a string representation
+ # Note: explain() prints but doesn't return a string, so we use logical_plan()
+ try:
+ plan = self.logical_plan()
+ plan_str = plan.display_indent()
+ except Exception:
+ # Fallback if logical_plan() fails
+ plan_str = "Unable to retrieve execution plan"
+ return ExecutionPlanVisualization(plan_str, analyze=analyze)
+
+ def collect_with_progress(
+ self,
+ callback: Optional[callable] = None,
+ poll_interval: float = 0.5,
+ ):
+ """
+ Collect results with progress indication.
+
+ For long-running queries, this method provides progress updates
+ through a callback function or displays a progress bar in Jupyter.
+
+ Args:
+ callback: Optional function to call with progress updates.
+ Signature: callback(status: str, progress: float)
+ poll_interval: How often to check progress (seconds).
+
+ Returns:
+ The collected result batches.
+
+ Example:
+ >>> def my_callback(status, progress):
+ ... print(f"{status}: {progress:.1%}")
+ >>> batches = df.collect_with_progress(callback=my_callback)
+ """
+ import threading
+ import time
+
+ result = [None]
+ error = [None]
+ done = threading.Event()
+
+ def execute():
+ try:
+ result[0] = self.collect()
+ except Exception as e:
+ error[0] = e
+ finally:
+ done.set()
+
+ thread = threading.Thread(target=execute)
+ thread.start()
+
+ # Check if we're in a Jupyter environment
+ try:
+ from IPython.display import display, clear_output
+ from IPython import get_ipython
+
+ in_jupyter = get_ipython() is not None
+ except (ImportError, AttributeError):
+ in_jupyter = False
+
+ start_time = time.time()
+
+ if in_jupyter and callback is None:
+ # Display a simple progress indicator
+ try:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ clear_output(wait=True)
+ print(f"⏳ Query executing... ({elapsed:.1f}s elapsed)")
+
+ clear_output(wait=True)
+ elapsed = time.time() - start_time
+ print(f"✓ Query completed in {elapsed:.1f}s")
+ except Exception:
+ pass # Ignore display errors
+ elif callback is not None:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ callback(f"Executing ({elapsed:.1f}s)", -1.0) # -1 means indeterminate
+
+ elapsed = time.time() - start_time
+ callback(f"Completed in {elapsed:.1f}s", 1.0)
+ else:
+ done.wait()
+
+ thread.join()
+
+ if error[0] is not None:
+ raise error[0]
+
+ return result[0]
+
+
+class ExecutionPlanVisualization:
+ """
+ A wrapper for execution plan visualizations that can render as SVG in Jupyter.
+
+ This class takes the text representation of an execution plan and converts
+ it to a Graphviz DOT format, which is then rendered as SVG.
+ """
+
+ def __init__(self, plan_str: str, analyze: bool = False):
+ self.plan_str = plan_str
+ self.analyze = analyze
+ self._svg_cache: Optional[str] = None
+
+ def _parse_plan_to_dot(self) -> str:
+ """Convert the plan string to DOT format for Graphviz."""
+ lines = self.plan_str.strip().split("\n")
+
+ dot_lines = [
+ "digraph ExecutionPlan {",
+ ' rankdir=TB;',
+ ' node [shape=box, style="rounded,filled", fontname="Helvetica"];',
+ ' edge [fontname="Helvetica"];',
+ "",
+ ]
+
+ nodes = []
+ edges = []
+ node_id = 0
+ stack = [] # (indent_level, node_id)
+
+ for line in lines:
+ if not line.strip():
+ continue
+
+ # Calculate indent level
+ indent = len(line) - len(line.lstrip())
+ content = line.strip()
+
+ # Skip non-plan lines
+ if content.startswith("physical_plan") or content.startswith("logical_plan"):
+ continue
+
+ # Create a node for this plan element
+ current_id = node_id
+ node_id += 1
+
+ # Determine node color based on operation type
+ color = "#E3F2FD" # Default light blue
+ if "Scan" in content or "TableScan" in content:
+ color = "#E8F5E9" # Light green for scans
+ elif "Filter" in content:
+ color = "#FFF3E0" # Light orange for filters
+ elif "Aggregate" in content or "HashAggregate" in content:
+ color = "#F3E5F5" # Light purple for aggregations
+ elif "Join" in content:
+ color = "#FFEBEE" # Light red for joins
+ elif "Sort" in content:
+ color = "#E0F7FA" # Light cyan for sorts
+ elif "Projection" in content:
+ color = "#FFF8E1" # Light amber for projections
+
+ # Escape special characters for DOT format
+ label = content.replace('"', '\\"').replace("\n", "\\n")
+ if len(label) > 60:
+ # Wrap long labels
+ label = label[:57] + "..."
+
+ nodes.append(f' node{current_id} [label="{label}", fillcolor="{color}"];')
+
+ # Connect to parent based on indentation
+ while stack and stack[-1][0] >= indent:
+ stack.pop()
+
+ if stack:
+ parent_id = stack[-1][1]
+ edges.append(f" node{parent_id} -> node{current_id};")
+
+ stack.append((indent, current_id))
+
+ dot_lines.extend(nodes)
+ dot_lines.append("")
+ dot_lines.extend(edges)
+ dot_lines.append("}")
+
+ return "\n".join(dot_lines)
+
+ def to_dot(self) -> str:
+ """Get the DOT representation of the execution plan."""
+ return self._parse_plan_to_dot()
+
+ def to_svg(self) -> str:
+ """
+ Convert the execution plan to SVG format.
+
+ Requires graphviz to be installed. If graphviz is not available,
+ returns a simple HTML representation instead.
+ """
+ if self._svg_cache is not None:
+ return self._svg_cache
+
+ dot_source = self._parse_plan_to_dot()
+
+ try:
+ import subprocess
+
+ # Try to use graphviz's dot command
+ process = subprocess.run(
+ ["dot", "-Tsvg"],
+ input=dot_source.encode(),
+ capture_output=True,
+ timeout=30,
+ )
+
+ if process.returncode == 0:
+ self._svg_cache = process.stdout.decode()
+ return self._svg_cache
+ except (subprocess.SubprocessError, FileNotFoundError, subprocess.TimeoutExpired):
+ pass
Review Comment:
Let's log warnings instead of just passing silently
##########
python/python/ballista/extension.py:
##########
@@ -229,9 +231,319 @@ def write_parquet(
df = self._to_internal_df()
df.write_parquet(str(path), compression.value, compression_level)
+ def explain_visual(self, analyze: bool = False) -> "ExecutionPlanVisualization":
+ """
+ Generate a visual representation of the execution plan.
+
+ This method creates an SVG visualization of the query execution plan,
+ which can be displayed directly in Jupyter notebooks.
+
+ Args:
+ analyze: If True, includes runtime statistics from actual execution.
+
+ Returns:
+ ExecutionPlanVisualization: An object that renders as SVG in Jupyter.
+
+ Example:
+ >>> df = ctx.sql("SELECT * FROM orders WHERE amount > 100")
+ >>> df.explain_visual() # Displays SVG in notebook
+ >>> viz = df.explain_visual(analyze=True)
+ >>> viz.save("plan.svg") # Save to file
+ """
+ # Get the execution plan as a string representation
+ # Note: explain() prints but doesn't return a string, so we use logical_plan()
+ try:
+ plan = self.logical_plan()
+ plan_str = plan.display_indent()
+ except Exception:
+ # Fallback if logical_plan() fails
+ plan_str = "Unable to retrieve execution plan"
+ return ExecutionPlanVisualization(plan_str, analyze=analyze)
+
+ def collect_with_progress(
+ self,
+ callback: Optional[callable] = None,
+ poll_interval: float = 0.5,
+ ):
+ """
+ Collect results with progress indication.
+
+ For long-running queries, this method provides progress updates
+ through a callback function or displays a progress bar in Jupyter.
+
+ Args:
+ callback: Optional function to call with progress updates.
+ Signature: callback(status: str, progress: float)
+ poll_interval: How often to check progress (seconds).
+
+ Returns:
+ The collected result batches.
+
+ Example:
+ >>> def my_callback(status, progress):
+ ... print(f"{status}: {progress:.1%}")
+ >>> batches = df.collect_with_progress(callback=my_callback)
+ """
+ import threading
+ import time
+
+ result = [None]
+ error = [None]
+ done = threading.Event()
+
+ def execute():
+ try:
+ result[0] = self.collect()
+ except Exception as e:
+ error[0] = e
+ finally:
+ done.set()
+
+ thread = threading.Thread(target=execute)
+ thread.start()
+
+ # Check if we're in a Jupyter environment
+ try:
+ from IPython.display import display, clear_output
+ from IPython import get_ipython
+
+ in_jupyter = get_ipython() is not None
+ except (ImportError, AttributeError):
+ in_jupyter = False
+
+ start_time = time.time()
+
+ if in_jupyter and callback is None:
+ # Display a simple progress indicator
+ try:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ clear_output(wait=True)
+ print(f"⏳ Query executing... ({elapsed:.1f}s elapsed)")
+
+ clear_output(wait=True)
+ elapsed = time.time() - start_time
+ print(f"✓ Query completed in {elapsed:.1f}s")
+ except Exception:
+ pass # Ignore display errors
Review Comment:
Let's log the exception as a warning
##########
python/python/ballista/jupyter.py:
##########
@@ -0,0 +1,493 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Ballista Jupyter Magic Commands.
+
+This module provides IPython magic commands for interacting with Ballista
+clusters directly from Jupyter notebooks.
+
+Usage:
+ %load_ext ballista.jupyter
+
+ # Connect to a Ballista cluster
+ %ballista connect df://localhost:50050
+
+ # Check connection status
+ %ballista status
+
+ # List registered tables
+ %ballista tables
+
+ # Show schema for a table
+ %ballista schema my_table
+
+ # Execute a simple SQL query (line magic)
+ %sql SELECT COUNT(*) FROM my_table
+
+ # Execute a complex SQL query (cell magic)
+ %%sql
+ SELECT
+ customer_id,
+ SUM(amount) as total
+ FROM orders
+ GROUP BY customer_id
+ ORDER BY total DESC
+ LIMIT 10
+"""
+
+from typing import Optional, List, Dict, Any
+import time
+import threading
+
+try:
+ from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic
+ from IPython.display import display, HTML
+
+ IPYTHON_AVAILABLE = True
+except ImportError:
+ IPYTHON_AVAILABLE = False
+
+ # Provide stub classes for when IPython is not available
+ class Magics:
+ def __init__(self, shell=None):
+ self.shell = shell
+
+ def magics_class(cls):
+ return cls
+
+ def line_magic(name_or_func):
+ """Stub line_magic decorator for when IPython is not available."""
+ # Handle both @line_magic and @line_magic("name") usage
+ if callable(name_or_func):
+ # Used as @line_magic without arguments
+ return name_or_func
+ else:
+ # Used as @line_magic("name") with arguments
+ def decorator(func):
+ return func
+ return decorator
+
+ def cell_magic(name_or_func):
+ """Stub cell_magic decorator for when IPython is not available."""
+ # Handle both @cell_magic and @cell_magic("name") usage
+ if callable(name_or_func):
+ return name_or_func
+ else:
+ def decorator(func):
+ return func
+ return decorator
+
+
+from .extension import BallistaSessionContext, DistributedDataFrame
+
+
+class BallistaConnectionError(Exception):
+ """Raised when not connected to a Ballista cluster."""
+
+ pass
+
+
+@magics_class
+class BallistaMagics(Magics):
+ """
+ IPython magic commands for Ballista.
+
+ Provides convenient commands for connecting to Ballista clusters,
+ executing SQL queries, and exploring table schemas.
+ """
+
+ def __init__(self, shell=None):
+ super().__init__(shell)
+ self._ctx: Optional[BallistaSessionContext] = None
+ self._address: Optional[str] = None
+ self._last_result: Optional[DistributedDataFrame] = None
+ self._query_history: List[Dict[str, Any]] = []
+
+ @property
+ def ctx(self) -> BallistaSessionContext:
+ """Get the current context, raising an error if not connected."""
+ if self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ return self._ctx
+
+ @property
+ def is_connected(self) -> bool:
+ """Check if connected to a Ballista cluster."""
+ return self._ctx is not None
+
+ @line_magic
+ def ballista(self, line: str) -> Optional[str]:
+ """
+ Ballista management commands.
+
+ Usage:
+ %ballista connect df://localhost:50050 - Connect to cluster
+ %ballista status - Show connection status
+ %ballista tables - List registered tables
+ %ballista schema <table> - Show table schema
+ %ballista disconnect - Disconnect from cluster
+ %ballista history - Show query history
+
+ Examples:
+ %ballista connect df://localhost:50050
+ %ballista tables
+ %ballista schema orders
+ """
+ parts = line.strip().split(maxsplit=1)
+ if not parts:
+ return self._show_help()
+
+ cmd = parts[0].lower()
+ args = parts[1] if len(parts) > 1 else ""
+
+ if cmd == "connect":
+ return self._connect(args)
+ elif cmd == "status":
+ return self._status()
+ elif cmd == "tables":
+ return self._tables()
+ elif cmd == "schema":
+ return self._schema(args)
+ elif cmd == "disconnect":
+ return self._disconnect()
+ elif cmd == "history":
+ return self._show_history()
+ elif cmd == "help":
+ return self._show_help()
+ else:
+ return f"Unknown command: {cmd}. Use '%ballista help' for available commands."
+
+ @line_magic
+ def register(self, line: str) -> Optional[str]:
+ """Register a new table"""
+ if not line:
+ return "You should provide file extension and table name to register"
+ elif self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ else:
+ args = line.strip().split()
+ file_type = args[0]
+ if len(args) < 2:
+ return f"You should provide table name for this .{file_type} file"
+ table_name = args[1]
+ if len(args) < 3:
+ return f"You should provide path to your file"
+ file_name = args[2]
+
+ if file_type == "parquet":
+ self._ctx.register_parquet(table_name, file_name)
+ elif file_type == "csv":
+ self._ctx.register_csv(table_name, file_name)
+ else:
+ raise NotImplemented("Currently not supporting the inserted file format")
+
+ @line_cell_magic
+ def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]:
+ """
+ Execute a SQL query (both line and cell magic).
+
+ Two cases possible: with cell or without cell
+
+ Examples:
+ 1. Without a cell (line_magic)
+ %sql SELECT * FROM test_table
+ 2. With a cell (cell_magic)
+ %%sql --no-display
+ SELECT
+ id,
+ bool_col,
+ tinyint_col
+ FROM test_data_v1
+ WHERE id > 2
+ ORDER BY id
+ LIMIT 5
+ """
+ if not cell:
+ return self._execute_sql(line.strip()) if line.strip() else None
+ else:
+ args = line.strip().split()
+ display_results = True
+ limit = 50
+ var_name = None
+
+ i = 0
+ while i < len(args):
+ if args[i] == "--no-display":
+ display_results = False
+ elif args[i] == "--limit" and i + 1 < len(args):
+ try:
+ limit = int(args[i + 1])
+ i += 1
+ except ValueError:
+ pass
+ elif not args[i].startswith("--"):
+ var_name = args[i]
+ i += 1
+
+ query = cell.strip()
+ if not query:
+ return None
+
+ result = self._execute_sql(query, display_results=display_results, limit=limit)
+
+ # Store in user namespace if variable name provided
+ if var_name and self.shell is not None:
+ self.shell.user_ns[var_name] = result
+
+ return result
+
+ def _connect(self, address: str) -> Optional[str]:
+ """Connect to a Ballista cluster."""
+ if not address:
+ return "Usage: %ballista connect df://host:port"
+
+ # Normalize address
+ if not address.startswith("df://"):
+ address = f"df://{address}"
+
+ try:
+ self._ctx = BallistaSessionContext(address)
+ self._address = address
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Connected to Ballista cluster at {address}"))
+ else:
+ return f"✓ Connected to Ballista cluster at {address}"
+ except Exception as e:
+ self._ctx = None
+ self._address = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✗ Failed to connect to {address}: {e}"))
+ else:
+ return f"✗ Failed to connect to {address}: {e}"
+
+ def _disconnect(self) -> Optional[str]:
+ """Disconnect from the Ballista cluster."""
+ if not self.is_connected:
+ return "Not connected to any cluster."
+
+ address = self._address
+ self._ctx = None
+ self._address = None
+ self._last_result = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Disconnected from {address}"))
+ else:
+ return f"✓ Disconnected from {address}"
+
+ def _status(self) -> Optional[str]:
+ """Show connection status."""
+ if not self.is_connected:
+ return "Status: Not connected\n\nUse '%ballista connect df://host:port' to connect."
+
+ status_lines = [
+ f"Status: Connected",
+ f"Address: {self._address}",
+ f"Session ID: {self._ctx.session_id}",
+ f"Queries executed: {len(self._query_history)}",
+ ]
+
+ if self._last_result is not None:
+ status_lines.append("Last result: Available (access via '_' or '_last_result')")
+
+ def _format_html_status_output(line: str) -> str:
+ name, value = line.split(":", 1)
+ return f"<p><strong>{name}:</strong> {value.strip()}</p>"
+
+ html = "".join(_format_html_status_output(line) for line in status_lines)
+ if IPYTHON_AVAILABLE:
+ display(HTML(html))
+ else:
+ print("\n".join(status_lines))
+
+ def _tables(self) -> Optional[str]:
+ """List all registered tables."""
+ try:
+ # Get table names from the catalog and their respective schemas
+ tables = self.ctx.get_tables()
+ 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())
Review Comment:
```suggestion
table_count = sum(len(v) for v in tables.values())
```
##########
python/python/ballista/jupyter.py:
##########
@@ -0,0 +1,493 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Ballista Jupyter Magic Commands.
+
+This module provides IPython magic commands for interacting with Ballista
+clusters directly from Jupyter notebooks.
+
+Usage:
+ %load_ext ballista.jupyter
+
+ # Connect to a Ballista cluster
+ %ballista connect df://localhost:50050
+
+ # Check connection status
+ %ballista status
+
+ # List registered tables
+ %ballista tables
+
+ # Show schema for a table
+ %ballista schema my_table
+
+ # Execute a simple SQL query (line magic)
+ %sql SELECT COUNT(*) FROM my_table
+
+ # Execute a complex SQL query (cell magic)
+ %%sql
+ SELECT
+ customer_id,
+ SUM(amount) as total
+ FROM orders
+ GROUP BY customer_id
+ ORDER BY total DESC
+ LIMIT 10
+"""
+
+from typing import Optional, List, Dict, Any
+import time
+import threading
+
+try:
+ from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic
+ from IPython.display import display, HTML
+
+ IPYTHON_AVAILABLE = True
+except ImportError:
+ IPYTHON_AVAILABLE = False
+
+ # Provide stub classes for when IPython is not available
+ class Magics:
+ def __init__(self, shell=None):
+ self.shell = shell
+
+ def magics_class(cls):
+ return cls
+
+ def line_magic(name_or_func):
+ """Stub line_magic decorator for when IPython is not available."""
+ # Handle both @line_magic and @line_magic("name") usage
+ if callable(name_or_func):
+ # Used as @line_magic without arguments
+ return name_or_func
+ else:
+ # Used as @line_magic("name") with arguments
+ def decorator(func):
+ return func
+ return decorator
+
+ def cell_magic(name_or_func):
+ """Stub cell_magic decorator for when IPython is not available."""
+ # Handle both @cell_magic and @cell_magic("name") usage
+ if callable(name_or_func):
+ return name_or_func
+ else:
+ def decorator(func):
+ return func
+ return decorator
Review Comment:
No fallback for `line_cell_magic`. `@line_cell_magic` will fail if IPython is not available
##########
python/examples/getting_started.ipynb:
##########
@@ -0,0 +1,353 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "<!---\n",
+ " Licensed to the Apache Software Foundation (ASF) under one\n",
+ " or more contributor license agreements. See the NOTICE file\n",
+ " distributed with this work for additional information\n",
+ " regarding copyright ownership. The ASF licenses this file\n",
+ " to you under the Apache License, Version 2.0 (the\n",
+ " \"License\"); you may not use this file except in compliance\n",
+ " with the License. You may obtain a copy of the License at\n",
+ "\n",
+ " http://www.apache.org/licenses/LICENSE-2.0\n",
+ "\n",
+ " Unless required by applicable law or agreed to in writing,\n",
+ " software distributed under the License is distributed on an\n",
+ " \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY\n",
+ " KIND, either express or implied. See the License for the\n",
+ " specific language governing permissions and limitations\n",
+ " under the License.\n",
+ "-->\n",
+ "\n",
+ "# Getting Started with PyBallista\n",
+ "\n",
+ "This notebook demonstrates how to get started with Ballista using Python.\n",
+ "\n",
+ "## Prerequisites\n",
+ "\n",
+ "1. Install PyBallista: `pip install ballista`\n",
+ "2. Have a Ballista cluster running (or use the built-in test cluster)\n",
+ "\n",
+ "## Overview\n",
+ "\n",
+ "Ballista is a distributed query engine built on Apache DataFusion. PyBallista provides:\n",
+ "\n",
+ "- **BallistaSessionContext**: Drop-in replacement for DataFusion's SessionContext\n",
+ "- **SQL Magic Commands**: Interactive SQL in Jupyter notebooks via `%sql` and `%%sql`\n",
+ "- **DataFrame API**: Full DataFrame API for data transformations\n",
+ "- **Rich HTML Display**: DataFrames render as styled HTML tables"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Method 1: Python API\n",
+ "\n",
+ "The most straightforward way to use Ballista is via the Python API."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import ballista\n",
+ "from ballista import BallistaSessionContext, setup_test_cluster\n",
+ "%load_ext autoreload\n",
+ "%autoreload 2\n",
+ "\n",
+ "# Check versions\n",
+ "print(f\"Ballista version: {ballista.__version__}\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# For this demo, we'll use the built-in test cluster\n",
+ "# In production, you would connect to your Ballista scheduler:\n",
+ "# ctx = BallistaSessionContext(\"df://your-scheduler:50050\")\n",
+ "\n",
+ "# host, port = setup_test_cluster()\n",
+ "host, port = \"localhost\", \"39431\"\n",
Review Comment:
Is this a debug leftover ? IMO line 79 should be used
##########
python/python/ballista/jupyter.py:
##########
@@ -0,0 +1,493 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Ballista Jupyter Magic Commands.
+
+This module provides IPython magic commands for interacting with Ballista
+clusters directly from Jupyter notebooks.
+
+Usage:
+ %load_ext ballista.jupyter
+
+ # Connect to a Ballista cluster
+ %ballista connect df://localhost:50050
+
+ # Check connection status
+ %ballista status
+
+ # List registered tables
+ %ballista tables
+
+ # Show schema for a table
+ %ballista schema my_table
+
+ # Execute a simple SQL query (line magic)
+ %sql SELECT COUNT(*) FROM my_table
+
+ # Execute a complex SQL query (cell magic)
+ %%sql
+ SELECT
+ customer_id,
+ SUM(amount) as total
+ FROM orders
+ GROUP BY customer_id
+ ORDER BY total DESC
+ LIMIT 10
+"""
+
+from typing import Optional, List, Dict, Any
+import time
+import threading
+
+try:
+ from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic
+ from IPython.display import display, HTML
+
+ IPYTHON_AVAILABLE = True
+except ImportError:
+ IPYTHON_AVAILABLE = False
+
+ # Provide stub classes for when IPython is not available
+ class Magics:
+ def __init__(self, shell=None):
+ self.shell = shell
+
+ def magics_class(cls):
+ return cls
+
+ def line_magic(name_or_func):
+ """Stub line_magic decorator for when IPython is not available."""
+ # Handle both @line_magic and @line_magic("name") usage
+ if callable(name_or_func):
+ # Used as @line_magic without arguments
+ return name_or_func
+ else:
+ # Used as @line_magic("name") with arguments
+ def decorator(func):
+ return func
+ return decorator
+
+ def cell_magic(name_or_func):
+ """Stub cell_magic decorator for when IPython is not available."""
+ # Handle both @cell_magic and @cell_magic("name") usage
+ if callable(name_or_func):
+ return name_or_func
+ else:
+ def decorator(func):
+ return func
+ return decorator
+
+
+from .extension import BallistaSessionContext, DistributedDataFrame
+
+
+class BallistaConnectionError(Exception):
+ """Raised when not connected to a Ballista cluster."""
+
+ pass
+
+
+@magics_class
+class BallistaMagics(Magics):
+ """
+ IPython magic commands for Ballista.
+
+ Provides convenient commands for connecting to Ballista clusters,
+ executing SQL queries, and exploring table schemas.
+ """
+
+ def __init__(self, shell=None):
+ super().__init__(shell)
+ self._ctx: Optional[BallistaSessionContext] = None
+ self._address: Optional[str] = None
+ self._last_result: Optional[DistributedDataFrame] = None
+ self._query_history: List[Dict[str, Any]] = []
+
+ @property
+ def ctx(self) -> BallistaSessionContext:
+ """Get the current context, raising an error if not connected."""
+ if self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ return self._ctx
+
+ @property
+ def is_connected(self) -> bool:
+ """Check if connected to a Ballista cluster."""
+ return self._ctx is not None
+
+ @line_magic
+ def ballista(self, line: str) -> Optional[str]:
+ """
+ Ballista management commands.
+
+ Usage:
+ %ballista connect df://localhost:50050 - Connect to cluster
+ %ballista status - Show connection status
+ %ballista tables - List registered tables
+ %ballista schema <table> - Show table schema
+ %ballista disconnect - Disconnect from cluster
+ %ballista history - Show query history
+
+ Examples:
+ %ballista connect df://localhost:50050
+ %ballista tables
+ %ballista schema orders
+ """
+ parts = line.strip().split(maxsplit=1)
+ if not parts:
+ return self._show_help()
+
+ cmd = parts[0].lower()
+ args = parts[1] if len(parts) > 1 else ""
+
+ if cmd == "connect":
+ return self._connect(args)
+ elif cmd == "status":
+ return self._status()
+ elif cmd == "tables":
+ return self._tables()
+ elif cmd == "schema":
+ return self._schema(args)
+ elif cmd == "disconnect":
+ return self._disconnect()
+ elif cmd == "history":
+ return self._show_history()
+ elif cmd == "help":
+ return self._show_help()
+ else:
+ return f"Unknown command: {cmd}. Use '%ballista help' for available commands."
+
+ @line_magic
+ def register(self, line: str) -> Optional[str]:
+ """Register a new table"""
+ if not line:
+ return "You should provide file extension and table name to register"
+ elif self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ else:
+ args = line.strip().split()
+ file_type = args[0]
+ if len(args) < 2:
+ return f"You should provide table name for this .{file_type} file"
+ table_name = args[1]
+ if len(args) < 3:
+ return f"You should provide path to your file"
+ file_name = args[2]
+
+ if file_type == "parquet":
+ self._ctx.register_parquet(table_name, file_name)
+ elif file_type == "csv":
+ self._ctx.register_csv(table_name, file_name)
+ else:
+ raise NotImplemented("Currently not supporting the inserted file format")
+
+ @line_cell_magic
+ def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]:
+ """
+ Execute a SQL query (both line and cell magic).
+
+ Two cases possible: with cell or without cell
+
+ Examples:
+ 1. Without a cell (line_magic)
+ %sql SELECT * FROM test_table
+ 2. With a cell (cell_magic)
+ %%sql --no-display
+ SELECT
+ id,
+ bool_col,
+ tinyint_col
+ FROM test_data_v1
+ WHERE id > 2
+ ORDER BY id
+ LIMIT 5
+ """
+ if not cell:
+ return self._execute_sql(line.strip()) if line.strip() else None
+ else:
+ args = line.strip().split()
+ display_results = True
+ limit = 50
+ var_name = None
+
+ i = 0
+ while i < len(args):
+ if args[i] == "--no-display":
+ display_results = False
+ elif args[i] == "--limit" and i + 1 < len(args):
+ try:
+ limit = int(args[i + 1])
+ i += 1
+ except ValueError:
+ pass
+ elif not args[i].startswith("--"):
+ var_name = args[i]
+ i += 1
+
+ query = cell.strip()
+ if not query:
+ return None
+
+ result = self._execute_sql(query, display_results=display_results, limit=limit)
+
+ # Store in user namespace if variable name provided
+ if var_name and self.shell is not None:
+ self.shell.user_ns[var_name] = result
+
+ return result
+
+ def _connect(self, address: str) -> Optional[str]:
+ """Connect to a Ballista cluster."""
+ if not address:
+ return "Usage: %ballista connect df://host:port"
+
+ # Normalize address
+ if not address.startswith("df://"):
+ address = f"df://{address}"
+
+ try:
+ self._ctx = BallistaSessionContext(address)
+ self._address = address
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Connected to Ballista cluster at {address}"))
+ else:
+ return f"✓ Connected to Ballista cluster at {address}"
+ except Exception as e:
+ self._ctx = None
+ self._address = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✗ Failed to connect to {address}: {e}"))
+ else:
+ return f"✗ Failed to connect to {address}: {e}"
+
+ def _disconnect(self) -> Optional[str]:
+ """Disconnect from the Ballista cluster."""
+ if not self.is_connected:
+ return "Not connected to any cluster."
+
+ address = self._address
+ self._ctx = None
+ self._address = None
+ self._last_result = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Disconnected from {address}"))
+ else:
+ return f"✓ Disconnected from {address}"
+
+ def _status(self) -> Optional[str]:
+ """Show connection status."""
+ if not self.is_connected:
+ return "Status: Not connected\n\nUse '%ballista connect df://host:port' to connect."
+
+ status_lines = [
+ f"Status: Connected",
+ f"Address: {self._address}",
+ f"Session ID: {self._ctx.session_id}",
+ f"Queries executed: {len(self._query_history)}",
+ ]
+
+ if self._last_result is not None:
+ status_lines.append("Last result: Available (access via '_' or '_last_result')")
+
+ def _format_html_status_output(line: str) -> str:
+ name, value = line.split(":", 1)
+ return f"<p><strong>{name}:</strong> {value.strip()}</p>"
+
+ html = "".join(_format_html_status_output(line) for line in status_lines)
+ if IPYTHON_AVAILABLE:
+ display(HTML(html))
+ else:
+ print("\n".join(status_lines))
+
+ def _tables(self) -> Optional[str]:
+ """List all registered tables."""
+ try:
+ # Get table names from the catalog and their respective schemas
+ tables = self.ctx.get_tables()
+ 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())
+ # Build a nice table display (HTML-formatted if applicable)
+ lines = [
+ {"content": f"Total: {table_count} table(s) in {schema_count} schema(s)", "is_info": True},
+ {"content": "Registered tables:", "is_info": True},
+ *[{"content": f"Schema: {schema_name}. Tables: {", ".join(list(table_names))}", "is_info": False}
+ for schema_name, table_names in tables.items()]
+ ]
+
+ def _format_html_tables_output(line: str, is_info: bool = False) -> str:
+ if is_info:
+ return f"<p><strong>{line}<strong></p>"
Review Comment:
```suggestion
return f"<p><strong>{line}</strong></p>"
```
##########
python/python/ballista/extension.py:
##########
@@ -229,9 +231,319 @@ def write_parquet(
df = self._to_internal_df()
df.write_parquet(str(path), compression.value, compression_level)
+ def explain_visual(self, analyze: bool = False) -> "ExecutionPlanVisualization":
+ """
+ Generate a visual representation of the execution plan.
+
+ This method creates an SVG visualization of the query execution plan,
+ which can be displayed directly in Jupyter notebooks.
+
+ Args:
+ analyze: If True, includes runtime statistics from actual execution.
+
+ Returns:
+ ExecutionPlanVisualization: An object that renders as SVG in Jupyter.
+
+ Example:
+ >>> df = ctx.sql("SELECT * FROM orders WHERE amount > 100")
+ >>> df.explain_visual() # Displays SVG in notebook
+ >>> viz = df.explain_visual(analyze=True)
+ >>> viz.save("plan.svg") # Save to file
+ """
+ # Get the execution plan as a string representation
+ # Note: explain() prints but doesn't return a string, so we use logical_plan()
+ try:
+ plan = self.logical_plan()
Review Comment:
Is `analyze` supposed to be used around here ?
Currently it is just used to print `with statistics` by ExecutionPlanVisualization, but it is not used for the logical plan
##########
python/python/ballista/jupyter.py:
##########
@@ -0,0 +1,493 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Ballista Jupyter Magic Commands.
+
+This module provides IPython magic commands for interacting with Ballista
+clusters directly from Jupyter notebooks.
+
+Usage:
+ %load_ext ballista.jupyter
+
+ # Connect to a Ballista cluster
+ %ballista connect df://localhost:50050
+
+ # Check connection status
+ %ballista status
+
+ # List registered tables
+ %ballista tables
+
+ # Show schema for a table
+ %ballista schema my_table
+
+ # Execute a simple SQL query (line magic)
+ %sql SELECT COUNT(*) FROM my_table
+
+ # Execute a complex SQL query (cell magic)
+ %%sql
+ SELECT
+ customer_id,
+ SUM(amount) as total
+ FROM orders
+ GROUP BY customer_id
+ ORDER BY total DESC
+ LIMIT 10
+"""
+
+from typing import Optional, List, Dict, Any
+import time
+import threading
+
+try:
+ from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic
+ from IPython.display import display, HTML
+
+ IPYTHON_AVAILABLE = True
+except ImportError:
+ IPYTHON_AVAILABLE = False
+
+ # Provide stub classes for when IPython is not available
+ class Magics:
+ def __init__(self, shell=None):
+ self.shell = shell
+
+ def magics_class(cls):
+ return cls
+
+ def line_magic(name_or_func):
+ """Stub line_magic decorator for when IPython is not available."""
+ # Handle both @line_magic and @line_magic("name") usage
+ if callable(name_or_func):
+ # Used as @line_magic without arguments
+ return name_or_func
+ else:
+ # Used as @line_magic("name") with arguments
+ def decorator(func):
+ return func
+ return decorator
+
+ def cell_magic(name_or_func):
+ """Stub cell_magic decorator for when IPython is not available."""
+ # Handle both @cell_magic and @cell_magic("name") usage
+ if callable(name_or_func):
+ return name_or_func
+ else:
+ def decorator(func):
+ return func
+ return decorator
+
+
+from .extension import BallistaSessionContext, DistributedDataFrame
+
+
+class BallistaConnectionError(Exception):
+ """Raised when not connected to a Ballista cluster."""
+
+ pass
+
+
+@magics_class
+class BallistaMagics(Magics):
+ """
+ IPython magic commands for Ballista.
+
+ Provides convenient commands for connecting to Ballista clusters,
+ executing SQL queries, and exploring table schemas.
+ """
+
+ def __init__(self, shell=None):
+ super().__init__(shell)
+ self._ctx: Optional[BallistaSessionContext] = None
+ self._address: Optional[str] = None
+ self._last_result: Optional[DistributedDataFrame] = None
+ self._query_history: List[Dict[str, Any]] = []
+
+ @property
+ def ctx(self) -> BallistaSessionContext:
+ """Get the current context, raising an error if not connected."""
+ if self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ return self._ctx
+
+ @property
+ def is_connected(self) -> bool:
+ """Check if connected to a Ballista cluster."""
+ return self._ctx is not None
+
+ @line_magic
+ def ballista(self, line: str) -> Optional[str]:
+ """
+ Ballista management commands.
+
+ Usage:
+ %ballista connect df://localhost:50050 - Connect to cluster
+ %ballista status - Show connection status
+ %ballista tables - List registered tables
+ %ballista schema <table> - Show table schema
+ %ballista disconnect - Disconnect from cluster
+ %ballista history - Show query history
+
+ Examples:
+ %ballista connect df://localhost:50050
+ %ballista tables
+ %ballista schema orders
+ """
+ parts = line.strip().split(maxsplit=1)
+ if not parts:
+ return self._show_help()
+
+ cmd = parts[0].lower()
+ args = parts[1] if len(parts) > 1 else ""
+
+ if cmd == "connect":
+ return self._connect(args)
+ elif cmd == "status":
+ return self._status()
+ elif cmd == "tables":
+ return self._tables()
+ elif cmd == "schema":
+ return self._schema(args)
+ elif cmd == "disconnect":
+ return self._disconnect()
+ elif cmd == "history":
+ return self._show_history()
+ elif cmd == "help":
+ return self._show_help()
+ else:
+ return f"Unknown command: {cmd}. Use '%ballista help' for available commands."
+
+ @line_magic
+ def register(self, line: str) -> Optional[str]:
+ """Register a new table"""
+ if not line:
+ return "You should provide file extension and table name to register"
+ elif self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ else:
+ args = line.strip().split()
+ file_type = args[0]
+ if len(args) < 2:
+ return f"You should provide table name for this .{file_type} file"
+ table_name = args[1]
+ if len(args) < 3:
+ return f"You should provide path to your file"
+ file_name = args[2]
+
+ if file_type == "parquet":
+ self._ctx.register_parquet(table_name, file_name)
+ elif file_type == "csv":
+ self._ctx.register_csv(table_name, file_name)
+ else:
+ raise NotImplemented("Currently not supporting the inserted file format")
+
+ @line_cell_magic
+ def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]:
+ """
+ Execute a SQL query (both line and cell magic).
+
+ Two cases possible: with cell or without cell
+
+ Examples:
+ 1. Without a cell (line_magic)
+ %sql SELECT * FROM test_table
+ 2. With a cell (cell_magic)
+ %%sql --no-display
+ SELECT
+ id,
+ bool_col,
+ tinyint_col
+ FROM test_data_v1
+ WHERE id > 2
+ ORDER BY id
+ LIMIT 5
+ """
+ if not cell:
+ return self._execute_sql(line.strip()) if line.strip() else None
+ else:
+ args = line.strip().split()
+ display_results = True
+ limit = 50
+ var_name = None
+
+ i = 0
+ while i < len(args):
+ if args[i] == "--no-display":
+ display_results = False
+ elif args[i] == "--limit" and i + 1 < len(args):
+ try:
+ limit = int(args[i + 1])
+ i += 1
+ except ValueError:
+ pass
+ elif not args[i].startswith("--"):
+ var_name = args[i]
+ i += 1
+
+ query = cell.strip()
+ if not query:
+ return None
+
+ result = self._execute_sql(query, display_results=display_results, limit=limit)
+
+ # Store in user namespace if variable name provided
+ if var_name and self.shell is not None:
+ self.shell.user_ns[var_name] = result
+
+ return result
+
+ def _connect(self, address: str) -> Optional[str]:
+ """Connect to a Ballista cluster."""
+ if not address:
+ return "Usage: %ballista connect df://host:port"
+
+ # Normalize address
+ if not address.startswith("df://"):
+ address = f"df://{address}"
+
+ try:
+ self._ctx = BallistaSessionContext(address)
+ self._address = address
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Connected to Ballista cluster at {address}"))
+ else:
+ return f"✓ Connected to Ballista cluster at {address}"
+ except Exception as e:
+ self._ctx = None
+ self._address = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✗ Failed to connect to {address}: {e}"))
+ else:
+ return f"✗ Failed to connect to {address}: {e}"
+
+ def _disconnect(self) -> Optional[str]:
+ """Disconnect from the Ballista cluster."""
+ if not self.is_connected:
+ return "Not connected to any cluster."
+
+ address = self._address
+ self._ctx = None
+ self._address = None
+ self._last_result = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Disconnected from {address}"))
+ else:
+ return f"✓ Disconnected from {address}"
+
+ def _status(self) -> Optional[str]:
+ """Show connection status."""
+ if not self.is_connected:
+ return "Status: Not connected\n\nUse '%ballista connect df://host:port' to connect."
+
+ status_lines = [
+ f"Status: Connected",
+ f"Address: {self._address}",
+ f"Session ID: {self._ctx.session_id}",
+ f"Queries executed: {len(self._query_history)}",
+ ]
+
+ if self._last_result is not None:
+ status_lines.append("Last result: Available (access via '_' or '_last_result')")
+
+ def _format_html_status_output(line: str) -> str:
+ name, value = line.split(":", 1)
+ return f"<p><strong>{name}:</strong> {value.strip()}</p>"
+
+ html = "".join(_format_html_status_output(line) for line in status_lines)
+ if IPYTHON_AVAILABLE:
+ display(HTML(html))
+ else:
+ print("\n".join(status_lines))
+
+ def _tables(self) -> Optional[str]:
+ """List all registered tables."""
+ try:
+ # Get table names from the catalog and their respective schemas
+ tables = self.ctx.get_tables()
+ 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())
+ # Build a nice table display (HTML-formatted if applicable)
+ lines = [
+ {"content": f"Total: {table_count} table(s) in {schema_count} schema(s)", "is_info": True},
+ {"content": "Registered tables:", "is_info": True},
+ *[{"content": f"Schema: {schema_name}. Tables: {", ".join(list(table_names))}", "is_info": False}
+ for schema_name, table_names in tables.items()]
+ ]
+
+ def _format_html_tables_output(line: str, is_info: bool = False) -> str:
+ if is_info:
+ return f"<p><strong>{line}<strong></p>"
+ else:
+ return f"<p><pre><i>{line}</i></pre></p>"
+
+ if IPYTHON_AVAILABLE:
+ display(
+ HTML("".join(_format_html_tables_output(val["content"], val["is_info"]) for val in lines))
+ )
+ else:
+ return "".join(val["content"] for val in lines)
+ except Exception as e:
+ return f"Error listing tables: {e}"
+
+ def _schema(self, table_name: str) -> Optional[str]:
+ """Show schema for a table."""
+ if not table_name:
+ return "Usage: %ballista schema <table_name>"
+
+ try:
+ # Query the table with LIMIT 0 to get schema without data
+ df = self.ctx.sql(f"SELECT * FROM {table_name} LIMIT 0")
+ schema = df.schema()
+
+ lines = [f"Schema for '{table_name}':", "-" * 50]
+ for field in schema:
+ nullable = "NULL" if field.nullable else "NOT NULL"
+ lines.append(f" {field.name:20} {str(field.type):15} {nullable}")
+ lines.append("-" * 50)
+ lines.append(f"Total: {len(schema)} column(s)")
+
+ return "\n".join(lines)
+ except Exception as e:
+ return f"Error getting schema for '{table_name}': {e}"
+
+ def _execute_sql(
+ self,
+ query: str,
+ display_results: bool = True,
+ limit: int = 50,
+ ) -> Optional[DistributedDataFrame]:
+ """Execute a SQL query and return the result."""
+ start_time = time.time()
+
+ try:
+ result = self.ctx.sql(query)
+ elapsed = time.time() - start_time
+
+ # Store result
+ self._last_result = result
+ if self.shell is not None and hasattr(self.shell, 'user_ns'):
+ self.shell.user_ns["_last_result"] = result
+
+ # Record in history
+ self._query_history.append(
+ {
+ "query": query,
+ "elapsed_seconds": elapsed,
+ "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
+ }
+ )
+
+ # Display if requested and in notebook environment
+ return result
+
+ except Exception as e:
+ error_msg = f"Query failed: {e}"
+ if IPYTHON_AVAILABLE:
+ try:
+ display(HTML(f'<div style="color: red; font-weight: bold;">{error_msg}</div>'))
+ except Exception:
+ print(error_msg)
+ else:
+ print(error_msg)
+ return None
+
+ def _show_history(self) -> Optional[str]:
+ """Show query history."""
+ if not self._query_history:
+ return "No queries executed yet."
+
+ lines = ["Query History:", "-" * 60]
+ for i, entry in enumerate(self._query_history[-10:], 1): # Last 10 queries
+ query_preview = entry["query"][:50] + "..." if len(entry["query"]) > 50 else entry["query"]
+ query_preview = query_preview.replace("\n", " ")
+ lines.append(f"{i}. [{entry['timestamp']}] ({entry['elapsed_seconds']:.2f}s)")
+ lines.append(f" {query_preview}")
+ lines.append("-" * 60)
+
+ for row in lines:
+ print(row)
+
+ def _show_help(self) -> Optional[str]:
Review Comment:
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
##########
python/python/ballista/jupyter.py:
##########
@@ -0,0 +1,493 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Ballista Jupyter Magic Commands.
+
+This module provides IPython magic commands for interacting with Ballista
+clusters directly from Jupyter notebooks.
+
+Usage:
+ %load_ext ballista.jupyter
+
+ # Connect to a Ballista cluster
+ %ballista connect df://localhost:50050
+
+ # Check connection status
+ %ballista status
+
+ # List registered tables
+ %ballista tables
+
+ # Show schema for a table
+ %ballista schema my_table
+
+ # Execute a simple SQL query (line magic)
+ %sql SELECT COUNT(*) FROM my_table
+
+ # Execute a complex SQL query (cell magic)
+ %%sql
+ SELECT
+ customer_id,
+ SUM(amount) as total
+ FROM orders
+ GROUP BY customer_id
+ ORDER BY total DESC
+ LIMIT 10
+"""
+
+from typing import Optional, List, Dict, Any
+import time
+import threading
+
+try:
+ from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic
+ from IPython.display import display, HTML
+
+ IPYTHON_AVAILABLE = True
+except ImportError:
+ IPYTHON_AVAILABLE = False
+
+ # Provide stub classes for when IPython is not available
+ class Magics:
+ def __init__(self, shell=None):
+ self.shell = shell
+
+ def magics_class(cls):
+ return cls
+
+ def line_magic(name_or_func):
+ """Stub line_magic decorator for when IPython is not available."""
+ # Handle both @line_magic and @line_magic("name") usage
+ if callable(name_or_func):
+ # Used as @line_magic without arguments
+ return name_or_func
+ else:
+ # Used as @line_magic("name") with arguments
+ def decorator(func):
+ return func
+ return decorator
+
+ def cell_magic(name_or_func):
+ """Stub cell_magic decorator for when IPython is not available."""
+ # Handle both @cell_magic and @cell_magic("name") usage
+ if callable(name_or_func):
+ return name_or_func
+ else:
+ def decorator(func):
+ return func
+ return decorator
+
+
+from .extension import BallistaSessionContext, DistributedDataFrame
+
+
+class BallistaConnectionError(Exception):
+ """Raised when not connected to a Ballista cluster."""
+
+ pass
+
+
+@magics_class
+class BallistaMagics(Magics):
+ """
+ IPython magic commands for Ballista.
+
+ Provides convenient commands for connecting to Ballista clusters,
+ executing SQL queries, and exploring table schemas.
+ """
+
+ def __init__(self, shell=None):
+ super().__init__(shell)
+ self._ctx: Optional[BallistaSessionContext] = None
+ self._address: Optional[str] = None
+ self._last_result: Optional[DistributedDataFrame] = None
+ self._query_history: List[Dict[str, Any]] = []
+
+ @property
+ def ctx(self) -> BallistaSessionContext:
+ """Get the current context, raising an error if not connected."""
+ if self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ return self._ctx
+
+ @property
+ def is_connected(self) -> bool:
+ """Check if connected to a Ballista cluster."""
+ return self._ctx is not None
+
+ @line_magic
+ def ballista(self, line: str) -> Optional[str]:
+ """
+ Ballista management commands.
+
+ Usage:
+ %ballista connect df://localhost:50050 - Connect to cluster
+ %ballista status - Show connection status
+ %ballista tables - List registered tables
+ %ballista schema <table> - Show table schema
+ %ballista disconnect - Disconnect from cluster
+ %ballista history - Show query history
+
+ Examples:
+ %ballista connect df://localhost:50050
+ %ballista tables
+ %ballista schema orders
+ """
+ parts = line.strip().split(maxsplit=1)
+ if not parts:
+ return self._show_help()
+
+ cmd = parts[0].lower()
+ args = parts[1] if len(parts) > 1 else ""
+
+ if cmd == "connect":
+ return self._connect(args)
+ elif cmd == "status":
+ return self._status()
+ elif cmd == "tables":
+ return self._tables()
+ elif cmd == "schema":
+ return self._schema(args)
+ elif cmd == "disconnect":
+ return self._disconnect()
+ elif cmd == "history":
+ return self._show_history()
+ elif cmd == "help":
+ return self._show_help()
+ else:
+ return f"Unknown command: {cmd}. Use '%ballista help' for available commands."
+
+ @line_magic
+ def register(self, line: str) -> Optional[str]:
+ """Register a new table"""
+ if not line:
+ return "You should provide file extension and table name to register"
+ elif self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ else:
+ args = line.strip().split()
+ file_type = args[0]
+ if len(args) < 2:
+ return f"You should provide table name for this .{file_type} file"
+ table_name = args[1]
+ if len(args) < 3:
+ return f"You should provide path to your file"
+ file_name = args[2]
+
+ if file_type == "parquet":
+ self._ctx.register_parquet(table_name, file_name)
+ elif file_type == "csv":
+ self._ctx.register_csv(table_name, file_name)
+ else:
+ raise NotImplemented("Currently not supporting the inserted file format")
+
+ @line_cell_magic
+ def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]:
+ """
+ Execute a SQL query (both line and cell magic).
+
+ Two cases possible: with cell or without cell
+
+ Examples:
+ 1. Without a cell (line_magic)
+ %sql SELECT * FROM test_table
+ 2. With a cell (cell_magic)
+ %%sql --no-display
+ SELECT
+ id,
+ bool_col,
+ tinyint_col
+ FROM test_data_v1
+ WHERE id > 2
+ ORDER BY id
+ LIMIT 5
+ """
+ if not cell:
+ return self._execute_sql(line.strip()) if line.strip() else None
+ else:
+ args = line.strip().split()
+ display_results = True
+ limit = 50
+ var_name = None
+
+ i = 0
+ while i < len(args):
+ if args[i] == "--no-display":
+ display_results = False
+ elif args[i] == "--limit" and i + 1 < len(args):
+ try:
+ limit = int(args[i + 1])
+ i += 1
+ except ValueError:
+ pass
+ elif not args[i].startswith("--"):
+ var_name = args[i]
+ i += 1
+
+ query = cell.strip()
+ if not query:
+ return None
+
+ result = self._execute_sql(query, display_results=display_results, limit=limit)
+
+ # Store in user namespace if variable name provided
+ if var_name and self.shell is not None:
+ self.shell.user_ns[var_name] = result
+
+ return result
+
+ def _connect(self, address: str) -> Optional[str]:
+ """Connect to a Ballista cluster."""
+ if not address:
+ return "Usage: %ballista connect df://host:port"
+
+ # Normalize address
+ if not address.startswith("df://"):
+ address = f"df://{address}"
+
+ try:
+ self._ctx = BallistaSessionContext(address)
+ self._address = address
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Connected to Ballista cluster at {address}"))
+ else:
+ return f"✓ Connected to Ballista cluster at {address}"
+ except Exception as e:
+ self._ctx = None
+ self._address = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✗ Failed to connect to {address}: {e}"))
+ else:
+ return f"✗ Failed to connect to {address}: {e}"
+
+ def _disconnect(self) -> Optional[str]:
+ """Disconnect from the Ballista cluster."""
+ if not self.is_connected:
+ return "Not connected to any cluster."
+
+ address = self._address
+ self._ctx = None
+ self._address = None
+ self._last_result = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Disconnected from {address}"))
+ else:
+ return f"✓ Disconnected from {address}"
+
+ def _status(self) -> Optional[str]:
+ """Show connection status."""
+ if not self.is_connected:
+ return "Status: Not connected\n\nUse '%ballista connect df://host:port' to connect."
+
+ status_lines = [
+ f"Status: Connected",
+ f"Address: {self._address}",
+ f"Session ID: {self._ctx.session_id}",
+ f"Queries executed: {len(self._query_history)}",
+ ]
+
+ if self._last_result is not None:
+ status_lines.append("Last result: Available (access via '_' or '_last_result')")
+
+ def _format_html_status_output(line: str) -> str:
+ name, value = line.split(":", 1)
+ return f"<p><strong>{name}:</strong> {value.strip()}</p>"
+
+ html = "".join(_format_html_status_output(line) for line in status_lines)
+ if IPYTHON_AVAILABLE:
+ display(HTML(html))
+ else:
+ print("\n".join(status_lines))
+
+ def _tables(self) -> Optional[str]:
+ """List all registered tables."""
+ try:
+ # Get table names from the catalog and their respective schemas
+ tables = self.ctx.get_tables()
+ 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())
+ # Build a nice table display (HTML-formatted if applicable)
+ lines = [
+ {"content": f"Total: {table_count} table(s) in {schema_count} schema(s)", "is_info": True},
+ {"content": "Registered tables:", "is_info": True},
+ *[{"content": f"Schema: {schema_name}. Tables: {", ".join(list(table_names))}", "is_info": False}
+ for schema_name, table_names in tables.items()]
+ ]
+
+ def _format_html_tables_output(line: str, is_info: bool = False) -> str:
+ if is_info:
+ return f"<p><strong>{line}<strong></p>"
+ else:
+ return f"<p><pre><i>{line}</i></pre></p>"
+
+ if IPYTHON_AVAILABLE:
+ display(
+ HTML("".join(_format_html_tables_output(val["content"], val["is_info"]) for val in lines))
+ )
+ else:
+ return "".join(val["content"] for val in lines)
+ except Exception as e:
+ return f"Error listing tables: {e}"
+
+ def _schema(self, table_name: str) -> Optional[str]:
+ """Show schema for a table."""
+ if not table_name:
+ return "Usage: %ballista schema <table_name>"
+
+ try:
+ # Query the table with LIMIT 0 to get schema without data
+ df = self.ctx.sql(f"SELECT * FROM {table_name} LIMIT 0")
+ schema = df.schema()
+
+ lines = [f"Schema for '{table_name}':", "-" * 50]
+ for field in schema:
+ nullable = "NULL" if field.nullable else "NOT NULL"
+ lines.append(f" {field.name:20} {str(field.type):15} {nullable}")
+ lines.append("-" * 50)
+ lines.append(f"Total: {len(schema)} column(s)")
+
+ return "\n".join(lines)
+ except Exception as e:
+ return f"Error getting schema for '{table_name}': {e}"
+
+ def _execute_sql(
+ self,
+ query: str,
+ display_results: bool = True,
+ limit: int = 50,
Review Comment:
`display_results` and `limit` seem unused in the method body
##########
python/python/ballista/extension.py:
##########
@@ -229,9 +231,319 @@ def write_parquet(
df = self._to_internal_df()
df.write_parquet(str(path), compression.value, compression_level)
+ def explain_visual(self, analyze: bool = False) -> "ExecutionPlanVisualization":
+ """
+ Generate a visual representation of the execution plan.
+
+ This method creates an SVG visualization of the query execution plan,
+ which can be displayed directly in Jupyter notebooks.
+
+ Args:
+ analyze: If True, includes runtime statistics from actual execution.
+
+ Returns:
+ ExecutionPlanVisualization: An object that renders as SVG in Jupyter.
+
+ Example:
+ >>> df = ctx.sql("SELECT * FROM orders WHERE amount > 100")
+ >>> df.explain_visual() # Displays SVG in notebook
+ >>> viz = df.explain_visual(analyze=True)
+ >>> viz.save("plan.svg") # Save to file
+ """
+ # Get the execution plan as a string representation
+ # Note: explain() prints but doesn't return a string, so we use logical_plan()
+ try:
+ plan = self.logical_plan()
+ plan_str = plan.display_indent()
+ except Exception:
+ # Fallback if logical_plan() fails
+ plan_str = "Unable to retrieve execution plan"
+ return ExecutionPlanVisualization(plan_str, analyze=analyze)
+
+ def collect_with_progress(
+ self,
+ callback: Optional[callable] = None,
Review Comment:
```suggestion
callback: Optional[Callable] = None,
```
##########
python/python/ballista/jupyter.py:
##########
@@ -0,0 +1,493 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Ballista Jupyter Magic Commands.
+
+This module provides IPython magic commands for interacting with Ballista
+clusters directly from Jupyter notebooks.
+
+Usage:
+ %load_ext ballista.jupyter
+
+ # Connect to a Ballista cluster
+ %ballista connect df://localhost:50050
+
+ # Check connection status
+ %ballista status
+
+ # List registered tables
+ %ballista tables
+
+ # Show schema for a table
+ %ballista schema my_table
+
+ # Execute a simple SQL query (line magic)
+ %sql SELECT COUNT(*) FROM my_table
+
+ # Execute a complex SQL query (cell magic)
+ %%sql
+ SELECT
+ customer_id,
+ SUM(amount) as total
+ FROM orders
+ GROUP BY customer_id
+ ORDER BY total DESC
+ LIMIT 10
+"""
+
+from typing import Optional, List, Dict, Any
+import time
+import threading
+
+try:
+ from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic
+ from IPython.display import display, HTML
+
+ IPYTHON_AVAILABLE = True
+except ImportError:
+ IPYTHON_AVAILABLE = False
+
+ # Provide stub classes for when IPython is not available
+ class Magics:
+ def __init__(self, shell=None):
+ self.shell = shell
+
+ def magics_class(cls):
+ return cls
+
+ def line_magic(name_or_func):
+ """Stub line_magic decorator for when IPython is not available."""
+ # Handle both @line_magic and @line_magic("name") usage
+ if callable(name_or_func):
+ # Used as @line_magic without arguments
+ return name_or_func
+ else:
+ # Used as @line_magic("name") with arguments
+ def decorator(func):
+ return func
+ return decorator
+
+ def cell_magic(name_or_func):
+ """Stub cell_magic decorator for when IPython is not available."""
+ # Handle both @cell_magic and @cell_magic("name") usage
+ if callable(name_or_func):
+ return name_or_func
+ else:
+ def decorator(func):
+ return func
+ return decorator
+
+
+from .extension import BallistaSessionContext, DistributedDataFrame
+
+
+class BallistaConnectionError(Exception):
+ """Raised when not connected to a Ballista cluster."""
+
+ pass
+
+
+@magics_class
+class BallistaMagics(Magics):
+ """
+ IPython magic commands for Ballista.
+
+ Provides convenient commands for connecting to Ballista clusters,
+ executing SQL queries, and exploring table schemas.
+ """
+
+ def __init__(self, shell=None):
+ super().__init__(shell)
+ self._ctx: Optional[BallistaSessionContext] = None
+ self._address: Optional[str] = None
+ self._last_result: Optional[DistributedDataFrame] = None
+ self._query_history: List[Dict[str, Any]] = []
+
+ @property
+ def ctx(self) -> BallistaSessionContext:
+ """Get the current context, raising an error if not connected."""
+ if self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ return self._ctx
+
+ @property
+ def is_connected(self) -> bool:
+ """Check if connected to a Ballista cluster."""
+ return self._ctx is not None
+
+ @line_magic
+ def ballista(self, line: str) -> Optional[str]:
+ """
+ Ballista management commands.
+
+ Usage:
+ %ballista connect df://localhost:50050 - Connect to cluster
+ %ballista status - Show connection status
+ %ballista tables - List registered tables
+ %ballista schema <table> - Show table schema
+ %ballista disconnect - Disconnect from cluster
+ %ballista history - Show query history
+
+ Examples:
+ %ballista connect df://localhost:50050
+ %ballista tables
+ %ballista schema orders
+ """
+ parts = line.strip().split(maxsplit=1)
+ if not parts:
+ return self._show_help()
+
+ cmd = parts[0].lower()
+ args = parts[1] if len(parts) > 1 else ""
+
+ if cmd == "connect":
+ return self._connect(args)
+ elif cmd == "status":
+ return self._status()
+ elif cmd == "tables":
+ return self._tables()
+ elif cmd == "schema":
+ return self._schema(args)
+ elif cmd == "disconnect":
+ return self._disconnect()
+ elif cmd == "history":
+ return self._show_history()
+ elif cmd == "help":
+ return self._show_help()
+ else:
+ return f"Unknown command: {cmd}. Use '%ballista help' for available commands."
+
+ @line_magic
+ def register(self, line: str) -> Optional[str]:
+ """Register a new table"""
+ if not line:
+ return "You should provide file extension and table name to register"
+ elif self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ else:
+ args = line.strip().split()
+ file_type = args[0]
+ if len(args) < 2:
+ return f"You should provide table name for this .{file_type} file"
+ table_name = args[1]
+ if len(args) < 3:
+ return f"You should provide path to your file"
+ file_name = args[2]
+
+ if file_type == "parquet":
+ self._ctx.register_parquet(table_name, file_name)
+ elif file_type == "csv":
+ self._ctx.register_csv(table_name, file_name)
+ else:
+ raise NotImplemented("Currently not supporting the inserted file format")
+
+ @line_cell_magic
+ def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]:
+ """
+ Execute a SQL query (both line and cell magic).
+
+ Two cases possible: with cell or without cell
+
+ Examples:
+ 1. Without a cell (line_magic)
+ %sql SELECT * FROM test_table
+ 2. With a cell (cell_magic)
+ %%sql --no-display
+ SELECT
+ id,
+ bool_col,
+ tinyint_col
+ FROM test_data_v1
+ WHERE id > 2
+ ORDER BY id
+ LIMIT 5
+ """
+ if not cell:
+ return self._execute_sql(line.strip()) if line.strip() else None
+ else:
+ args = line.strip().split()
+ display_results = True
+ limit = 50
+ var_name = None
+
+ i = 0
+ while i < len(args):
+ if args[i] == "--no-display":
+ display_results = False
+ elif args[i] == "--limit" and i + 1 < len(args):
+ try:
+ limit = int(args[i + 1])
+ i += 1
+ except ValueError:
+ pass
+ elif not args[i].startswith("--"):
+ var_name = args[i]
+ i += 1
+
+ query = cell.strip()
+ if not query:
+ return None
+
+ result = self._execute_sql(query, display_results=display_results, limit=limit)
+
+ # Store in user namespace if variable name provided
+ if var_name and self.shell is not None:
+ self.shell.user_ns[var_name] = result
+
+ return result
+
+ def _connect(self, address: str) -> Optional[str]:
+ """Connect to a Ballista cluster."""
+ if not address:
+ return "Usage: %ballista connect df://host:port"
+
+ # Normalize address
+ if not address.startswith("df://"):
+ address = f"df://{address}"
+
+ try:
+ self._ctx = BallistaSessionContext(address)
+ self._address = address
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Connected to Ballista cluster at {address}"))
+ else:
+ return f"✓ Connected to Ballista cluster at {address}"
+ except Exception as e:
+ self._ctx = None
+ self._address = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✗ Failed to connect to {address}: {e}"))
+ else:
+ return f"✗ Failed to connect to {address}: {e}"
+
+ def _disconnect(self) -> Optional[str]:
+ """Disconnect from the Ballista cluster."""
+ if not self.is_connected:
+ return "Not connected to any cluster."
+
+ address = self._address
+ self._ctx = None
+ self._address = None
+ self._last_result = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Disconnected from {address}"))
+ else:
+ return f"✓ Disconnected from {address}"
+
+ def _status(self) -> Optional[str]:
+ """Show connection status."""
+ if not self.is_connected:
+ return "Status: Not connected\n\nUse '%ballista connect df://host:port' to connect."
+
+ status_lines = [
+ f"Status: Connected",
+ f"Address: {self._address}",
+ f"Session ID: {self._ctx.session_id}",
+ f"Queries executed: {len(self._query_history)}",
+ ]
+
+ if self._last_result is not None:
+ status_lines.append("Last result: Available (access via '_' or '_last_result')")
+
+ def _format_html_status_output(line: str) -> str:
+ name, value = line.split(":", 1)
+ return f"<p><strong>{name}:</strong> {value.strip()}</p>"
+
+ html = "".join(_format_html_status_output(line) for line in status_lines)
+ if IPYTHON_AVAILABLE:
+ display(HTML(html))
+ else:
+ print("\n".join(status_lines))
+
+ def _tables(self) -> Optional[str]:
+ """List all registered tables."""
+ try:
+ # Get table names from the catalog and their respective schemas
+ tables = self.ctx.get_tables()
+ 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())
+ # Build a nice table display (HTML-formatted if applicable)
+ lines = [
+ {"content": f"Total: {table_count} table(s) in {schema_count} schema(s)", "is_info": True},
+ {"content": "Registered tables:", "is_info": True},
+ *[{"content": f"Schema: {schema_name}. Tables: {", ".join(list(table_names))}", "is_info": False}
Review Comment:
The quotes here look wrong.
##########
python/python/ballista/jupyter.py:
##########
@@ -0,0 +1,493 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Ballista Jupyter Magic Commands.
+
+This module provides IPython magic commands for interacting with Ballista
+clusters directly from Jupyter notebooks.
+
+Usage:
+ %load_ext ballista.jupyter
+
+ # Connect to a Ballista cluster
+ %ballista connect df://localhost:50050
+
+ # Check connection status
+ %ballista status
+
+ # List registered tables
+ %ballista tables
+
+ # Show schema for a table
+ %ballista schema my_table
+
+ # Execute a simple SQL query (line magic)
+ %sql SELECT COUNT(*) FROM my_table
+
+ # Execute a complex SQL query (cell magic)
+ %%sql
+ SELECT
+ customer_id,
+ SUM(amount) as total
+ FROM orders
+ GROUP BY customer_id
+ ORDER BY total DESC
+ LIMIT 10
+"""
+
+from typing import Optional, List, Dict, Any
+import time
+import threading
+
+try:
+ from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic
+ from IPython.display import display, HTML
+
+ IPYTHON_AVAILABLE = True
+except ImportError:
+ IPYTHON_AVAILABLE = False
+
+ # Provide stub classes for when IPython is not available
+ class Magics:
+ def __init__(self, shell=None):
+ self.shell = shell
+
+ def magics_class(cls):
+ return cls
+
+ def line_magic(name_or_func):
+ """Stub line_magic decorator for when IPython is not available."""
+ # Handle both @line_magic and @line_magic("name") usage
+ if callable(name_or_func):
+ # Used as @line_magic without arguments
+ return name_or_func
+ else:
+ # Used as @line_magic("name") with arguments
+ def decorator(func):
+ return func
+ return decorator
+
+ def cell_magic(name_or_func):
+ """Stub cell_magic decorator for when IPython is not available."""
+ # Handle both @cell_magic and @cell_magic("name") usage
+ if callable(name_or_func):
+ return name_or_func
+ else:
+ def decorator(func):
+ return func
+ return decorator
+
+
+from .extension import BallistaSessionContext, DistributedDataFrame
+
+
+class BallistaConnectionError(Exception):
+ """Raised when not connected to a Ballista cluster."""
+
+ pass
+
+
+@magics_class
+class BallistaMagics(Magics):
+ """
+ IPython magic commands for Ballista.
+
+ Provides convenient commands for connecting to Ballista clusters,
+ executing SQL queries, and exploring table schemas.
+ """
+
+ def __init__(self, shell=None):
+ super().__init__(shell)
+ self._ctx: Optional[BallistaSessionContext] = None
+ self._address: Optional[str] = None
+ self._last_result: Optional[DistributedDataFrame] = None
+ self._query_history: List[Dict[str, Any]] = []
+
+ @property
+ def ctx(self) -> BallistaSessionContext:
+ """Get the current context, raising an error if not connected."""
+ if self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ return self._ctx
+
+ @property
+ def is_connected(self) -> bool:
+ """Check if connected to a Ballista cluster."""
+ return self._ctx is not None
+
+ @line_magic
+ def ballista(self, line: str) -> Optional[str]:
+ """
+ Ballista management commands.
+
+ Usage:
+ %ballista connect df://localhost:50050 - Connect to cluster
+ %ballista status - Show connection status
+ %ballista tables - List registered tables
+ %ballista schema <table> - Show table schema
+ %ballista disconnect - Disconnect from cluster
+ %ballista history - Show query history
+
+ Examples:
+ %ballista connect df://localhost:50050
+ %ballista tables
+ %ballista schema orders
+ """
+ parts = line.strip().split(maxsplit=1)
+ if not parts:
+ return self._show_help()
+
+ cmd = parts[0].lower()
+ args = parts[1] if len(parts) > 1 else ""
+
+ if cmd == "connect":
+ return self._connect(args)
+ elif cmd == "status":
+ return self._status()
+ elif cmd == "tables":
+ return self._tables()
+ elif cmd == "schema":
+ return self._schema(args)
+ elif cmd == "disconnect":
+ return self._disconnect()
+ elif cmd == "history":
+ return self._show_history()
+ elif cmd == "help":
+ return self._show_help()
+ else:
+ return f"Unknown command: {cmd}. Use '%ballista help' for available commands."
+
+ @line_magic
+ def register(self, line: str) -> Optional[str]:
+ """Register a new table"""
+ if not line:
+ return "You should provide file extension and table name to register"
+ elif self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ else:
+ args = line.strip().split()
+ file_type = args[0]
+ if len(args) < 2:
+ return f"You should provide table name for this .{file_type} file"
+ table_name = args[1]
+ if len(args) < 3:
+ return f"You should provide path to your file"
+ file_name = args[2]
+
+ if file_type == "parquet":
+ self._ctx.register_parquet(table_name, file_name)
+ elif file_type == "csv":
+ self._ctx.register_csv(table_name, file_name)
+ else:
+ raise NotImplemented("Currently not supporting the inserted file format")
+
+ @line_cell_magic
+ def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]:
+ """
+ Execute a SQL query (both line and cell magic).
+
+ Two cases possible: with cell or without cell
+
+ Examples:
+ 1. Without a cell (line_magic)
+ %sql SELECT * FROM test_table
+ 2. With a cell (cell_magic)
+ %%sql --no-display
+ SELECT
+ id,
+ bool_col,
+ tinyint_col
+ FROM test_data_v1
+ WHERE id > 2
+ ORDER BY id
+ LIMIT 5
+ """
+ if not cell:
+ return self._execute_sql(line.strip()) if line.strip() else None
+ else:
+ args = line.strip().split()
+ display_results = True
+ limit = 50
+ var_name = None
+
+ i = 0
+ while i < len(args):
+ if args[i] == "--no-display":
+ display_results = False
+ elif args[i] == "--limit" and i + 1 < len(args):
+ try:
+ limit = int(args[i + 1])
+ i += 1
+ except ValueError:
+ pass
+ elif not args[i].startswith("--"):
+ var_name = args[i]
+ i += 1
+
+ query = cell.strip()
+ if not query:
+ return None
+
+ result = self._execute_sql(query, display_results=display_results, limit=limit)
+
+ # Store in user namespace if variable name provided
+ if var_name and self.shell is not None:
+ self.shell.user_ns[var_name] = result
+
+ return result
+
+ def _connect(self, address: str) -> Optional[str]:
+ """Connect to a Ballista cluster."""
+ if not address:
+ return "Usage: %ballista connect df://host:port"
+
+ # Normalize address
+ if not address.startswith("df://"):
+ address = f"df://{address}"
+
+ try:
+ self._ctx = BallistaSessionContext(address)
+ self._address = address
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Connected to Ballista cluster at {address}"))
+ else:
+ return f"✓ Connected to Ballista cluster at {address}"
+ except Exception as e:
+ self._ctx = None
+ self._address = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✗ Failed to connect to {address}: {e}"))
+ else:
+ return f"✗ Failed to connect to {address}: {e}"
+
+ def _disconnect(self) -> Optional[str]:
+ """Disconnect from the Ballista cluster."""
+ if not self.is_connected:
+ return "Not connected to any cluster."
+
+ address = self._address
+ self._ctx = None
+ self._address = None
+ self._last_result = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Disconnected from {address}"))
+ else:
+ return f"✓ Disconnected from {address}"
+
+ def _status(self) -> Optional[str]:
+ """Show connection status."""
+ if not self.is_connected:
+ return "Status: Not connected\n\nUse '%ballista connect df://host:port' to connect."
+
+ status_lines = [
+ f"Status: Connected",
+ f"Address: {self._address}",
+ f"Session ID: {self._ctx.session_id}",
+ f"Queries executed: {len(self._query_history)}",
+ ]
+
+ if self._last_result is not None:
+ status_lines.append("Last result: Available (access via '_' or '_last_result')")
+
+ def _format_html_status_output(line: str) -> str:
+ name, value = line.split(":", 1)
+ return f"<p><strong>{name}:</strong> {value.strip()}</p>"
+
+ html = "".join(_format_html_status_output(line) for line in status_lines)
+ if IPYTHON_AVAILABLE:
+ display(HTML(html))
+ else:
+ print("\n".join(status_lines))
+
+ def _tables(self) -> Optional[str]:
+ """List all registered tables."""
+ try:
+ # Get table names from the catalog and their respective schemas
+ tables = self.ctx.get_tables()
+ 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())
+ # Build a nice table display (HTML-formatted if applicable)
+ lines = [
+ {"content": f"Total: {table_count} table(s) in {schema_count} schema(s)", "is_info": True},
+ {"content": "Registered tables:", "is_info": True},
+ *[{"content": f"Schema: {schema_name}. Tables: {", ".join(list(table_names))}", "is_info": False}
+ for schema_name, table_names in tables.items()]
+ ]
+
+ def _format_html_tables_output(line: str, is_info: bool = False) -> str:
+ if is_info:
+ return f"<p><strong>{line}<strong></p>"
+ else:
+ return f"<p><pre><i>{line}</i></pre></p>"
+
+ if IPYTHON_AVAILABLE:
+ display(
+ HTML("".join(_format_html_tables_output(val["content"], val["is_info"]) for val in lines))
+ )
+ else:
+ return "".join(val["content"] for val in lines)
+ except Exception as e:
+ return f"Error listing tables: {e}"
+
+ def _schema(self, table_name: str) -> Optional[str]:
+ """Show schema for a table."""
+ if not table_name:
+ return "Usage: %ballista schema <table_name>"
+
+ try:
+ # Query the table with LIMIT 0 to get schema without data
+ df = self.ctx.sql(f"SELECT * FROM {table_name} LIMIT 0")
Review Comment:
Should `table_name` be sanitized ? Currently it could lead to SQL injection
##########
python/python/ballista/extension.py:
##########
@@ -229,9 +231,319 @@ def write_parquet(
df = self._to_internal_df()
df.write_parquet(str(path), compression.value, compression_level)
+ def explain_visual(self, analyze: bool = False) -> "ExecutionPlanVisualization":
+ """
+ Generate a visual representation of the execution plan.
+
+ This method creates an SVG visualization of the query execution plan,
+ which can be displayed directly in Jupyter notebooks.
+
+ Args:
+ analyze: If True, includes runtime statistics from actual execution.
+
+ Returns:
+ ExecutionPlanVisualization: An object that renders as SVG in Jupyter.
+
+ Example:
+ >>> df = ctx.sql("SELECT * FROM orders WHERE amount > 100")
+ >>> df.explain_visual() # Displays SVG in notebook
+ >>> viz = df.explain_visual(analyze=True)
+ >>> viz.save("plan.svg") # Save to file
+ """
+ # Get the execution plan as a string representation
+ # Note: explain() prints but doesn't return a string, so we use logical_plan()
+ try:
+ plan = self.logical_plan()
+ plan_str = plan.display_indent()
+ except Exception:
+ # Fallback if logical_plan() fails
+ plan_str = "Unable to retrieve execution plan"
+ return ExecutionPlanVisualization(plan_str, analyze=analyze)
+
+ def collect_with_progress(
+ self,
+ callback: Optional[callable] = None,
+ poll_interval: float = 0.5,
+ ):
+ """
+ Collect results with progress indication.
+
+ For long-running queries, this method provides progress updates
+ through a callback function or displays a progress bar in Jupyter.
+
+ Args:
+ callback: Optional function to call with progress updates.
+ Signature: callback(status: str, progress: float)
+ poll_interval: How often to check progress (seconds).
+
+ Returns:
+ The collected result batches.
+
+ Example:
+ >>> def my_callback(status, progress):
+ ... print(f"{status}: {progress:.1%}")
+ >>> batches = df.collect_with_progress(callback=my_callback)
+ """
+ import threading
+ import time
+
+ result = [None]
+ error = [None]
+ done = threading.Event()
+
+ def execute():
+ try:
+ result[0] = self.collect()
+ except Exception as e:
+ error[0] = e
+ finally:
+ done.set()
+
+ thread = threading.Thread(target=execute)
+ thread.start()
+
+ # Check if we're in a Jupyter environment
+ try:
+ from IPython.display import display, clear_output
+ from IPython import get_ipython
+
+ in_jupyter = get_ipython() is not None
+ except (ImportError, AttributeError):
+ in_jupyter = False
+
+ start_time = time.time()
+
+ if in_jupyter and callback is None:
+ # Display a simple progress indicator
+ try:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ clear_output(wait=True)
+ print(f"⏳ Query executing... ({elapsed:.1f}s elapsed)")
+
+ clear_output(wait=True)
+ elapsed = time.time() - start_time
+ print(f"✓ Query completed in {elapsed:.1f}s")
+ except Exception:
+ pass # Ignore display errors
+ elif callback is not None:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ callback(f"Executing ({elapsed:.1f}s)", -1.0) # -1 means indeterminate
+
+ elapsed = time.time() - start_time
+ callback(f"Completed in {elapsed:.1f}s", 1.0)
+ else:
+ done.wait()
+
+ thread.join()
+
+ if error[0] is not None:
+ raise error[0]
+
+ return result[0]
+
+
+class ExecutionPlanVisualization:
+ """
+ A wrapper for execution plan visualizations that can render as SVG in Jupyter.
+
+ This class takes the text representation of an execution plan and converts
+ it to a Graphviz DOT format, which is then rendered as SVG.
+ """
+
+ def __init__(self, plan_str: str, analyze: bool = False):
+ self.plan_str = plan_str
+ self.analyze = analyze
+ self._svg_cache: Optional[str] = None
+
+ def _parse_plan_to_dot(self) -> str:
+ """Convert the plan string to DOT format for Graphviz."""
+ lines = self.plan_str.strip().split("\n")
+
+ dot_lines = [
+ "digraph ExecutionPlan {",
+ ' rankdir=TB;',
+ ' node [shape=box, style="rounded,filled", fontname="Helvetica"];',
+ ' edge [fontname="Helvetica"];',
+ "",
+ ]
+
+ nodes = []
+ edges = []
+ node_id = 0
+ stack = [] # (indent_level, node_id)
+
+ for line in lines:
+ if not line.strip():
+ continue
+
+ # Calculate indent level
+ indent = len(line) - len(line.lstrip())
+ content = line.strip()
+
+ # Skip non-plan lines
+ if content.startswith("physical_plan") or content.startswith("logical_plan"):
+ continue
+
+ # Create a node for this plan element
+ current_id = node_id
+ node_id += 1
+
+ # Determine node color based on operation type
+ color = "#E3F2FD" # Default light blue
+ if "Scan" in content or "TableScan" in content:
+ color = "#E8F5E9" # Light green for scans
+ elif "Filter" in content:
+ color = "#FFF3E0" # Light orange for filters
+ elif "Aggregate" in content or "HashAggregate" in content:
+ color = "#F3E5F5" # Light purple for aggregations
+ elif "Join" in content:
+ color = "#FFEBEE" # Light red for joins
+ elif "Sort" in content:
+ color = "#E0F7FA" # Light cyan for sorts
+ elif "Projection" in content:
+ color = "#FFF8E1" # Light amber for projections
+
+ # Escape special characters for DOT format
+ label = content.replace('"', '\\"').replace("\n", "\\n")
+ if len(label) > 60:
+ # Wrap long labels
+ label = label[:57] + "..."
+
+ nodes.append(f' node{current_id} [label="{label}", fillcolor="{color}"];')
+
+ # Connect to parent based on indentation
+ while stack and stack[-1][0] >= indent:
+ stack.pop()
+
+ if stack:
+ parent_id = stack[-1][1]
+ edges.append(f" node{parent_id} -> node{current_id};")
+
+ stack.append((indent, current_id))
+
+ dot_lines.extend(nodes)
+ dot_lines.append("")
+ dot_lines.extend(edges)
+ dot_lines.append("}")
+
+ return "\n".join(dot_lines)
+
+ def to_dot(self) -> str:
+ """Get the DOT representation of the execution plan."""
+ return self._parse_plan_to_dot()
+
+ def to_svg(self) -> str:
+ """
+ Convert the execution plan to SVG format.
+
+ Requires graphviz to be installed. If graphviz is not available,
+ returns a simple HTML representation instead.
+ """
+ if self._svg_cache is not None:
+ return self._svg_cache
+
+ dot_source = self._parse_plan_to_dot()
+
+ try:
+ import subprocess
+
+ # Try to use graphviz's dot command
+ process = subprocess.run(
+ ["dot", "-Tsvg"],
+ input=dot_source.encode(),
+ capture_output=True,
+ timeout=30,
+ )
+
+ if process.returncode == 0:
+ self._svg_cache = process.stdout.decode()
+ return self._svg_cache
+ except (subprocess.SubprocessError, FileNotFoundError, subprocess.TimeoutExpired):
+ pass
+
+ # Fallback: return a pre-formatted HTML representation
+ escaped_plan = (
+ self.plan_str.replace("&", "&")
+ .replace("<", "<")
+ .replace(">", ">")
+ )
+ self._svg_cache = f"""
+ <div style="font-family: monospace; background: #f5f5f5; padding: 10px;
+ border-radius: 5px; overflow-x: auto;">
+ <div style="color: #666; margin-bottom: 5px;">
+ Execution Plan {'(with statistics)' if self.analyze else ''}
+ <br><small>Install graphviz for visual diagram: brew install graphviz</small>
+ </div>
+ <pre style="margin: 0;">{escaped_plan}</pre>
+ </div>
+ """
+ return self._svg_cache
+
+ def save(self, path: str) -> None:
+ """Save the visualization to a file (SVG or DOT format)."""
+ if path.endswith(".dot"):
+ content = self.to_dot()
+ else:
+ content = self.to_svg()
+
+ with open(path, "w") as f:
+ f.write(content)
+
+ def _repr_html_(self) -> str:
+ """HTML representation for Jupyter notebooks."""
+ return self.to_svg()
+
+ def _repr_svg_(self) -> str:
+ """SVG representation for Jupyter notebooks."""
+ svg = self.to_svg()
+ # Only return if it's actual SVG content
+ if svg.strip().startswith("<svg") or svg.strip().startswith("<?xml"):
+ return svg
+ return ""
+
+ def __repr__(self) -> str:
+ """String representation."""
+ return f"ExecutionPlanVisualization(analyze={self.analyze})\n{self.plan_str}"
+
class BallistaSessionContext(SessionContext, metaclass=RedefiningSessionContextMeta):
+ """
+ A session context for connecting to and querying a Ballista cluster.
+
+ This class extends DataFusion's SessionContext to work with distributed
+ Ballista clusters, automatically routing query execution to the cluster
+ while maintaining API compatibility with local DataFusion usage.
+
+ Example:
+ >>> from ballista import BallistaSessionContext
+ >>> ctx = BallistaSessionContext("df://localhost:50050")
+ >>> df = ctx.sql("SELECT * FROM my_table LIMIT 10")
+ >>> df.show()
+
+ For Jupyter notebook users:
+ >>> %load_ext ballista.jupyter
+ >>> %ballista connect df://localhost:50050
+ >>> %sql SELECT * FROM my_table
+ """
+
def __init__(self, address: str, config=None, runtime=None):
super().__init__(config, runtime)
self.address = address
self.session_id = self.session_id()
+
+ def get_tables(self) -> Optional[dict[str, str]]:
+ """Get tables and their respective schemas (in terms of database schema)."""
+ try:
+ catalog = self.catalog()
+ schema_names = list(catalog.schema_names())
+ if schema_names is not None:
+ tables_info = {}
+ for schema_name in schema_names:
+ tables_info[schema_name] = list(catalog.schema(name=schema_name).table_names())
+ return tables_info
+ except (AttributeError, NotImplementedError):
+ pass
Review Comment:
```suggestion
import warnings
warnings.warn(f"Could not retrieve tables from catalog: {e}")
pass
```
##########
python/python/ballista/jupyter.py:
##########
@@ -0,0 +1,493 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Ballista Jupyter Magic Commands.
+
+This module provides IPython magic commands for interacting with Ballista
+clusters directly from Jupyter notebooks.
+
+Usage:
+ %load_ext ballista.jupyter
+
+ # Connect to a Ballista cluster
+ %ballista connect df://localhost:50050
+
+ # Check connection status
+ %ballista status
+
+ # List registered tables
+ %ballista tables
+
+ # Show schema for a table
+ %ballista schema my_table
+
+ # Execute a simple SQL query (line magic)
+ %sql SELECT COUNT(*) FROM my_table
+
+ # Execute a complex SQL query (cell magic)
+ %%sql
+ SELECT
+ customer_id,
+ SUM(amount) as total
+ FROM orders
+ GROUP BY customer_id
+ ORDER BY total DESC
+ LIMIT 10
+"""
+
+from typing import Optional, List, Dict, Any
+import time
+import threading
+
+try:
+ from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic
+ from IPython.display import display, HTML
+
+ IPYTHON_AVAILABLE = True
+except ImportError:
+ IPYTHON_AVAILABLE = False
+
+ # Provide stub classes for when IPython is not available
+ class Magics:
+ def __init__(self, shell=None):
+ self.shell = shell
+
+ def magics_class(cls):
+ return cls
+
+ def line_magic(name_or_func):
+ """Stub line_magic decorator for when IPython is not available."""
+ # Handle both @line_magic and @line_magic("name") usage
+ if callable(name_or_func):
+ # Used as @line_magic without arguments
+ return name_or_func
+ else:
+ # Used as @line_magic("name") with arguments
+ def decorator(func):
+ return func
+ return decorator
+
+ def cell_magic(name_or_func):
+ """Stub cell_magic decorator for when IPython is not available."""
+ # Handle both @cell_magic and @cell_magic("name") usage
+ if callable(name_or_func):
+ return name_or_func
+ else:
+ def decorator(func):
+ return func
+ return decorator
+
+
+from .extension import BallistaSessionContext, DistributedDataFrame
+
+
+class BallistaConnectionError(Exception):
+ """Raised when not connected to a Ballista cluster."""
+
+ pass
+
+
+@magics_class
+class BallistaMagics(Magics):
+ """
+ IPython magic commands for Ballista.
+
+ Provides convenient commands for connecting to Ballista clusters,
+ executing SQL queries, and exploring table schemas.
+ """
+
+ def __init__(self, shell=None):
+ super().__init__(shell)
+ self._ctx: Optional[BallistaSessionContext] = None
+ self._address: Optional[str] = None
+ self._last_result: Optional[DistributedDataFrame] = None
+ self._query_history: List[Dict[str, Any]] = []
+
+ @property
+ def ctx(self) -> BallistaSessionContext:
+ """Get the current context, raising an error if not connected."""
+ if self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ return self._ctx
+
+ @property
+ def is_connected(self) -> bool:
+ """Check if connected to a Ballista cluster."""
+ return self._ctx is not None
+
+ @line_magic
+ def ballista(self, line: str) -> Optional[str]:
+ """
+ Ballista management commands.
+
+ Usage:
+ %ballista connect df://localhost:50050 - Connect to cluster
+ %ballista status - Show connection status
+ %ballista tables - List registered tables
+ %ballista schema <table> - Show table schema
+ %ballista disconnect - Disconnect from cluster
+ %ballista history - Show query history
+
+ Examples:
+ %ballista connect df://localhost:50050
+ %ballista tables
+ %ballista schema orders
+ """
+ parts = line.strip().split(maxsplit=1)
+ if not parts:
+ return self._show_help()
+
+ cmd = parts[0].lower()
+ args = parts[1] if len(parts) > 1 else ""
+
+ if cmd == "connect":
+ return self._connect(args)
+ elif cmd == "status":
+ return self._status()
+ elif cmd == "tables":
+ return self._tables()
+ elif cmd == "schema":
+ return self._schema(args)
+ elif cmd == "disconnect":
+ return self._disconnect()
+ elif cmd == "history":
+ return self._show_history()
+ elif cmd == "help":
+ return self._show_help()
+ else:
+ return f"Unknown command: {cmd}. Use '%ballista help' for available commands."
+
+ @line_magic
+ def register(self, line: str) -> Optional[str]:
+ """Register a new table"""
+ if not line:
+ return "You should provide file extension and table name to register"
+ elif self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ else:
+ args = line.strip().split()
+ file_type = args[0]
+ if len(args) < 2:
+ return f"You should provide table name for this .{file_type} file"
+ table_name = args[1]
+ if len(args) < 3:
+ return f"You should provide path to your file"
+ file_name = args[2]
+
+ if file_type == "parquet":
+ self._ctx.register_parquet(table_name, file_name)
+ elif file_type == "csv":
+ self._ctx.register_csv(table_name, file_name)
+ else:
+ raise NotImplemented("Currently not supporting the inserted file format")
+
+ @line_cell_magic
+ def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]:
+ """
+ Execute a SQL query (both line and cell magic).
+
+ Two cases possible: with cell or without cell
+
+ Examples:
+ 1. Without a cell (line_magic)
+ %sql SELECT * FROM test_table
+ 2. With a cell (cell_magic)
+ %%sql --no-display
+ SELECT
+ id,
+ bool_col,
+ tinyint_col
+ FROM test_data_v1
+ WHERE id > 2
+ ORDER BY id
+ LIMIT 5
+ """
+ if not cell:
+ return self._execute_sql(line.strip()) if line.strip() else None
+ else:
+ args = line.strip().split()
+ display_results = True
+ limit = 50
+ var_name = None
+
+ i = 0
+ while i < len(args):
+ if args[i] == "--no-display":
+ display_results = False
+ elif args[i] == "--limit" and i + 1 < len(args):
+ try:
+ limit = int(args[i + 1])
+ i += 1
+ except ValueError:
+ pass
+ elif not args[i].startswith("--"):
+ var_name = args[i]
+ i += 1
+
+ query = cell.strip()
+ if not query:
+ return None
+
+ result = self._execute_sql(query, display_results=display_results, limit=limit)
+
+ # Store in user namespace if variable name provided
+ if var_name and self.shell is not None:
+ self.shell.user_ns[var_name] = result
+
+ return result
+
+ def _connect(self, address: str) -> Optional[str]:
+ """Connect to a Ballista cluster."""
+ if not address:
+ return "Usage: %ballista connect df://host:port"
+
+ # Normalize address
+ if not address.startswith("df://"):
+ address = f"df://{address}"
+
+ try:
+ self._ctx = BallistaSessionContext(address)
+ self._address = address
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Connected to Ballista cluster at {address}"))
+ else:
+ return f"✓ Connected to Ballista cluster at {address}"
+ except Exception as e:
+ self._ctx = None
+ self._address = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✗ Failed to connect to {address}: {e}"))
+ else:
+ return f"✗ Failed to connect to {address}: {e}"
+
+ def _disconnect(self) -> Optional[str]:
+ """Disconnect from the Ballista cluster."""
+ if not self.is_connected:
+ return "Not connected to any cluster."
+
+ address = self._address
+ self._ctx = None
+ self._address = None
+ self._last_result = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Disconnected from {address}"))
+ else:
+ return f"✓ Disconnected from {address}"
+
+ def _status(self) -> Optional[str]:
+ """Show connection status."""
+ if not self.is_connected:
+ return "Status: Not connected\n\nUse '%ballista connect df://host:port' to connect."
+
+ status_lines = [
+ f"Status: Connected",
+ f"Address: {self._address}",
+ f"Session ID: {self._ctx.session_id}",
+ f"Queries executed: {len(self._query_history)}",
+ ]
+
+ if self._last_result is not None:
+ status_lines.append("Last result: Available (access via '_' or '_last_result')")
+
+ def _format_html_status_output(line: str) -> str:
+ name, value = line.split(":", 1)
+ return f"<p><strong>{name}:</strong> {value.strip()}</p>"
+
+ html = "".join(_format_html_status_output(line) for line in status_lines)
+ if IPYTHON_AVAILABLE:
+ display(HTML(html))
+ else:
+ print("\n".join(status_lines))
+
+ def _tables(self) -> Optional[str]:
+ """List all registered tables."""
+ try:
+ # Get table names from the catalog and their respective schemas
+ tables = self.ctx.get_tables()
+ 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())
+ # Build a nice table display (HTML-formatted if applicable)
+ lines = [
+ {"content": f"Total: {table_count} table(s) in {schema_count} schema(s)", "is_info": True},
+ {"content": "Registered tables:", "is_info": True},
+ *[{"content": f"Schema: {schema_name}. Tables: {", ".join(list(table_names))}", "is_info": False}
+ for schema_name, table_names in tables.items()]
+ ]
+
+ def _format_html_tables_output(line: str, is_info: bool = False) -> str:
+ if is_info:
+ return f"<p><strong>{line}<strong></p>"
+ else:
+ return f"<p><pre><i>{line}</i></pre></p>"
+
+ if IPYTHON_AVAILABLE:
+ display(
+ HTML("".join(_format_html_tables_output(val["content"], val["is_info"]) for val in lines))
+ )
+ else:
+ return "".join(val["content"] for val in lines)
+ except Exception as e:
+ return f"Error listing tables: {e}"
+
+ def _schema(self, table_name: str) -> Optional[str]:
+ """Show schema for a table."""
+ if not table_name:
+ return "Usage: %ballista schema <table_name>"
+
+ try:
+ # Query the table with LIMIT 0 to get schema without data
+ df = self.ctx.sql(f"SELECT * FROM {table_name} LIMIT 0")
+ schema = df.schema()
+
+ lines = [f"Schema for '{table_name}':", "-" * 50]
+ for field in schema:
+ nullable = "NULL" if field.nullable else "NOT NULL"
+ lines.append(f" {field.name:20} {str(field.type):15} {nullable}")
+ lines.append("-" * 50)
+ lines.append(f"Total: {len(schema)} column(s)")
+
+ return "\n".join(lines)
+ except Exception as e:
+ return f"Error getting schema for '{table_name}': {e}"
+
+ def _execute_sql(
+ self,
+ query: str,
+ display_results: bool = True,
+ limit: int = 50,
+ ) -> Optional[DistributedDataFrame]:
+ """Execute a SQL query and return the result."""
+ start_time = time.time()
+
+ try:
+ result = self.ctx.sql(query)
+ elapsed = time.time() - start_time
+
+ # Store result
+ self._last_result = result
+ if self.shell is not None and hasattr(self.shell, 'user_ns'):
+ self.shell.user_ns["_last_result"] = result
+
+ # Record in history
+ self._query_history.append(
+ {
+ "query": query,
+ "elapsed_seconds": elapsed,
+ "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
+ }
+ )
+
+ # Display if requested and in notebook environment
+ return result
+
+ except Exception as e:
+ error_msg = f"Query failed: {e}"
+ if IPYTHON_AVAILABLE:
+ try:
+ display(HTML(f'<div style="color: red; font-weight: bold;">{error_msg}</div>'))
+ except Exception:
+ print(error_msg)
+ else:
+ print(error_msg)
+ return None
+
+ def _show_history(self) -> Optional[str]:
+ """Show query history."""
+ if not self._query_history:
+ return "No queries executed yet."
+
+ lines = ["Query History:", "-" * 60]
+ for i, entry in enumerate(self._query_history[-10:], 1): # Last 10 queries
+ query_preview = entry["query"][:50] + "..." if len(entry["query"]) > 50 else entry["query"]
+ query_preview = query_preview.replace("\n", " ")
+ lines.append(f"{i}. [{entry['timestamp']}] ({entry['elapsed_seconds']:.2f}s)")
+ lines.append(f" {query_preview}")
+ lines.append("-" * 60)
+
+ for row in lines:
+ print(row)
+
+ def _show_help(self) -> Optional[str]:
+ """Show help for Ballista magic commands."""
+ help_info = """
+Ballista Jupyter Magic Commands
+================================
+
+Connection:
+ %ballista connect <url> - Connect to Ballista cluster
+ %ballista disconnect - Disconnect from cluster
+ %ballista status - Show connection status
+
+Exploration:
+ %ballista tables - List registered tables
+ %ballista schema <table> - Show table schema
+
+Table-register:
+ %register [format] [schema.table_name] [file_path] - Register a new table in the current Ballista Context
+
+Query:
+ %sql <query> - Execute single-line SQL query
+
+ %%sql [options] [var] - Execute multi-line SQL query
+ Options:
+ --no-display - Don't display results
+ --limit N - Limit displayed rows (default: 50)
+ var - Store result in variable
+
+History:
+ %ballista history - Show recent query history
+
+Examples:
+ %ballista connect df://localhost:50050
+ %ballista tables
+ %ballista schema orders
+
+ %sql SELECT COUNT(*) FROM orders
+
+ %%sql my_result
+ SELECT customer_id, SUM(amount) as total
+ FROM orders
+ GROUP BY customer_id
+ ORDER BY total DESC
+ LIMIT 10
+ """
+ for row in help_info.split("\n"):
+ print(row)
Review Comment:
```suggestion
print(row)
return help_info
```
##########
python/python/ballista/extension.py:
##########
@@ -229,9 +231,319 @@ def write_parquet(
df = self._to_internal_df()
df.write_parquet(str(path), compression.value, compression_level)
+ def explain_visual(self, analyze: bool = False) -> "ExecutionPlanVisualization":
+ """
+ Generate a visual representation of the execution plan.
+
+ This method creates an SVG visualization of the query execution plan,
+ which can be displayed directly in Jupyter notebooks.
+
+ Args:
+ analyze: If True, includes runtime statistics from actual execution.
+
+ Returns:
+ ExecutionPlanVisualization: An object that renders as SVG in Jupyter.
+
+ Example:
+ >>> df = ctx.sql("SELECT * FROM orders WHERE amount > 100")
+ >>> df.explain_visual() # Displays SVG in notebook
+ >>> viz = df.explain_visual(analyze=True)
+ >>> viz.save("plan.svg") # Save to file
+ """
+ # Get the execution plan as a string representation
+ # Note: explain() prints but doesn't return a string, so we use logical_plan()
+ try:
+ plan = self.logical_plan()
+ plan_str = plan.display_indent()
+ except Exception:
+ # Fallback if logical_plan() fails
+ plan_str = "Unable to retrieve execution plan"
+ return ExecutionPlanVisualization(plan_str, analyze=analyze)
+
+ def collect_with_progress(
+ self,
+ callback: Optional[callable] = None,
+ poll_interval: float = 0.5,
+ ):
+ """
+ Collect results with progress indication.
+
+ For long-running queries, this method provides progress updates
+ through a callback function or displays a progress bar in Jupyter.
+
+ Args:
+ callback: Optional function to call with progress updates.
+ Signature: callback(status: str, progress: float)
+ poll_interval: How often to check progress (seconds).
+
+ Returns:
+ The collected result batches.
+
+ Example:
+ >>> def my_callback(status, progress):
+ ... print(f"{status}: {progress:.1%}")
+ >>> batches = df.collect_with_progress(callback=my_callback)
+ """
+ import threading
+ import time
+
+ result = [None]
+ error = [None]
+ done = threading.Event()
+
+ def execute():
+ try:
+ result[0] = self.collect()
+ except Exception as e:
+ error[0] = e
+ finally:
+ done.set()
+
+ thread = threading.Thread(target=execute)
+ thread.start()
+
+ # Check if we're in a Jupyter environment
+ try:
+ from IPython.display import display, clear_output
+ from IPython import get_ipython
+
+ in_jupyter = get_ipython() is not None
+ except (ImportError, AttributeError):
+ in_jupyter = False
+
+ start_time = time.time()
+
+ if in_jupyter and callback is None:
+ # Display a simple progress indicator
+ try:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ clear_output(wait=True)
+ print(f"⏳ Query executing... ({elapsed:.1f}s elapsed)")
+
+ clear_output(wait=True)
+ elapsed = time.time() - start_time
+ print(f"✓ Query completed in {elapsed:.1f}s")
+ except Exception:
+ pass # Ignore display errors
+ elif callback is not None:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ callback(f"Executing ({elapsed:.1f}s)", -1.0) # -1 means indeterminate
+
+ elapsed = time.time() - start_time
+ callback(f"Completed in {elapsed:.1f}s", 1.0)
+ else:
+ done.wait()
+
+ thread.join()
+
+ if error[0] is not None:
+ raise error[0]
+
+ return result[0]
+
+
+class ExecutionPlanVisualization:
+ """
+ A wrapper for execution plan visualizations that can render as SVG in Jupyter.
+
+ This class takes the text representation of an execution plan and converts
+ it to a Graphviz DOT format, which is then rendered as SVG.
+ """
+
+ def __init__(self, plan_str: str, analyze: bool = False):
+ self.plan_str = plan_str
+ self.analyze = analyze
+ self._svg_cache: Optional[str] = None
+
+ def _parse_plan_to_dot(self) -> str:
+ """Convert the plan string to DOT format for Graphviz."""
+ lines = self.plan_str.strip().split("\n")
+
+ dot_lines = [
+ "digraph ExecutionPlan {",
+ ' rankdir=TB;',
+ ' node [shape=box, style="rounded,filled", fontname="Helvetica"];',
+ ' edge [fontname="Helvetica"];',
+ "",
+ ]
+
+ nodes = []
+ edges = []
+ node_id = 0
+ stack = [] # (indent_level, node_id)
+
+ for line in lines:
+ if not line.strip():
+ continue
+
+ # Calculate indent level
+ indent = len(line) - len(line.lstrip())
+ content = line.strip()
+
+ # Skip non-plan lines
+ if content.startswith("physical_plan") or content.startswith("logical_plan"):
+ continue
+
+ # Create a node for this plan element
+ current_id = node_id
+ node_id += 1
+
+ # Determine node color based on operation type
+ color = "#E3F2FD" # Default light blue
+ if "Scan" in content or "TableScan" in content:
+ color = "#E8F5E9" # Light green for scans
+ elif "Filter" in content:
+ color = "#FFF3E0" # Light orange for filters
+ elif "Aggregate" in content or "HashAggregate" in content:
+ color = "#F3E5F5" # Light purple for aggregations
+ elif "Join" in content:
+ color = "#FFEBEE" # Light red for joins
+ elif "Sort" in content:
+ color = "#E0F7FA" # Light cyan for sorts
+ elif "Projection" in content:
+ color = "#FFF8E1" # Light amber for projections
+
+ # Escape special characters for DOT format
+ label = content.replace('"', '\\"').replace("\n", "\\n")
+ if len(label) > 60:
+ # Wrap long labels
+ label = label[:57] + "..."
+
+ nodes.append(f' node{current_id} [label="{label}", fillcolor="{color}"];')
+
+ # Connect to parent based on indentation
+ while stack and stack[-1][0] >= indent:
+ stack.pop()
+
+ if stack:
+ parent_id = stack[-1][1]
+ edges.append(f" node{parent_id} -> node{current_id};")
+
+ stack.append((indent, current_id))
+
+ dot_lines.extend(nodes)
+ dot_lines.append("")
+ dot_lines.extend(edges)
+ dot_lines.append("}")
+
+ return "\n".join(dot_lines)
+
+ def to_dot(self) -> str:
+ """Get the DOT representation of the execution plan."""
+ return self._parse_plan_to_dot()
+
+ def to_svg(self) -> str:
+ """
+ Convert the execution plan to SVG format.
+
+ Requires graphviz to be installed. If graphviz is not available,
+ returns a simple HTML representation instead.
+ """
+ if self._svg_cache is not None:
+ return self._svg_cache
+
+ dot_source = self._parse_plan_to_dot()
+
+ try:
+ import subprocess
+
+ # Try to use graphviz's dot command
+ process = subprocess.run(
+ ["dot", "-Tsvg"],
+ input=dot_source.encode(),
+ capture_output=True,
+ timeout=30,
+ )
+
+ if process.returncode == 0:
+ self._svg_cache = process.stdout.decode()
+ return self._svg_cache
+ except (subprocess.SubprocessError, FileNotFoundError, subprocess.TimeoutExpired):
+ pass
+
+ # Fallback: return a pre-formatted HTML representation
+ escaped_plan = (
+ self.plan_str.replace("&", "&")
+ .replace("<", "<")
+ .replace(">", ">")
+ )
+ self._svg_cache = f"""
+ <div style="font-family: monospace; background: #f5f5f5; padding: 10px;
+ border-radius: 5px; overflow-x: auto;">
+ <div style="color: #666; margin-bottom: 5px;">
+ Execution Plan {'(with statistics)' if self.analyze else ''}
+ <br><small>Install graphviz for visual diagram: brew install graphviz</small>
+ </div>
+ <pre style="margin: 0;">{escaped_plan}</pre>
+ </div>
+ """
+ return self._svg_cache
+
+ def save(self, path: str) -> None:
+ """Save the visualization to a file (SVG or DOT format)."""
+ if path.endswith(".dot"):
+ content = self.to_dot()
+ else:
+ content = self.to_svg()
+
+ with open(path, "w") as f:
+ f.write(content)
+
+ def _repr_html_(self) -> str:
+ """HTML representation for Jupyter notebooks."""
+ return self.to_svg()
+
+ def _repr_svg_(self) -> str:
+ """SVG representation for Jupyter notebooks."""
+ svg = self.to_svg()
+ # Only return if it's actual SVG content
+ if svg.strip().startswith("<svg") or svg.strip().startswith("<?xml"):
+ return svg
+ return ""
+
+ def __repr__(self) -> str:
+ """String representation."""
+ return f"ExecutionPlanVisualization(analyze={self.analyze})\n{self.plan_str}"
+
class BallistaSessionContext(SessionContext, metaclass=RedefiningSessionContextMeta):
+ """
+ A session context for connecting to and querying a Ballista cluster.
+
+ This class extends DataFusion's SessionContext to work with distributed
+ Ballista clusters, automatically routing query execution to the cluster
+ while maintaining API compatibility with local DataFusion usage.
+
+ Example:
+ >>> from ballista import BallistaSessionContext
+ >>> ctx = BallistaSessionContext("df://localhost:50050")
+ >>> df = ctx.sql("SELECT * FROM my_table LIMIT 10")
+ >>> df.show()
+
+ For Jupyter notebook users:
+ >>> %load_ext ballista.jupyter
+ >>> %ballista connect df://localhost:50050
+ >>> %sql SELECT * FROM my_table
+ """
+
def __init__(self, address: str, config=None, runtime=None):
super().__init__(config, runtime)
self.address = address
self.session_id = self.session_id()
Review Comment:
This overwrites/shadows the method with a field. Any attempt later to use the method will fail with something like `'str' object is not callable-style error.`
##########
python/README.md:
##########
@@ -79,6 +79,94 @@ pyarrow_batches = df.collect()
Check [DataFusion python](https://datafusion.apache.org/python/) provides more examples and manuals.
+## Jupyter Notebook Support
+
+PyBallista provides first-class Jupyter notebook support with SQL magic commands and rich HTML rendering.
+
Review Comment:
```suggestion
Install Jupyter extras first:
```bash
pip install "ballista[jupyter]"
```
```
##########
python/python/ballista/extension.py:
##########
@@ -229,9 +231,319 @@ def write_parquet(
df = self._to_internal_df()
df.write_parquet(str(path), compression.value, compression_level)
+ def explain_visual(self, analyze: bool = False) -> "ExecutionPlanVisualization":
+ """
+ Generate a visual representation of the execution plan.
+
+ This method creates an SVG visualization of the query execution plan,
+ which can be displayed directly in Jupyter notebooks.
+
+ Args:
+ analyze: If True, includes runtime statistics from actual execution.
+
+ Returns:
+ ExecutionPlanVisualization: An object that renders as SVG in Jupyter.
+
+ Example:
+ >>> df = ctx.sql("SELECT * FROM orders WHERE amount > 100")
+ >>> df.explain_visual() # Displays SVG in notebook
+ >>> viz = df.explain_visual(analyze=True)
+ >>> viz.save("plan.svg") # Save to file
+ """
+ # Get the execution plan as a string representation
+ # Note: explain() prints but doesn't return a string, so we use logical_plan()
+ try:
+ plan = self.logical_plan()
+ plan_str = plan.display_indent()
+ except Exception:
+ # Fallback if logical_plan() fails
+ plan_str = "Unable to retrieve execution plan"
+ return ExecutionPlanVisualization(plan_str, analyze=analyze)
+
+ def collect_with_progress(
+ self,
+ callback: Optional[callable] = None,
+ poll_interval: float = 0.5,
+ ):
+ """
+ Collect results with progress indication.
+
+ For long-running queries, this method provides progress updates
+ through a callback function or displays a progress bar in Jupyter.
+
+ Args:
+ callback: Optional function to call with progress updates.
+ Signature: callback(status: str, progress: float)
+ poll_interval: How often to check progress (seconds).
+
+ Returns:
+ The collected result batches.
+
+ Example:
+ >>> def my_callback(status, progress):
+ ... print(f"{status}: {progress:.1%}")
+ >>> batches = df.collect_with_progress(callback=my_callback)
+ """
+ import threading
+ import time
+
+ result = [None]
+ error = [None]
+ done = threading.Event()
+
+ def execute():
+ try:
+ result[0] = self.collect()
+ except Exception as e:
+ error[0] = e
+ finally:
+ done.set()
+
+ thread = threading.Thread(target=execute)
+ thread.start()
+
+ # Check if we're in a Jupyter environment
+ try:
+ from IPython.display import display, clear_output
+ from IPython import get_ipython
+
+ in_jupyter = get_ipython() is not None
+ except (ImportError, AttributeError):
+ in_jupyter = False
+
+ start_time = time.time()
+
+ if in_jupyter and callback is None:
+ # Display a simple progress indicator
+ try:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ clear_output(wait=True)
+ print(f"⏳ Query executing... ({elapsed:.1f}s elapsed)")
+
+ clear_output(wait=True)
+ elapsed = time.time() - start_time
+ print(f"✓ Query completed in {elapsed:.1f}s")
+ except Exception:
+ pass # Ignore display errors
+ elif callback is not None:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ callback(f"Executing ({elapsed:.1f}s)", -1.0) # -1 means indeterminate
+
+ elapsed = time.time() - start_time
+ callback(f"Completed in {elapsed:.1f}s", 1.0)
+ else:
+ done.wait()
+
+ thread.join()
+
+ if error[0] is not None:
+ raise error[0]
+
+ return result[0]
+
+
+class ExecutionPlanVisualization:
+ """
+ A wrapper for execution plan visualizations that can render as SVG in Jupyter.
+
+ This class takes the text representation of an execution plan and converts
+ it to a Graphviz DOT format, which is then rendered as SVG.
+ """
+
+ def __init__(self, plan_str: str, analyze: bool = False):
+ self.plan_str = plan_str
+ self.analyze = analyze
+ self._svg_cache: Optional[str] = None
+
+ def _parse_plan_to_dot(self) -> str:
+ """Convert the plan string to DOT format for Graphviz."""
+ lines = self.plan_str.strip().split("\n")
+
+ dot_lines = [
+ "digraph ExecutionPlan {",
+ ' rankdir=TB;',
+ ' node [shape=box, style="rounded,filled", fontname="Helvetica"];',
+ ' edge [fontname="Helvetica"];',
+ "",
+ ]
+
+ nodes = []
+ edges = []
+ node_id = 0
+ stack = [] # (indent_level, node_id)
+
+ for line in lines:
+ if not line.strip():
+ continue
+
+ # Calculate indent level
+ indent = len(line) - len(line.lstrip())
+ content = line.strip()
+
+ # Skip non-plan lines
+ if content.startswith("physical_plan") or content.startswith("logical_plan"):
+ continue
+
+ # Create a node for this plan element
+ current_id = node_id
+ node_id += 1
+
+ # Determine node color based on operation type
+ color = "#E3F2FD" # Default light blue
+ if "Scan" in content or "TableScan" in content:
+ color = "#E8F5E9" # Light green for scans
+ elif "Filter" in content:
+ color = "#FFF3E0" # Light orange for filters
+ elif "Aggregate" in content or "HashAggregate" in content:
+ color = "#F3E5F5" # Light purple for aggregations
+ elif "Join" in content:
+ color = "#FFEBEE" # Light red for joins
+ elif "Sort" in content:
+ color = "#E0F7FA" # Light cyan for sorts
+ elif "Projection" in content:
+ color = "#FFF8E1" # Light amber for projections
+
+ # Escape special characters for DOT format
+ label = content.replace('"', '\\"').replace("\n", "\\n")
+ if len(label) > 60:
+ # Wrap long labels
+ label = label[:57] + "..."
+
+ nodes.append(f' node{current_id} [label="{label}", fillcolor="{color}"];')
+
+ # Connect to parent based on indentation
+ while stack and stack[-1][0] >= indent:
+ stack.pop()
+
+ if stack:
+ parent_id = stack[-1][1]
+ edges.append(f" node{parent_id} -> node{current_id};")
+
+ stack.append((indent, current_id))
+
+ dot_lines.extend(nodes)
+ dot_lines.append("")
+ dot_lines.extend(edges)
+ dot_lines.append("}")
+
+ return "\n".join(dot_lines)
+
+ def to_dot(self) -> str:
+ """Get the DOT representation of the execution plan."""
+ return self._parse_plan_to_dot()
+
+ def to_svg(self) -> str:
+ """
+ Convert the execution plan to SVG format.
+
+ Requires graphviz to be installed. If graphviz is not available,
+ returns a simple HTML representation instead.
+ """
+ if self._svg_cache is not None:
+ return self._svg_cache
+
+ dot_source = self._parse_plan_to_dot()
+
+ try:
+ import subprocess
+
+ # Try to use graphviz's dot command
+ process = subprocess.run(
+ ["dot", "-Tsvg"],
+ input=dot_source.encode(),
+ capture_output=True,
+ timeout=30,
+ )
+
+ if process.returncode == 0:
+ self._svg_cache = process.stdout.decode()
+ return self._svg_cache
+ except (subprocess.SubprocessError, FileNotFoundError, subprocess.TimeoutExpired):
+ pass
+
+ # Fallback: return a pre-formatted HTML representation
+ escaped_plan = (
+ self.plan_str.replace("&", "&")
+ .replace("<", "<")
+ .replace(">", ">")
+ )
+ self._svg_cache = f"""
+ <div style="font-family: monospace; background: #f5f5f5; padding: 10px;
+ border-radius: 5px; overflow-x: auto;">
+ <div style="color: #666; margin-bottom: 5px;">
+ Execution Plan {'(with statistics)' if self.analyze else ''}
+ <br><small>Install graphviz for visual diagram: brew install graphviz</small>
+ </div>
+ <pre style="margin: 0;">{escaped_plan}</pre>
+ </div>
+ """
+ return self._svg_cache
+
+ def save(self, path: str) -> None:
+ """Save the visualization to a file (SVG or DOT format)."""
+ if path.endswith(".dot"):
+ content = self.to_dot()
+ else:
+ content = self.to_svg()
+
+ with open(path, "w") as f:
+ f.write(content)
+
+ def _repr_html_(self) -> str:
+ """HTML representation for Jupyter notebooks."""
+ return self.to_svg()
+
+ def _repr_svg_(self) -> str:
+ """SVG representation for Jupyter notebooks."""
+ svg = self.to_svg()
+ # Only return if it's actual SVG content
+ if svg.strip().startswith("<svg") or svg.strip().startswith("<?xml"):
+ return svg
+ return ""
+
+ def __repr__(self) -> str:
+ """String representation."""
+ return f"ExecutionPlanVisualization(analyze={self.analyze})\n{self.plan_str}"
+
class BallistaSessionContext(SessionContext, metaclass=RedefiningSessionContextMeta):
+ """
+ A session context for connecting to and querying a Ballista cluster.
+
+ This class extends DataFusion's SessionContext to work with distributed
+ Ballista clusters, automatically routing query execution to the cluster
+ while maintaining API compatibility with local DataFusion usage.
+
+ Example:
+ >>> from ballista import BallistaSessionContext
+ >>> ctx = BallistaSessionContext("df://localhost:50050")
+ >>> df = ctx.sql("SELECT * FROM my_table LIMIT 10")
+ >>> df.show()
+
+ For Jupyter notebook users:
+ >>> %load_ext ballista.jupyter
+ >>> %ballista connect df://localhost:50050
+ >>> %sql SELECT * FROM my_table
+ """
+
def __init__(self, address: str, config=None, runtime=None):
super().__init__(config, runtime)
self.address = address
self.session_id = self.session_id()
+
+ def get_tables(self) -> Optional[dict[str, str]]:
+ """Get tables and their respective schemas (in terms of database schema)."""
+ try:
+ catalog = self.catalog()
+ schema_names = list(catalog.schema_names())
+ if schema_names is not None:
Review Comment:
A list is never None.
```suggestion
if schema_names:
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feature/jupyter notebook support [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513#issuecomment-4149868324
Converting to draft, going to fix it
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feature/jupyter notebook support [datafusion-ballista]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513#issuecomment-4149866257
hey @sandugood there are few tests failing here, if you wish to progress with this PR please have a look
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feature/jupyter notebook support [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513#issuecomment-4074367515
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feature/jupyter notebook support [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on code in PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513#discussion_r2946170716
##########
python/python/ballista/jupyter.py:
##########
@@ -0,0 +1,493 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Ballista Jupyter Magic Commands.
+
+This module provides IPython magic commands for interacting with Ballista
+clusters directly from Jupyter notebooks.
+
+Usage:
+ %load_ext ballista.jupyter
+
+ # Connect to a Ballista cluster
+ %ballista connect df://localhost:50050
+
+ # Check connection status
+ %ballista status
+
+ # List registered tables
+ %ballista tables
+
+ # Show schema for a table
+ %ballista schema my_table
+
+ # Execute a simple SQL query (line magic)
+ %sql SELECT COUNT(*) FROM my_table
+
+ # Execute a complex SQL query (cell magic)
+ %%sql
+ SELECT
+ customer_id,
+ SUM(amount) as total
+ FROM orders
+ GROUP BY customer_id
+ ORDER BY total DESC
+ LIMIT 10
+"""
+
+from typing import Optional, List, Dict, Any
+import time
+import threading
+
+try:
+ from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic
+ from IPython.display import display, HTML
+
+ IPYTHON_AVAILABLE = True
+except ImportError:
+ IPYTHON_AVAILABLE = False
+
+ # Provide stub classes for when IPython is not available
+ class Magics:
+ def __init__(self, shell=None):
+ self.shell = shell
+
+ def magics_class(cls):
+ return cls
+
+ def line_magic(name_or_func):
+ """Stub line_magic decorator for when IPython is not available."""
+ # Handle both @line_magic and @line_magic("name") usage
+ if callable(name_or_func):
+ # Used as @line_magic without arguments
+ return name_or_func
+ else:
+ # Used as @line_magic("name") with arguments
+ def decorator(func):
+ return func
+ return decorator
+
+ def cell_magic(name_or_func):
+ """Stub cell_magic decorator for when IPython is not available."""
+ # Handle both @cell_magic and @cell_magic("name") usage
+ if callable(name_or_func):
+ return name_or_func
+ else:
+ def decorator(func):
+ return func
+ return decorator
+
+
+from .extension import BallistaSessionContext, DistributedDataFrame
+
+
+class BallistaConnectionError(Exception):
+ """Raised when not connected to a Ballista cluster."""
+
+ pass
+
+
+@magics_class
+class BallistaMagics(Magics):
+ """
+ IPython magic commands for Ballista.
+
+ Provides convenient commands for connecting to Ballista clusters,
+ executing SQL queries, and exploring table schemas.
+ """
+
+ def __init__(self, shell=None):
+ super().__init__(shell)
+ self._ctx: Optional[BallistaSessionContext] = None
+ self._address: Optional[str] = None
+ self._last_result: Optional[DistributedDataFrame] = None
+ self._query_history: List[Dict[str, Any]] = []
+
+ @property
+ def ctx(self) -> BallistaSessionContext:
+ """Get the current context, raising an error if not connected."""
+ if self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ return self._ctx
+
+ @property
+ def is_connected(self) -> bool:
+ """Check if connected to a Ballista cluster."""
+ return self._ctx is not None
+
+ @line_magic
+ def ballista(self, line: str) -> Optional[str]:
+ """
+ Ballista management commands.
+
+ Usage:
+ %ballista connect df://localhost:50050 - Connect to cluster
+ %ballista status - Show connection status
+ %ballista tables - List registered tables
+ %ballista schema <table> - Show table schema
+ %ballista disconnect - Disconnect from cluster
+ %ballista history - Show query history
+
+ Examples:
+ %ballista connect df://localhost:50050
+ %ballista tables
+ %ballista schema orders
+ """
+ parts = line.strip().split(maxsplit=1)
+ if not parts:
+ return self._show_help()
+
+ cmd = parts[0].lower()
+ args = parts[1] if len(parts) > 1 else ""
+
+ if cmd == "connect":
+ return self._connect(args)
+ elif cmd == "status":
+ return self._status()
+ elif cmd == "tables":
+ return self._tables()
+ elif cmd == "schema":
+ return self._schema(args)
+ elif cmd == "disconnect":
+ return self._disconnect()
+ elif cmd == "history":
+ return self._show_history()
+ elif cmd == "help":
+ return self._show_help()
+ else:
+ return f"Unknown command: {cmd}. Use '%ballista help' for available commands."
+
+ @line_magic
+ def register(self, line: str) -> Optional[str]:
+ """Register a new table"""
+ if not line:
+ return "You should provide file extension and table name to register"
+ elif self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ else:
+ args = line.strip().split()
+ file_type = args[0]
+ if len(args) < 2:
+ return f"You should provide table name for this .{file_type} file"
+ table_name = args[1]
+ if len(args) < 3:
+ return f"You should provide path to your file"
+ file_name = args[2]
+
+ if file_type == "parquet":
+ self._ctx.register_parquet(table_name, file_name)
+ elif file_type == "csv":
+ self._ctx.register_csv(table_name, file_name)
+ else:
+ raise NotImplemented("Currently not supporting the inserted file format")
+
+ @line_cell_magic
+ def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]:
+ """
+ Execute a SQL query (both line and cell magic).
+
+ Two cases possible: with cell or without cell
+
+ Examples:
+ 1. Without a cell (line_magic)
+ %sql SELECT * FROM test_table
+ 2. With a cell (cell_magic)
+ %%sql --no-display
+ SELECT
+ id,
+ bool_col,
+ tinyint_col
+ FROM test_data_v1
+ WHERE id > 2
+ ORDER BY id
+ LIMIT 5
+ """
+ if not cell:
+ return self._execute_sql(line.strip()) if line.strip() else None
+ else:
+ args = line.strip().split()
+ display_results = True
+ limit = 50
+ var_name = None
+
+ i = 0
+ while i < len(args):
+ if args[i] == "--no-display":
+ display_results = False
+ elif args[i] == "--limit" and i + 1 < len(args):
+ try:
+ limit = int(args[i + 1])
+ i += 1
+ except ValueError:
+ pass
+ elif not args[i].startswith("--"):
+ var_name = args[i]
+ i += 1
+
+ query = cell.strip()
+ if not query:
+ return None
+
+ result = self._execute_sql(query, display_results=display_results, limit=limit)
+
+ # Store in user namespace if variable name provided
+ if var_name and self.shell is not None:
+ self.shell.user_ns[var_name] = result
+
+ return result
+
+ def _connect(self, address: str) -> Optional[str]:
+ """Connect to a Ballista cluster."""
+ if not address:
+ return "Usage: %ballista connect df://host:port"
+
+ # Normalize address
+ if not address.startswith("df://"):
+ address = f"df://{address}"
+
+ try:
+ self._ctx = BallistaSessionContext(address)
+ self._address = address
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Connected to Ballista cluster at {address}"))
+ else:
+ return f"✓ Connected to Ballista cluster at {address}"
+ except Exception as e:
+ self._ctx = None
+ self._address = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✗ Failed to connect to {address}: {e}"))
+ else:
+ return f"✗ Failed to connect to {address}: {e}"
+
+ def _disconnect(self) -> Optional[str]:
+ """Disconnect from the Ballista cluster."""
+ if not self.is_connected:
+ return "Not connected to any cluster."
+
+ address = self._address
+ self._ctx = None
+ self._address = None
+ self._last_result = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Disconnected from {address}"))
+ else:
+ return f"✓ Disconnected from {address}"
+
+ def _status(self) -> Optional[str]:
+ """Show connection status."""
+ if not self.is_connected:
+ return "Status: Not connected\n\nUse '%ballista connect df://host:port' to connect."
+
+ status_lines = [
+ f"Status: Connected",
+ f"Address: {self._address}",
+ f"Session ID: {self._ctx.session_id}",
+ f"Queries executed: {len(self._query_history)}",
+ ]
+
+ if self._last_result is not None:
+ status_lines.append("Last result: Available (access via '_' or '_last_result')")
+
+ def _format_html_status_output(line: str) -> str:
+ name, value = line.split(":", 1)
+ return f"<p><strong>{name}:</strong> {value.strip()}</p>"
+
+ html = "".join(_format_html_status_output(line) for line in status_lines)
+ if IPYTHON_AVAILABLE:
+ display(HTML(html))
+ else:
+ print("\n".join(status_lines))
+
+ def _tables(self) -> Optional[str]:
+ """List all registered tables."""
+ try:
+ # Get table names from the catalog and their respective schemas
+ tables = self.ctx.get_tables()
+ 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())
Review Comment:
Resolved.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feature/jupyter notebook support [datafusion-ballista]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513#issuecomment-4149885691
thanks
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feature/jupyter notebook support [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on code in PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513#discussion_r2946194387
##########
python/python/ballista/extension.py:
##########
@@ -229,9 +231,319 @@ def write_parquet(
df = self._to_internal_df()
df.write_parquet(str(path), compression.value, compression_level)
+ def explain_visual(self, analyze: bool = False) -> "ExecutionPlanVisualization":
+ """
+ Generate a visual representation of the execution plan.
+
+ This method creates an SVG visualization of the query execution plan,
+ which can be displayed directly in Jupyter notebooks.
+
+ Args:
+ analyze: If True, includes runtime statistics from actual execution.
+
+ Returns:
+ ExecutionPlanVisualization: An object that renders as SVG in Jupyter.
+
+ Example:
+ >>> df = ctx.sql("SELECT * FROM orders WHERE amount > 100")
+ >>> df.explain_visual() # Displays SVG in notebook
+ >>> viz = df.explain_visual(analyze=True)
+ >>> viz.save("plan.svg") # Save to file
+ """
+ # Get the execution plan as a string representation
+ # Note: explain() prints but doesn't return a string, so we use logical_plan()
+ try:
+ plan = self.logical_plan()
+ plan_str = plan.display_indent()
+ except Exception:
+ # Fallback if logical_plan() fails
+ plan_str = "Unable to retrieve execution plan"
+ return ExecutionPlanVisualization(plan_str, analyze=analyze)
+
+ def collect_with_progress(
+ self,
+ callback: Optional[callable] = None,
+ poll_interval: float = 0.5,
+ ):
+ """
+ Collect results with progress indication.
+
+ For long-running queries, this method provides progress updates
+ through a callback function or displays a progress bar in Jupyter.
+
+ Args:
+ callback: Optional function to call with progress updates.
+ Signature: callback(status: str, progress: float)
+ poll_interval: How often to check progress (seconds).
+
+ Returns:
+ The collected result batches.
+
+ Example:
+ >>> def my_callback(status, progress):
+ ... print(f"{status}: {progress:.1%}")
+ >>> batches = df.collect_with_progress(callback=my_callback)
+ """
+ import threading
+ import time
+
+ result = [None]
+ error = [None]
+ done = threading.Event()
+
+ def execute():
+ try:
+ result[0] = self.collect()
+ except Exception as e:
+ error[0] = e
+ finally:
+ done.set()
+
+ thread = threading.Thread(target=execute)
+ thread.start()
+
+ # Check if we're in a Jupyter environment
+ try:
+ from IPython.display import display, clear_output
+ from IPython import get_ipython
+
+ in_jupyter = get_ipython() is not None
+ except (ImportError, AttributeError):
+ in_jupyter = False
+
+ start_time = time.time()
+
+ if in_jupyter and callback is None:
+ # Display a simple progress indicator
+ try:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ clear_output(wait=True)
+ print(f"⏳ Query executing... ({elapsed:.1f}s elapsed)")
+
+ clear_output(wait=True)
+ elapsed = time.time() - start_time
+ print(f"✓ Query completed in {elapsed:.1f}s")
+ except Exception:
+ pass # Ignore display errors
+ elif callback is not None:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ callback(f"Executing ({elapsed:.1f}s)", -1.0) # -1 means indeterminate
+
+ elapsed = time.time() - start_time
+ callback(f"Completed in {elapsed:.1f}s", 1.0)
+ else:
+ done.wait()
+
+ thread.join()
+
+ if error[0] is not None:
+ raise error[0]
+
+ return result[0]
+
+
+class ExecutionPlanVisualization:
+ """
+ A wrapper for execution plan visualizations that can render as SVG in Jupyter.
+
+ This class takes the text representation of an execution plan and converts
+ it to a Graphviz DOT format, which is then rendered as SVG.
+ """
+
+ def __init__(self, plan_str: str, analyze: bool = False):
+ self.plan_str = plan_str
+ self.analyze = analyze
+ self._svg_cache: Optional[str] = None
+
+ def _parse_plan_to_dot(self) -> str:
+ """Convert the plan string to DOT format for Graphviz."""
+ lines = self.plan_str.strip().split("\n")
+
+ dot_lines = [
+ "digraph ExecutionPlan {",
+ ' rankdir=TB;',
+ ' node [shape=box, style="rounded,filled", fontname="Helvetica"];',
+ ' edge [fontname="Helvetica"];',
+ "",
+ ]
+
+ nodes = []
+ edges = []
+ node_id = 0
+ stack = [] # (indent_level, node_id)
+
+ for line in lines:
+ if not line.strip():
+ continue
+
+ # Calculate indent level
+ indent = len(line) - len(line.lstrip())
+ content = line.strip()
+
+ # Skip non-plan lines
+ if content.startswith("physical_plan") or content.startswith("logical_plan"):
+ continue
+
+ # Create a node for this plan element
+ current_id = node_id
+ node_id += 1
+
+ # Determine node color based on operation type
+ color = "#E3F2FD" # Default light blue
+ if "Scan" in content or "TableScan" in content:
+ color = "#E8F5E9" # Light green for scans
+ elif "Filter" in content:
+ color = "#FFF3E0" # Light orange for filters
+ elif "Aggregate" in content or "HashAggregate" in content:
+ color = "#F3E5F5" # Light purple for aggregations
+ elif "Join" in content:
+ color = "#FFEBEE" # Light red for joins
+ elif "Sort" in content:
+ color = "#E0F7FA" # Light cyan for sorts
+ elif "Projection" in content:
+ color = "#FFF8E1" # Light amber for projections
+
+ # Escape special characters for DOT format
+ label = content.replace('"', '\\"').replace("\n", "\\n")
+ if len(label) > 60:
+ # Wrap long labels
+ label = label[:57] + "..."
+
+ nodes.append(f' node{current_id} [label="{label}", fillcolor="{color}"];')
+
+ # Connect to parent based on indentation
+ while stack and stack[-1][0] >= indent:
+ stack.pop()
+
+ if stack:
+ parent_id = stack[-1][1]
+ edges.append(f" node{parent_id} -> node{current_id};")
+
+ stack.append((indent, current_id))
+
+ dot_lines.extend(nodes)
+ dot_lines.append("")
+ dot_lines.extend(edges)
+ dot_lines.append("}")
+
+ return "\n".join(dot_lines)
+
+ def to_dot(self) -> str:
+ """Get the DOT representation of the execution plan."""
+ return self._parse_plan_to_dot()
+
+ def to_svg(self) -> str:
+ """
+ Convert the execution plan to SVG format.
+
+ Requires graphviz to be installed. If graphviz is not available,
+ returns a simple HTML representation instead.
+ """
+ if self._svg_cache is not None:
+ return self._svg_cache
+
+ dot_source = self._parse_plan_to_dot()
+
+ try:
+ import subprocess
+
+ # Try to use graphviz's dot command
+ process = subprocess.run(
+ ["dot", "-Tsvg"],
+ input=dot_source.encode(),
+ capture_output=True,
+ timeout=30,
+ )
+
+ if process.returncode == 0:
+ self._svg_cache = process.stdout.decode()
+ return self._svg_cache
+ except (subprocess.SubprocessError, FileNotFoundError, subprocess.TimeoutExpired):
+ pass
+
+ # Fallback: return a pre-formatted HTML representation
+ escaped_plan = (
+ self.plan_str.replace("&", "&")
+ .replace("<", "<")
+ .replace(">", ">")
+ )
+ self._svg_cache = f"""
+ <div style="font-family: monospace; background: #f5f5f5; padding: 10px;
+ border-radius: 5px; overflow-x: auto;">
+ <div style="color: #666; margin-bottom: 5px;">
+ Execution Plan {'(with statistics)' if self.analyze else ''}
+ <br><small>Install graphviz for visual diagram: brew install graphviz</small>
+ </div>
+ <pre style="margin: 0;">{escaped_plan}</pre>
+ </div>
+ """
+ return self._svg_cache
+
+ def save(self, path: str) -> None:
+ """Save the visualization to a file (SVG or DOT format)."""
+ if path.endswith(".dot"):
+ content = self.to_dot()
+ else:
+ content = self.to_svg()
+
+ with open(path, "w") as f:
+ f.write(content)
+
+ def _repr_html_(self) -> str:
+ """HTML representation for Jupyter notebooks."""
+ return self.to_svg()
+
+ def _repr_svg_(self) -> str:
+ """SVG representation for Jupyter notebooks."""
+ svg = self.to_svg()
+ # Only return if it's actual SVG content
+ if svg.strip().startswith("<svg") or svg.strip().startswith("<?xml"):
+ return svg
+ return ""
+
+ def __repr__(self) -> str:
+ """String representation."""
+ return f"ExecutionPlanVisualization(analyze={self.analyze})\n{self.plan_str}"
+
class BallistaSessionContext(SessionContext, metaclass=RedefiningSessionContextMeta):
+ """
+ A session context for connecting to and querying a Ballista cluster.
+
+ This class extends DataFusion's SessionContext to work with distributed
+ Ballista clusters, automatically routing query execution to the cluster
+ while maintaining API compatibility with local DataFusion usage.
+
+ Example:
+ >>> from ballista import BallistaSessionContext
+ >>> ctx = BallistaSessionContext("df://localhost:50050")
+ >>> df = ctx.sql("SELECT * FROM my_table LIMIT 10")
+ >>> df.show()
+
+ For Jupyter notebook users:
+ >>> %load_ext ballista.jupyter
+ >>> %ballista connect df://localhost:50050
+ >>> %sql SELECT * FROM my_table
+ """
+
def __init__(self, address: str, config=None, runtime=None):
super().__init__(config, runtime)
self.address = address
self.session_id = self.session_id()
Review Comment:
Changed to @property decorator
##########
python/python/ballista/extension.py:
##########
@@ -229,9 +231,319 @@ def write_parquet(
df = self._to_internal_df()
df.write_parquet(str(path), compression.value, compression_level)
+ def explain_visual(self, analyze: bool = False) -> "ExecutionPlanVisualization":
+ """
+ Generate a visual representation of the execution plan.
+
+ This method creates an SVG visualization of the query execution plan,
+ which can be displayed directly in Jupyter notebooks.
+
+ Args:
+ analyze: If True, includes runtime statistics from actual execution.
+
+ Returns:
+ ExecutionPlanVisualization: An object that renders as SVG in Jupyter.
+
+ Example:
+ >>> df = ctx.sql("SELECT * FROM orders WHERE amount > 100")
+ >>> df.explain_visual() # Displays SVG in notebook
+ >>> viz = df.explain_visual(analyze=True)
+ >>> viz.save("plan.svg") # Save to file
+ """
+ # Get the execution plan as a string representation
+ # Note: explain() prints but doesn't return a string, so we use logical_plan()
+ try:
+ plan = self.logical_plan()
+ plan_str = plan.display_indent()
+ except Exception:
+ # Fallback if logical_plan() fails
+ plan_str = "Unable to retrieve execution plan"
+ return ExecutionPlanVisualization(plan_str, analyze=analyze)
+
+ def collect_with_progress(
+ self,
+ callback: Optional[callable] = None,
+ poll_interval: float = 0.5,
+ ):
+ """
+ Collect results with progress indication.
+
+ For long-running queries, this method provides progress updates
+ through a callback function or displays a progress bar in Jupyter.
+
+ Args:
+ callback: Optional function to call with progress updates.
+ Signature: callback(status: str, progress: float)
+ poll_interval: How often to check progress (seconds).
+
+ Returns:
+ The collected result batches.
+
+ Example:
+ >>> def my_callback(status, progress):
+ ... print(f"{status}: {progress:.1%}")
+ >>> batches = df.collect_with_progress(callback=my_callback)
+ """
+ import threading
+ import time
+
+ result = [None]
+ error = [None]
+ done = threading.Event()
+
+ def execute():
+ try:
+ result[0] = self.collect()
+ except Exception as e:
+ error[0] = e
+ finally:
+ done.set()
+
+ thread = threading.Thread(target=execute)
+ thread.start()
+
+ # Check if we're in a Jupyter environment
+ try:
+ from IPython.display import display, clear_output
+ from IPython import get_ipython
+
+ in_jupyter = get_ipython() is not None
+ except (ImportError, AttributeError):
+ in_jupyter = False
+
+ start_time = time.time()
+
+ if in_jupyter and callback is None:
+ # Display a simple progress indicator
+ try:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ clear_output(wait=True)
+ print(f"⏳ Query executing... ({elapsed:.1f}s elapsed)")
+
+ clear_output(wait=True)
+ elapsed = time.time() - start_time
+ print(f"✓ Query completed in {elapsed:.1f}s")
+ except Exception:
+ pass # Ignore display errors
+ elif callback is not None:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ callback(f"Executing ({elapsed:.1f}s)", -1.0) # -1 means indeterminate
+
+ elapsed = time.time() - start_time
+ callback(f"Completed in {elapsed:.1f}s", 1.0)
+ else:
+ done.wait()
+
+ thread.join()
+
+ if error[0] is not None:
+ raise error[0]
+
+ return result[0]
+
+
+class ExecutionPlanVisualization:
+ """
+ A wrapper for execution plan visualizations that can render as SVG in Jupyter.
+
+ This class takes the text representation of an execution plan and converts
+ it to a Graphviz DOT format, which is then rendered as SVG.
+ """
+
+ def __init__(self, plan_str: str, analyze: bool = False):
+ self.plan_str = plan_str
+ self.analyze = analyze
+ self._svg_cache: Optional[str] = None
+
+ def _parse_plan_to_dot(self) -> str:
+ """Convert the plan string to DOT format for Graphviz."""
+ lines = self.plan_str.strip().split("\n")
+
+ dot_lines = [
+ "digraph ExecutionPlan {",
+ ' rankdir=TB;',
+ ' node [shape=box, style="rounded,filled", fontname="Helvetica"];',
+ ' edge [fontname="Helvetica"];',
+ "",
+ ]
+
+ nodes = []
+ edges = []
+ node_id = 0
+ stack = [] # (indent_level, node_id)
+
+ for line in lines:
+ if not line.strip():
+ continue
+
+ # Calculate indent level
+ indent = len(line) - len(line.lstrip())
+ content = line.strip()
+
+ # Skip non-plan lines
+ if content.startswith("physical_plan") or content.startswith("logical_plan"):
+ continue
+
+ # Create a node for this plan element
+ current_id = node_id
+ node_id += 1
+
+ # Determine node color based on operation type
+ color = "#E3F2FD" # Default light blue
+ if "Scan" in content or "TableScan" in content:
+ color = "#E8F5E9" # Light green for scans
+ elif "Filter" in content:
+ color = "#FFF3E0" # Light orange for filters
+ elif "Aggregate" in content or "HashAggregate" in content:
+ color = "#F3E5F5" # Light purple for aggregations
+ elif "Join" in content:
+ color = "#FFEBEE" # Light red for joins
+ elif "Sort" in content:
+ color = "#E0F7FA" # Light cyan for sorts
+ elif "Projection" in content:
+ color = "#FFF8E1" # Light amber for projections
+
+ # Escape special characters for DOT format
+ label = content.replace('"', '\\"').replace("\n", "\\n")
+ if len(label) > 60:
+ # Wrap long labels
+ label = label[:57] + "..."
+
+ nodes.append(f' node{current_id} [label="{label}", fillcolor="{color}"];')
+
+ # Connect to parent based on indentation
+ while stack and stack[-1][0] >= indent:
+ stack.pop()
+
+ if stack:
+ parent_id = stack[-1][1]
+ edges.append(f" node{parent_id} -> node{current_id};")
+
+ stack.append((indent, current_id))
+
+ dot_lines.extend(nodes)
+ dot_lines.append("")
+ dot_lines.extend(edges)
+ dot_lines.append("}")
+
+ return "\n".join(dot_lines)
+
+ def to_dot(self) -> str:
+ """Get the DOT representation of the execution plan."""
+ return self._parse_plan_to_dot()
+
+ def to_svg(self) -> str:
+ """
+ Convert the execution plan to SVG format.
+
+ Requires graphviz to be installed. If graphviz is not available,
+ returns a simple HTML representation instead.
+ """
+ if self._svg_cache is not None:
+ return self._svg_cache
+
+ dot_source = self._parse_plan_to_dot()
+
+ try:
+ import subprocess
+
+ # Try to use graphviz's dot command
+ process = subprocess.run(
+ ["dot", "-Tsvg"],
+ input=dot_source.encode(),
+ capture_output=True,
+ timeout=30,
+ )
+
+ if process.returncode == 0:
+ self._svg_cache = process.stdout.decode()
+ return self._svg_cache
+ except (subprocess.SubprocessError, FileNotFoundError, subprocess.TimeoutExpired):
+ pass
+
+ # Fallback: return a pre-formatted HTML representation
+ escaped_plan = (
+ self.plan_str.replace("&", "&")
+ .replace("<", "<")
+ .replace(">", ">")
+ )
+ self._svg_cache = f"""
+ <div style="font-family: monospace; background: #f5f5f5; padding: 10px;
+ border-radius: 5px; overflow-x: auto;">
+ <div style="color: #666; margin-bottom: 5px;">
+ Execution Plan {'(with statistics)' if self.analyze else ''}
+ <br><small>Install graphviz for visual diagram: brew install graphviz</small>
+ </div>
+ <pre style="margin: 0;">{escaped_plan}</pre>
+ </div>
+ """
+ return self._svg_cache
+
+ def save(self, path: str) -> None:
+ """Save the visualization to a file (SVG or DOT format)."""
+ if path.endswith(".dot"):
+ content = self.to_dot()
+ else:
+ content = self.to_svg()
+
+ with open(path, "w") as f:
+ f.write(content)
+
+ def _repr_html_(self) -> str:
+ """HTML representation for Jupyter notebooks."""
+ return self.to_svg()
+
+ def _repr_svg_(self) -> str:
+ """SVG representation for Jupyter notebooks."""
+ svg = self.to_svg()
+ # Only return if it's actual SVG content
+ if svg.strip().startswith("<svg") or svg.strip().startswith("<?xml"):
+ return svg
+ return ""
+
+ def __repr__(self) -> str:
+ """String representation."""
+ return f"ExecutionPlanVisualization(analyze={self.analyze})\n{self.plan_str}"
+
class BallistaSessionContext(SessionContext, metaclass=RedefiningSessionContextMeta):
+ """
+ A session context for connecting to and querying a Ballista cluster.
+
+ This class extends DataFusion's SessionContext to work with distributed
+ Ballista clusters, automatically routing query execution to the cluster
+ while maintaining API compatibility with local DataFusion usage.
+
+ Example:
+ >>> from ballista import BallistaSessionContext
+ >>> ctx = BallistaSessionContext("df://localhost:50050")
+ >>> df = ctx.sql("SELECT * FROM my_table LIMIT 10")
+ >>> df.show()
+
+ For Jupyter notebook users:
+ >>> %load_ext ballista.jupyter
+ >>> %ballista connect df://localhost:50050
+ >>> %sql SELECT * FROM my_table
+ """
+
def __init__(self, address: str, config=None, runtime=None):
super().__init__(config, runtime)
self.address = address
self.session_id = self.session_id()
+
+ def get_tables(self) -> Optional[dict[str, str]]:
Review Comment:
Fixed. Thank you
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feature/jupyter notebook support [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on code in PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513#discussion_r2946184117
##########
python/examples/getting_started.ipynb:
##########
@@ -0,0 +1,353 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "<!---\n",
+ " Licensed to the Apache Software Foundation (ASF) under one\n",
+ " or more contributor license agreements. See the NOTICE file\n",
+ " distributed with this work for additional information\n",
+ " regarding copyright ownership. The ASF licenses this file\n",
+ " to you under the Apache License, Version 2.0 (the\n",
+ " \"License\"); you may not use this file except in compliance\n",
+ " with the License. You may obtain a copy of the License at\n",
+ "\n",
+ " http://www.apache.org/licenses/LICENSE-2.0\n",
+ "\n",
+ " Unless required by applicable law or agreed to in writing,\n",
+ " software distributed under the License is distributed on an\n",
+ " \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY\n",
+ " KIND, either express or implied. See the License for the\n",
+ " specific language governing permissions and limitations\n",
+ " under the License.\n",
+ "-->\n",
+ "\n",
+ "# Getting Started with PyBallista\n",
+ "\n",
+ "This notebook demonstrates how to get started with Ballista using Python.\n",
+ "\n",
+ "## Prerequisites\n",
+ "\n",
+ "1. Install PyBallista: `pip install ballista`\n",
+ "2. Have a Ballista cluster running (or use the built-in test cluster)\n",
+ "\n",
+ "## Overview\n",
+ "\n",
+ "Ballista is a distributed query engine built on Apache DataFusion. PyBallista provides:\n",
+ "\n",
+ "- **BallistaSessionContext**: Drop-in replacement for DataFusion's SessionContext\n",
+ "- **SQL Magic Commands**: Interactive SQL in Jupyter notebooks via `%sql` and `%%sql`\n",
+ "- **DataFrame API**: Full DataFrame API for data transformations\n",
+ "- **Rich HTML Display**: DataFrames render as styled HTML tables"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Method 1: Python API\n",
+ "\n",
+ "The most straightforward way to use Ballista is via the Python API."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import ballista\n",
+ "from ballista import BallistaSessionContext, setup_test_cluster\n",
+ "%load_ext autoreload\n",
+ "%autoreload 2\n",
+ "\n",
+ "# Check versions\n",
+ "print(f\"Ballista version: {ballista.__version__}\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# For this demo, we'll use the built-in test cluster\n",
+ "# In production, you would connect to your Ballista scheduler:\n",
+ "# ctx = BallistaSessionContext(\"df://your-scheduler:50050\")\n",
+ "\n",
+ "# host, port = setup_test_cluster()\n",
+ "host, port = \"localhost\", \"39431\"\n",
Review Comment:
Yes, debug leftover. Fixed
##########
python/python/ballista/extension.py:
##########
@@ -229,9 +231,319 @@ def write_parquet(
df = self._to_internal_df()
df.write_parquet(str(path), compression.value, compression_level)
+ def explain_visual(self, analyze: bool = False) -> "ExecutionPlanVisualization":
+ """
+ Generate a visual representation of the execution plan.
+
+ This method creates an SVG visualization of the query execution plan,
+ which can be displayed directly in Jupyter notebooks.
+
+ Args:
+ analyze: If True, includes runtime statistics from actual execution.
+
+ Returns:
+ ExecutionPlanVisualization: An object that renders as SVG in Jupyter.
+
+ Example:
+ >>> df = ctx.sql("SELECT * FROM orders WHERE amount > 100")
+ >>> df.explain_visual() # Displays SVG in notebook
+ >>> viz = df.explain_visual(analyze=True)
+ >>> viz.save("plan.svg") # Save to file
+ """
+ # Get the execution plan as a string representation
+ # Note: explain() prints but doesn't return a string, so we use logical_plan()
+ try:
+ plan = self.logical_plan()
+ plan_str = plan.display_indent()
+ except Exception:
+ # Fallback if logical_plan() fails
+ plan_str = "Unable to retrieve execution plan"
+ return ExecutionPlanVisualization(plan_str, analyze=analyze)
+
+ def collect_with_progress(
+ self,
+ callback: Optional[callable] = None,
Review Comment:
Fixed. Thank you
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feature/jupyter notebook support [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on code in PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513#discussion_r2946196370
##########
python/python/ballista/extension.py:
##########
@@ -229,9 +231,319 @@ def write_parquet(
df = self._to_internal_df()
df.write_parquet(str(path), compression.value, compression_level)
+ def explain_visual(self, analyze: bool = False) -> "ExecutionPlanVisualization":
+ """
+ Generate a visual representation of the execution plan.
+
+ This method creates an SVG visualization of the query execution plan,
+ which can be displayed directly in Jupyter notebooks.
+
+ Args:
+ analyze: If True, includes runtime statistics from actual execution.
+
+ Returns:
+ ExecutionPlanVisualization: An object that renders as SVG in Jupyter.
+
+ Example:
+ >>> df = ctx.sql("SELECT * FROM orders WHERE amount > 100")
+ >>> df.explain_visual() # Displays SVG in notebook
+ >>> viz = df.explain_visual(analyze=True)
+ >>> viz.save("plan.svg") # Save to file
+ """
+ # Get the execution plan as a string representation
+ # Note: explain() prints but doesn't return a string, so we use logical_plan()
+ try:
+ plan = self.logical_plan()
+ plan_str = plan.display_indent()
+ except Exception:
+ # Fallback if logical_plan() fails
+ plan_str = "Unable to retrieve execution plan"
+ return ExecutionPlanVisualization(plan_str, analyze=analyze)
+
+ def collect_with_progress(
+ self,
+ callback: Optional[callable] = None,
+ poll_interval: float = 0.5,
+ ):
+ """
+ Collect results with progress indication.
+
+ For long-running queries, this method provides progress updates
+ through a callback function or displays a progress bar in Jupyter.
+
+ Args:
+ callback: Optional function to call with progress updates.
+ Signature: callback(status: str, progress: float)
+ poll_interval: How often to check progress (seconds).
+
+ Returns:
+ The collected result batches.
+
+ Example:
+ >>> def my_callback(status, progress):
+ ... print(f"{status}: {progress:.1%}")
+ >>> batches = df.collect_with_progress(callback=my_callback)
+ """
+ import threading
+ import time
+
+ result = [None]
+ error = [None]
+ done = threading.Event()
+
+ def execute():
+ try:
+ result[0] = self.collect()
+ except Exception as e:
+ error[0] = e
+ finally:
+ done.set()
+
+ thread = threading.Thread(target=execute)
+ thread.start()
+
+ # Check if we're in a Jupyter environment
+ try:
+ from IPython.display import display, clear_output
+ from IPython import get_ipython
+
+ in_jupyter = get_ipython() is not None
+ except (ImportError, AttributeError):
+ in_jupyter = False
+
+ start_time = time.time()
+
+ if in_jupyter and callback is None:
+ # Display a simple progress indicator
+ try:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ clear_output(wait=True)
+ print(f"⏳ Query executing... ({elapsed:.1f}s elapsed)")
+
+ clear_output(wait=True)
+ elapsed = time.time() - start_time
+ print(f"✓ Query completed in {elapsed:.1f}s")
+ except Exception:
+ pass # Ignore display errors
Review Comment:
Changed all passes to warnings.warn() were applicable
##########
python/python/ballista/extension.py:
##########
@@ -229,9 +231,319 @@ def write_parquet(
df = self._to_internal_df()
df.write_parquet(str(path), compression.value, compression_level)
+ def explain_visual(self, analyze: bool = False) -> "ExecutionPlanVisualization":
+ """
+ Generate a visual representation of the execution plan.
+
+ This method creates an SVG visualization of the query execution plan,
+ which can be displayed directly in Jupyter notebooks.
+
+ Args:
+ analyze: If True, includes runtime statistics from actual execution.
+
+ Returns:
+ ExecutionPlanVisualization: An object that renders as SVG in Jupyter.
+
+ Example:
+ >>> df = ctx.sql("SELECT * FROM orders WHERE amount > 100")
+ >>> df.explain_visual() # Displays SVG in notebook
+ >>> viz = df.explain_visual(analyze=True)
+ >>> viz.save("plan.svg") # Save to file
+ """
+ # Get the execution plan as a string representation
+ # Note: explain() prints but doesn't return a string, so we use logical_plan()
+ try:
+ plan = self.logical_plan()
+ plan_str = plan.display_indent()
+ except Exception:
+ # Fallback if logical_plan() fails
+ plan_str = "Unable to retrieve execution plan"
+ return ExecutionPlanVisualization(plan_str, analyze=analyze)
+
+ def collect_with_progress(
+ self,
+ callback: Optional[callable] = None,
+ poll_interval: float = 0.5,
+ ):
+ """
+ Collect results with progress indication.
+
+ For long-running queries, this method provides progress updates
+ through a callback function or displays a progress bar in Jupyter.
+
+ Args:
+ callback: Optional function to call with progress updates.
+ Signature: callback(status: str, progress: float)
+ poll_interval: How often to check progress (seconds).
+
+ Returns:
+ The collected result batches.
+
+ Example:
+ >>> def my_callback(status, progress):
+ ... print(f"{status}: {progress:.1%}")
+ >>> batches = df.collect_with_progress(callback=my_callback)
+ """
+ import threading
+ import time
+
+ result = [None]
+ error = [None]
+ done = threading.Event()
+
+ def execute():
+ try:
+ result[0] = self.collect()
+ except Exception as e:
+ error[0] = e
+ finally:
+ done.set()
+
+ thread = threading.Thread(target=execute)
+ thread.start()
+
+ # Check if we're in a Jupyter environment
+ try:
+ from IPython.display import display, clear_output
+ from IPython import get_ipython
+
+ in_jupyter = get_ipython() is not None
+ except (ImportError, AttributeError):
+ in_jupyter = False
+
+ start_time = time.time()
+
+ if in_jupyter and callback is None:
+ # Display a simple progress indicator
+ try:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ clear_output(wait=True)
+ print(f"⏳ Query executing... ({elapsed:.1f}s elapsed)")
+
+ clear_output(wait=True)
+ elapsed = time.time() - start_time
+ print(f"✓ Query completed in {elapsed:.1f}s")
+ except Exception:
+ pass # Ignore display errors
+ elif callback is not None:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ callback(f"Executing ({elapsed:.1f}s)", -1.0) # -1 means indeterminate
+
+ elapsed = time.time() - start_time
+ callback(f"Completed in {elapsed:.1f}s", 1.0)
+ else:
+ done.wait()
+
+ thread.join()
+
+ if error[0] is not None:
+ raise error[0]
+
+ return result[0]
+
+
+class ExecutionPlanVisualization:
+ """
+ A wrapper for execution plan visualizations that can render as SVG in Jupyter.
+
+ This class takes the text representation of an execution plan and converts
+ it to a Graphviz DOT format, which is then rendered as SVG.
+ """
+
+ def __init__(self, plan_str: str, analyze: bool = False):
+ self.plan_str = plan_str
+ self.analyze = analyze
+ self._svg_cache: Optional[str] = None
+
+ def _parse_plan_to_dot(self) -> str:
+ """Convert the plan string to DOT format for Graphviz."""
+ lines = self.plan_str.strip().split("\n")
+
+ dot_lines = [
+ "digraph ExecutionPlan {",
+ ' rankdir=TB;',
+ ' node [shape=box, style="rounded,filled", fontname="Helvetica"];',
+ ' edge [fontname="Helvetica"];',
+ "",
+ ]
+
+ nodes = []
+ edges = []
+ node_id = 0
+ stack = [] # (indent_level, node_id)
+
+ for line in lines:
+ if not line.strip():
+ continue
+
+ # Calculate indent level
+ indent = len(line) - len(line.lstrip())
+ content = line.strip()
+
+ # Skip non-plan lines
+ if content.startswith("physical_plan") or content.startswith("logical_plan"):
+ continue
+
+ # Create a node for this plan element
+ current_id = node_id
+ node_id += 1
+
+ # Determine node color based on operation type
+ color = "#E3F2FD" # Default light blue
+ if "Scan" in content or "TableScan" in content:
+ color = "#E8F5E9" # Light green for scans
+ elif "Filter" in content:
+ color = "#FFF3E0" # Light orange for filters
+ elif "Aggregate" in content or "HashAggregate" in content:
+ color = "#F3E5F5" # Light purple for aggregations
+ elif "Join" in content:
+ color = "#FFEBEE" # Light red for joins
+ elif "Sort" in content:
+ color = "#E0F7FA" # Light cyan for sorts
+ elif "Projection" in content:
+ color = "#FFF8E1" # Light amber for projections
+
+ # Escape special characters for DOT format
+ label = content.replace('"', '\\"').replace("\n", "\\n")
+ if len(label) > 60:
+ # Wrap long labels
+ label = label[:57] + "..."
+
+ nodes.append(f' node{current_id} [label="{label}", fillcolor="{color}"];')
+
+ # Connect to parent based on indentation
+ while stack and stack[-1][0] >= indent:
+ stack.pop()
+
+ if stack:
+ parent_id = stack[-1][1]
+ edges.append(f" node{parent_id} -> node{current_id};")
+
+ stack.append((indent, current_id))
+
+ dot_lines.extend(nodes)
+ dot_lines.append("")
+ dot_lines.extend(edges)
+ dot_lines.append("}")
+
+ return "\n".join(dot_lines)
+
+ def to_dot(self) -> str:
+ """Get the DOT representation of the execution plan."""
+ return self._parse_plan_to_dot()
+
+ def to_svg(self) -> str:
+ """
+ Convert the execution plan to SVG format.
+
+ Requires graphviz to be installed. If graphviz is not available,
+ returns a simple HTML representation instead.
+ """
+ if self._svg_cache is not None:
+ return self._svg_cache
+
+ dot_source = self._parse_plan_to_dot()
+
+ try:
+ import subprocess
+
+ # Try to use graphviz's dot command
+ process = subprocess.run(
+ ["dot", "-Tsvg"],
+ input=dot_source.encode(),
+ capture_output=True,
+ timeout=30,
+ )
+
+ if process.returncode == 0:
+ self._svg_cache = process.stdout.decode()
+ return self._svg_cache
+ except (subprocess.SubprocessError, FileNotFoundError, subprocess.TimeoutExpired):
+ pass
Review Comment:
Fixed
##########
python/python/ballista/extension.py:
##########
@@ -229,9 +231,319 @@ def write_parquet(
df = self._to_internal_df()
df.write_parquet(str(path), compression.value, compression_level)
+ def explain_visual(self, analyze: bool = False) -> "ExecutionPlanVisualization":
+ """
+ Generate a visual representation of the execution plan.
+
+ This method creates an SVG visualization of the query execution plan,
+ which can be displayed directly in Jupyter notebooks.
+
+ Args:
+ analyze: If True, includes runtime statistics from actual execution.
+
+ Returns:
+ ExecutionPlanVisualization: An object that renders as SVG in Jupyter.
+
+ Example:
+ >>> df = ctx.sql("SELECT * FROM orders WHERE amount > 100")
+ >>> df.explain_visual() # Displays SVG in notebook
+ >>> viz = df.explain_visual(analyze=True)
+ >>> viz.save("plan.svg") # Save to file
+ """
+ # Get the execution plan as a string representation
+ # Note: explain() prints but doesn't return a string, so we use logical_plan()
+ try:
+ plan = self.logical_plan()
+ plan_str = plan.display_indent()
+ except Exception:
+ # Fallback if logical_plan() fails
+ plan_str = "Unable to retrieve execution plan"
+ return ExecutionPlanVisualization(plan_str, analyze=analyze)
+
+ def collect_with_progress(
+ self,
+ callback: Optional[callable] = None,
+ poll_interval: float = 0.5,
+ ):
+ """
+ Collect results with progress indication.
+
+ For long-running queries, this method provides progress updates
+ through a callback function or displays a progress bar in Jupyter.
+
+ Args:
+ callback: Optional function to call with progress updates.
+ Signature: callback(status: str, progress: float)
+ poll_interval: How often to check progress (seconds).
+
+ Returns:
+ The collected result batches.
+
+ Example:
+ >>> def my_callback(status, progress):
+ ... print(f"{status}: {progress:.1%}")
+ >>> batches = df.collect_with_progress(callback=my_callback)
+ """
+ import threading
+ import time
+
+ result = [None]
+ error = [None]
+ done = threading.Event()
+
+ def execute():
+ try:
+ result[0] = self.collect()
+ except Exception as e:
+ error[0] = e
+ finally:
+ done.set()
+
+ thread = threading.Thread(target=execute)
+ thread.start()
+
+ # Check if we're in a Jupyter environment
+ try:
+ from IPython.display import display, clear_output
+ from IPython import get_ipython
+
+ in_jupyter = get_ipython() is not None
+ except (ImportError, AttributeError):
+ in_jupyter = False
+
+ start_time = time.time()
+
+ if in_jupyter and callback is None:
+ # Display a simple progress indicator
+ try:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ clear_output(wait=True)
+ print(f"⏳ Query executing... ({elapsed:.1f}s elapsed)")
+
+ clear_output(wait=True)
+ elapsed = time.time() - start_time
+ print(f"✓ Query completed in {elapsed:.1f}s")
+ except Exception:
+ pass # Ignore display errors
+ elif callback is not None:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ callback(f"Executing ({elapsed:.1f}s)", -1.0) # -1 means indeterminate
+
+ elapsed = time.time() - start_time
+ callback(f"Completed in {elapsed:.1f}s", 1.0)
+ else:
+ done.wait()
+
+ thread.join()
+
+ if error[0] is not None:
+ raise error[0]
+
+ return result[0]
+
+
+class ExecutionPlanVisualization:
+ """
+ A wrapper for execution plan visualizations that can render as SVG in Jupyter.
+
+ This class takes the text representation of an execution plan and converts
+ it to a Graphviz DOT format, which is then rendered as SVG.
+ """
+
+ def __init__(self, plan_str: str, analyze: bool = False):
+ self.plan_str = plan_str
+ self.analyze = analyze
+ self._svg_cache: Optional[str] = None
+
+ def _parse_plan_to_dot(self) -> str:
+ """Convert the plan string to DOT format for Graphviz."""
+ lines = self.plan_str.strip().split("\n")
+
+ dot_lines = [
+ "digraph ExecutionPlan {",
+ ' rankdir=TB;',
+ ' node [shape=box, style="rounded,filled", fontname="Helvetica"];',
+ ' edge [fontname="Helvetica"];',
+ "",
+ ]
+
+ nodes = []
+ edges = []
+ node_id = 0
+ stack = [] # (indent_level, node_id)
+
+ for line in lines:
+ if not line.strip():
+ continue
+
+ # Calculate indent level
+ indent = len(line) - len(line.lstrip())
+ content = line.strip()
+
+ # Skip non-plan lines
+ if content.startswith("physical_plan") or content.startswith("logical_plan"):
+ continue
+
+ # Create a node for this plan element
+ current_id = node_id
+ node_id += 1
+
+ # Determine node color based on operation type
+ color = "#E3F2FD" # Default light blue
+ if "Scan" in content or "TableScan" in content:
+ color = "#E8F5E9" # Light green for scans
+ elif "Filter" in content:
+ color = "#FFF3E0" # Light orange for filters
+ elif "Aggregate" in content or "HashAggregate" in content:
+ color = "#F3E5F5" # Light purple for aggregations
+ elif "Join" in content:
+ color = "#FFEBEE" # Light red for joins
+ elif "Sort" in content:
+ color = "#E0F7FA" # Light cyan for sorts
+ elif "Projection" in content:
+ color = "#FFF8E1" # Light amber for projections
+
+ # Escape special characters for DOT format
+ label = content.replace('"', '\\"').replace("\n", "\\n")
+ if len(label) > 60:
+ # Wrap long labels
+ label = label[:57] + "..."
+
+ nodes.append(f' node{current_id} [label="{label}", fillcolor="{color}"];')
+
+ # Connect to parent based on indentation
+ while stack and stack[-1][0] >= indent:
+ stack.pop()
+
+ if stack:
+ parent_id = stack[-1][1]
+ edges.append(f" node{parent_id} -> node{current_id};")
+
+ stack.append((indent, current_id))
+
+ dot_lines.extend(nodes)
+ dot_lines.append("")
+ dot_lines.extend(edges)
+ dot_lines.append("}")
+
+ return "\n".join(dot_lines)
+
+ def to_dot(self) -> str:
+ """Get the DOT representation of the execution plan."""
+ return self._parse_plan_to_dot()
+
+ def to_svg(self) -> str:
+ """
+ Convert the execution plan to SVG format.
+
+ Requires graphviz to be installed. If graphviz is not available,
+ returns a simple HTML representation instead.
+ """
+ if self._svg_cache is not None:
+ return self._svg_cache
+
+ dot_source = self._parse_plan_to_dot()
+
+ try:
+ import subprocess
+
+ # Try to use graphviz's dot command
+ process = subprocess.run(
+ ["dot", "-Tsvg"],
+ input=dot_source.encode(),
+ capture_output=True,
+ timeout=30,
+ )
+
+ if process.returncode == 0:
+ self._svg_cache = process.stdout.decode()
+ return self._svg_cache
+ except (subprocess.SubprocessError, FileNotFoundError, subprocess.TimeoutExpired):
+ pass
+
+ # Fallback: return a pre-formatted HTML representation
+ escaped_plan = (
+ self.plan_str.replace("&", "&")
+ .replace("<", "<")
+ .replace(">", ">")
+ )
+ self._svg_cache = f"""
+ <div style="font-family: monospace; background: #f5f5f5; padding: 10px;
+ border-radius: 5px; overflow-x: auto;">
+ <div style="color: #666; margin-bottom: 5px;">
+ Execution Plan {'(with statistics)' if self.analyze else ''}
+ <br><small>Install graphviz for visual diagram: brew install graphviz</small>
+ </div>
+ <pre style="margin: 0;">{escaped_plan}</pre>
+ </div>
+ """
+ return self._svg_cache
+
+ def save(self, path: str) -> None:
+ """Save the visualization to a file (SVG or DOT format)."""
+ if path.endswith(".dot"):
+ content = self.to_dot()
+ else:
+ content = self.to_svg()
+
+ with open(path, "w") as f:
+ f.write(content)
+
+ def _repr_html_(self) -> str:
+ """HTML representation for Jupyter notebooks."""
+ return self.to_svg()
+
+ def _repr_svg_(self) -> str:
+ """SVG representation for Jupyter notebooks."""
+ svg = self.to_svg()
+ # Only return if it's actual SVG content
+ if svg.strip().startswith("<svg") or svg.strip().startswith("<?xml"):
+ return svg
+ return ""
+
+ def __repr__(self) -> str:
+ """String representation."""
+ return f"ExecutionPlanVisualization(analyze={self.analyze})\n{self.plan_str}"
+
class BallistaSessionContext(SessionContext, metaclass=RedefiningSessionContextMeta):
+ """
+ A session context for connecting to and querying a Ballista cluster.
+
+ This class extends DataFusion's SessionContext to work with distributed
+ Ballista clusters, automatically routing query execution to the cluster
+ while maintaining API compatibility with local DataFusion usage.
+
+ Example:
+ >>> from ballista import BallistaSessionContext
+ >>> ctx = BallistaSessionContext("df://localhost:50050")
+ >>> df = ctx.sql("SELECT * FROM my_table LIMIT 10")
+ >>> df.show()
+
+ For Jupyter notebook users:
+ >>> %load_ext ballista.jupyter
+ >>> %ballista connect df://localhost:50050
+ >>> %sql SELECT * FROM my_table
+ """
+
def __init__(self, address: str, config=None, runtime=None):
super().__init__(config, runtime)
self.address = address
self.session_id = self.session_id()
+
+ def get_tables(self) -> Optional[dict[str, str]]:
+ """Get tables and their respective schemas (in terms of database schema)."""
+ try:
+ catalog = self.catalog()
+ schema_names = list(catalog.schema_names())
+ if schema_names is not None:
+ tables_info = {}
+ for schema_name in schema_names:
+ tables_info[schema_name] = list(catalog.schema(name=schema_name).table_names())
+ return tables_info
+ except (AttributeError, NotImplementedError):
+ pass
Review Comment:
Fixed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feature/jupyter notebook support [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on code in PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513#discussion_r2946197574
##########
python/python/ballista/jupyter.py:
##########
@@ -0,0 +1,493 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Ballista Jupyter Magic Commands.
+
+This module provides IPython magic commands for interacting with Ballista
+clusters directly from Jupyter notebooks.
+
+Usage:
+ %load_ext ballista.jupyter
+
+ # Connect to a Ballista cluster
+ %ballista connect df://localhost:50050
+
+ # Check connection status
+ %ballista status
+
+ # List registered tables
+ %ballista tables
+
+ # Show schema for a table
+ %ballista schema my_table
+
+ # Execute a simple SQL query (line magic)
+ %sql SELECT COUNT(*) FROM my_table
+
+ # Execute a complex SQL query (cell magic)
+ %%sql
+ SELECT
+ customer_id,
+ SUM(amount) as total
+ FROM orders
+ GROUP BY customer_id
+ ORDER BY total DESC
+ LIMIT 10
+"""
+
+from typing import Optional, List, Dict, Any
+import time
+import threading
+
+try:
+ from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic
+ from IPython.display import display, HTML
+
+ IPYTHON_AVAILABLE = True
+except ImportError:
+ IPYTHON_AVAILABLE = False
+
+ # Provide stub classes for when IPython is not available
+ class Magics:
+ def __init__(self, shell=None):
+ self.shell = shell
+
+ def magics_class(cls):
+ return cls
+
+ def line_magic(name_or_func):
+ """Stub line_magic decorator for when IPython is not available."""
+ # Handle both @line_magic and @line_magic("name") usage
+ if callable(name_or_func):
+ # Used as @line_magic without arguments
+ return name_or_func
+ else:
+ # Used as @line_magic("name") with arguments
+ def decorator(func):
+ return func
+ return decorator
+
+ def cell_magic(name_or_func):
+ """Stub cell_magic decorator for when IPython is not available."""
+ # Handle both @cell_magic and @cell_magic("name") usage
+ if callable(name_or_func):
+ return name_or_func
+ else:
+ def decorator(func):
+ return func
+ return decorator
+
+
+from .extension import BallistaSessionContext, DistributedDataFrame
+
+
+class BallistaConnectionError(Exception):
+ """Raised when not connected to a Ballista cluster."""
+
+ pass
+
+
+@magics_class
+class BallistaMagics(Magics):
+ """
+ IPython magic commands for Ballista.
+
+ Provides convenient commands for connecting to Ballista clusters,
+ executing SQL queries, and exploring table schemas.
+ """
+
+ def __init__(self, shell=None):
+ super().__init__(shell)
+ self._ctx: Optional[BallistaSessionContext] = None
+ self._address: Optional[str] = None
+ self._last_result: Optional[DistributedDataFrame] = None
+ self._query_history: List[Dict[str, Any]] = []
+
+ @property
+ def ctx(self) -> BallistaSessionContext:
+ """Get the current context, raising an error if not connected."""
+ if self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ return self._ctx
+
+ @property
+ def is_connected(self) -> bool:
+ """Check if connected to a Ballista cluster."""
+ return self._ctx is not None
+
+ @line_magic
+ def ballista(self, line: str) -> Optional[str]:
+ """
+ Ballista management commands.
+
+ Usage:
+ %ballista connect df://localhost:50050 - Connect to cluster
+ %ballista status - Show connection status
+ %ballista tables - List registered tables
+ %ballista schema <table> - Show table schema
+ %ballista disconnect - Disconnect from cluster
+ %ballista history - Show query history
+
+ Examples:
+ %ballista connect df://localhost:50050
+ %ballista tables
+ %ballista schema orders
+ """
+ parts = line.strip().split(maxsplit=1)
+ if not parts:
+ return self._show_help()
+
+ cmd = parts[0].lower()
+ args = parts[1] if len(parts) > 1 else ""
+
+ if cmd == "connect":
+ return self._connect(args)
+ elif cmd == "status":
+ return self._status()
+ elif cmd == "tables":
+ return self._tables()
+ elif cmd == "schema":
+ return self._schema(args)
+ elif cmd == "disconnect":
+ return self._disconnect()
+ elif cmd == "history":
+ return self._show_history()
+ elif cmd == "help":
+ return self._show_help()
+ else:
+ return f"Unknown command: {cmd}. Use '%ballista help' for available commands."
+
+ @line_magic
+ def register(self, line: str) -> Optional[str]:
+ """Register a new table"""
+ if not line:
+ return "You should provide file extension and table name to register"
+ elif self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ else:
+ args = line.strip().split()
+ file_type = args[0]
+ if len(args) < 2:
+ return f"You should provide table name for this .{file_type} file"
+ table_name = args[1]
+ if len(args) < 3:
+ return f"You should provide path to your file"
+ file_name = args[2]
+
+ if file_type == "parquet":
+ self._ctx.register_parquet(table_name, file_name)
+ elif file_type == "csv":
+ self._ctx.register_csv(table_name, file_name)
+ else:
+ raise NotImplemented("Currently not supporting the inserted file format")
+
+ @line_cell_magic
+ def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]:
+ """
+ Execute a SQL query (both line and cell magic).
+
+ Two cases possible: with cell or without cell
+
+ Examples:
+ 1. Without a cell (line_magic)
+ %sql SELECT * FROM test_table
+ 2. With a cell (cell_magic)
+ %%sql --no-display
+ SELECT
+ id,
+ bool_col,
+ tinyint_col
+ FROM test_data_v1
+ WHERE id > 2
+ ORDER BY id
+ LIMIT 5
+ """
+ if not cell:
+ return self._execute_sql(line.strip()) if line.strip() else None
+ else:
+ args = line.strip().split()
+ display_results = True
+ limit = 50
+ var_name = None
+
+ i = 0
+ while i < len(args):
+ if args[i] == "--no-display":
+ display_results = False
+ elif args[i] == "--limit" and i + 1 < len(args):
+ try:
+ limit = int(args[i + 1])
+ i += 1
+ except ValueError:
+ pass
+ elif not args[i].startswith("--"):
+ var_name = args[i]
+ i += 1
+
+ query = cell.strip()
+ if not query:
+ return None
+
+ result = self._execute_sql(query, display_results=display_results, limit=limit)
+
+ # Store in user namespace if variable name provided
+ if var_name and self.shell is not None:
+ self.shell.user_ns[var_name] = result
+
+ return result
+
+ def _connect(self, address: str) -> Optional[str]:
+ """Connect to a Ballista cluster."""
+ if not address:
+ return "Usage: %ballista connect df://host:port"
+
+ # Normalize address
+ if not address.startswith("df://"):
+ address = f"df://{address}"
+
+ try:
+ self._ctx = BallistaSessionContext(address)
+ self._address = address
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Connected to Ballista cluster at {address}"))
+ else:
+ return f"✓ Connected to Ballista cluster at {address}"
+ except Exception as e:
+ self._ctx = None
+ self._address = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✗ Failed to connect to {address}: {e}"))
+ else:
+ return f"✗ Failed to connect to {address}: {e}"
+
+ def _disconnect(self) -> Optional[str]:
+ """Disconnect from the Ballista cluster."""
+ if not self.is_connected:
+ return "Not connected to any cluster."
+
+ address = self._address
+ self._ctx = None
+ self._address = None
+ self._last_result = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Disconnected from {address}"))
+ else:
+ return f"✓ Disconnected from {address}"
+
+ def _status(self) -> Optional[str]:
+ """Show connection status."""
+ if not self.is_connected:
+ return "Status: Not connected\n\nUse '%ballista connect df://host:port' to connect."
+
+ status_lines = [
+ f"Status: Connected",
+ f"Address: {self._address}",
+ f"Session ID: {self._ctx.session_id}",
+ f"Queries executed: {len(self._query_history)}",
+ ]
+
+ if self._last_result is not None:
+ status_lines.append("Last result: Available (access via '_' or '_last_result')")
+
+ def _format_html_status_output(line: str) -> str:
+ name, value = line.split(":", 1)
+ return f"<p><strong>{name}:</strong> {value.strip()}</p>"
+
+ html = "".join(_format_html_status_output(line) for line in status_lines)
+ if IPYTHON_AVAILABLE:
+ display(HTML(html))
+ else:
+ print("\n".join(status_lines))
+
+ def _tables(self) -> Optional[str]:
+ """List all registered tables."""
+ try:
+ # Get table names from the catalog and their respective schemas
+ tables = self.ctx.get_tables()
+ 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())
+ # Build a nice table display (HTML-formatted if applicable)
+ lines = [
+ {"content": f"Total: {table_count} table(s) in {schema_count} schema(s)", "is_info": True},
+ {"content": "Registered tables:", "is_info": True},
+ *[{"content": f"Schema: {schema_name}. Tables: {", ".join(list(table_names))}", "is_info": False}
+ for schema_name, table_names in tables.items()]
+ ]
+
+ def _format_html_tables_output(line: str, is_info: bool = False) -> str:
+ if is_info:
+ return f"<p><strong>{line}<strong></p>"
Review Comment:
Fixed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feature/jupyter notebook support [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on code in PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513#discussion_r2946185195
##########
python/python/ballista/jupyter.py:
##########
@@ -0,0 +1,493 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Ballista Jupyter Magic Commands.
+
+This module provides IPython magic commands for interacting with Ballista
+clusters directly from Jupyter notebooks.
+
+Usage:
+ %load_ext ballista.jupyter
+
+ # Connect to a Ballista cluster
+ %ballista connect df://localhost:50050
+
+ # Check connection status
+ %ballista status
+
+ # List registered tables
+ %ballista tables
+
+ # Show schema for a table
+ %ballista schema my_table
+
+ # Execute a simple SQL query (line magic)
+ %sql SELECT COUNT(*) FROM my_table
+
+ # Execute a complex SQL query (cell magic)
+ %%sql
+ SELECT
+ customer_id,
+ SUM(amount) as total
+ FROM orders
+ GROUP BY customer_id
+ ORDER BY total DESC
+ LIMIT 10
+"""
+
+from typing import Optional, List, Dict, Any
+import time
+import threading
+
+try:
+ from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic
+ from IPython.display import display, HTML
+
+ IPYTHON_AVAILABLE = True
+except ImportError:
+ IPYTHON_AVAILABLE = False
+
+ # Provide stub classes for when IPython is not available
+ class Magics:
+ def __init__(self, shell=None):
+ self.shell = shell
+
+ def magics_class(cls):
+ return cls
+
+ def line_magic(name_or_func):
+ """Stub line_magic decorator for when IPython is not available."""
+ # Handle both @line_magic and @line_magic("name") usage
+ if callable(name_or_func):
+ # Used as @line_magic without arguments
+ return name_or_func
+ else:
+ # Used as @line_magic("name") with arguments
+ def decorator(func):
+ return func
+ return decorator
+
+ def cell_magic(name_or_func):
+ """Stub cell_magic decorator for when IPython is not available."""
+ # Handle both @cell_magic and @cell_magic("name") usage
+ if callable(name_or_func):
+ return name_or_func
+ else:
+ def decorator(func):
+ return func
+ return decorator
+
+
+from .extension import BallistaSessionContext, DistributedDataFrame
+
+
+class BallistaConnectionError(Exception):
+ """Raised when not connected to a Ballista cluster."""
+
+ pass
+
+
+@magics_class
+class BallistaMagics(Magics):
+ """
+ IPython magic commands for Ballista.
+
+ Provides convenient commands for connecting to Ballista clusters,
+ executing SQL queries, and exploring table schemas.
+ """
+
+ def __init__(self, shell=None):
+ super().__init__(shell)
+ self._ctx: Optional[BallistaSessionContext] = None
+ self._address: Optional[str] = None
+ self._last_result: Optional[DistributedDataFrame] = None
+ self._query_history: List[Dict[str, Any]] = []
+
+ @property
+ def ctx(self) -> BallistaSessionContext:
+ """Get the current context, raising an error if not connected."""
+ if self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ return self._ctx
+
+ @property
+ def is_connected(self) -> bool:
+ """Check if connected to a Ballista cluster."""
+ return self._ctx is not None
+
+ @line_magic
+ def ballista(self, line: str) -> Optional[str]:
+ """
+ Ballista management commands.
+
+ Usage:
+ %ballista connect df://localhost:50050 - Connect to cluster
+ %ballista status - Show connection status
+ %ballista tables - List registered tables
+ %ballista schema <table> - Show table schema
+ %ballista disconnect - Disconnect from cluster
+ %ballista history - Show query history
+
+ Examples:
+ %ballista connect df://localhost:50050
+ %ballista tables
+ %ballista schema orders
+ """
+ parts = line.strip().split(maxsplit=1)
+ if not parts:
+ return self._show_help()
+
+ cmd = parts[0].lower()
+ args = parts[1] if len(parts) > 1 else ""
+
+ if cmd == "connect":
+ return self._connect(args)
+ elif cmd == "status":
+ return self._status()
+ elif cmd == "tables":
+ return self._tables()
+ elif cmd == "schema":
+ return self._schema(args)
+ elif cmd == "disconnect":
+ return self._disconnect()
+ elif cmd == "history":
+ return self._show_history()
+ elif cmd == "help":
+ return self._show_help()
+ else:
+ return f"Unknown command: {cmd}. Use '%ballista help' for available commands."
+
+ @line_magic
+ def register(self, line: str) -> Optional[str]:
+ """Register a new table"""
+ if not line:
+ return "You should provide file extension and table name to register"
+ elif self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ else:
+ args = line.strip().split()
+ file_type = args[0]
+ if len(args) < 2:
+ return f"You should provide table name for this .{file_type} file"
+ table_name = args[1]
+ if len(args) < 3:
+ return f"You should provide path to your file"
+ file_name = args[2]
+
+ if file_type == "parquet":
+ self._ctx.register_parquet(table_name, file_name)
+ elif file_type == "csv":
+ self._ctx.register_csv(table_name, file_name)
+ else:
+ raise NotImplemented("Currently not supporting the inserted file format")
+
+ @line_cell_magic
+ def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]:
+ """
+ Execute a SQL query (both line and cell magic).
+
+ Two cases possible: with cell or without cell
+
+ Examples:
+ 1. Without a cell (line_magic)
+ %sql SELECT * FROM test_table
+ 2. With a cell (cell_magic)
+ %%sql --no-display
+ SELECT
+ id,
+ bool_col,
+ tinyint_col
+ FROM test_data_v1
+ WHERE id > 2
+ ORDER BY id
+ LIMIT 5
+ """
+ if not cell:
+ return self._execute_sql(line.strip()) if line.strip() else None
+ else:
+ args = line.strip().split()
+ display_results = True
+ limit = 50
+ var_name = None
+
+ i = 0
+ while i < len(args):
+ if args[i] == "--no-display":
+ display_results = False
+ elif args[i] == "--limit" and i + 1 < len(args):
+ try:
+ limit = int(args[i + 1])
+ i += 1
+ except ValueError:
+ pass
+ elif not args[i].startswith("--"):
+ var_name = args[i]
+ i += 1
+
+ query = cell.strip()
+ if not query:
+ return None
+
+ result = self._execute_sql(query, display_results=display_results, limit=limit)
+
+ # Store in user namespace if variable name provided
+ if var_name and self.shell is not None:
+ self.shell.user_ns[var_name] = result
+
+ return result
+
+ def _connect(self, address: str) -> Optional[str]:
+ """Connect to a Ballista cluster."""
+ if not address:
+ return "Usage: %ballista connect df://host:port"
+
+ # Normalize address
+ if not address.startswith("df://"):
+ address = f"df://{address}"
+
+ try:
+ self._ctx = BallistaSessionContext(address)
+ self._address = address
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Connected to Ballista cluster at {address}"))
+ else:
+ return f"✓ Connected to Ballista cluster at {address}"
+ except Exception as e:
+ self._ctx = None
+ self._address = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✗ Failed to connect to {address}: {e}"))
+ else:
+ return f"✗ Failed to connect to {address}: {e}"
+
+ def _disconnect(self) -> Optional[str]:
+ """Disconnect from the Ballista cluster."""
+ if not self.is_connected:
+ return "Not connected to any cluster."
+
+ address = self._address
+ self._ctx = None
+ self._address = None
+ self._last_result = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Disconnected from {address}"))
+ else:
+ return f"✓ Disconnected from {address}"
+
+ def _status(self) -> Optional[str]:
+ """Show connection status."""
+ if not self.is_connected:
+ return "Status: Not connected\n\nUse '%ballista connect df://host:port' to connect."
+
+ status_lines = [
+ f"Status: Connected",
+ f"Address: {self._address}",
+ f"Session ID: {self._ctx.session_id}",
+ f"Queries executed: {len(self._query_history)}",
+ ]
+
+ if self._last_result is not None:
+ status_lines.append("Last result: Available (access via '_' or '_last_result')")
+
+ def _format_html_status_output(line: str) -> str:
+ name, value = line.split(":", 1)
+ return f"<p><strong>{name}:</strong> {value.strip()}</p>"
+
+ html = "".join(_format_html_status_output(line) for line in status_lines)
+ if IPYTHON_AVAILABLE:
+ display(HTML(html))
+ else:
+ print("\n".join(status_lines))
+
+ def _tables(self) -> Optional[str]:
+ """List all registered tables."""
+ try:
+ # Get table names from the catalog and their respective schemas
+ tables = self.ctx.get_tables()
+ 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())
+ # Build a nice table display (HTML-formatted if applicable)
+ lines = [
+ {"content": f"Total: {table_count} table(s) in {schema_count} schema(s)", "is_info": True},
+ {"content": "Registered tables:", "is_info": True},
+ *[{"content": f"Schema: {schema_name}. Tables: {", ".join(list(table_names))}", "is_info": False}
+ for schema_name, table_names in tables.items()]
+ ]
+
+ def _format_html_tables_output(line: str, is_info: bool = False) -> str:
+ if is_info:
+ return f"<p><strong>{line}<strong></p>"
+ else:
+ return f"<p><pre><i>{line}</i></pre></p>"
+
+ if IPYTHON_AVAILABLE:
+ display(
+ HTML("".join(_format_html_tables_output(val["content"], val["is_info"]) for val in lines))
+ )
+ else:
+ return "".join(val["content"] for val in lines)
+ except Exception as e:
+ return f"Error listing tables: {e}"
+
+ def _schema(self, table_name: str) -> Optional[str]:
+ """Show schema for a table."""
+ if not table_name:
+ return "Usage: %ballista schema <table_name>"
+
+ try:
+ # Query the table with LIMIT 0 to get schema without data
+ df = self.ctx.sql(f"SELECT * FROM {table_name} LIMIT 0")
+ schema = df.schema()
+
+ lines = [f"Schema for '{table_name}':", "-" * 50]
+ for field in schema:
+ nullable = "NULL" if field.nullable else "NOT NULL"
+ lines.append(f" {field.name:20} {str(field.type):15} {nullable}")
+ lines.append("-" * 50)
+ lines.append(f"Total: {len(schema)} column(s)")
+
+ return "\n".join(lines)
+ except Exception as e:
+ return f"Error getting schema for '{table_name}': {e}"
+
+ def _execute_sql(
+ self,
+ query: str,
+ display_results: bool = True,
+ limit: int = 50,
+ ) -> Optional[DistributedDataFrame]:
+ """Execute a SQL query and return the result."""
+ start_time = time.time()
+
+ try:
+ result = self.ctx.sql(query)
+ elapsed = time.time() - start_time
+
+ # Store result
+ self._last_result = result
+ if self.shell is not None and hasattr(self.shell, 'user_ns'):
+ self.shell.user_ns["_last_result"] = result
+
+ # Record in history
+ self._query_history.append(
+ {
+ "query": query,
+ "elapsed_seconds": elapsed,
+ "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
+ }
+ )
+
+ # Display if requested and in notebook environment
+ return result
+
+ except Exception as e:
+ error_msg = f"Query failed: {e}"
+ if IPYTHON_AVAILABLE:
+ try:
+ display(HTML(f'<div style="color: red; font-weight: bold;">{error_msg}</div>'))
+ except Exception:
+ print(error_msg)
+ else:
+ print(error_msg)
+ return None
+
+ def _show_history(self) -> Optional[str]:
+ """Show query history."""
+ if not self._query_history:
+ return "No queries executed yet."
+
+ lines = ["Query History:", "-" * 60]
+ for i, entry in enumerate(self._query_history[-10:], 1): # Last 10 queries
+ query_preview = entry["query"][:50] + "..." if len(entry["query"]) > 50 else entry["query"]
+ query_preview = query_preview.replace("\n", " ")
+ lines.append(f"{i}. [{entry['timestamp']}] ({entry['elapsed_seconds']:.2f}s)")
+ lines.append(f" {query_preview}")
+ lines.append("-" * 60)
+
+ for row in lines:
+ print(row)
+
+ def _show_help(self) -> Optional[str]:
+ """Show help for Ballista magic commands."""
+ help_info = """
+Ballista Jupyter Magic Commands
+================================
+
+Connection:
+ %ballista connect <url> - Connect to Ballista cluster
+ %ballista disconnect - Disconnect from cluster
+ %ballista status - Show connection status
+
+Exploration:
+ %ballista tables - List registered tables
+ %ballista schema <table> - Show table schema
+
+Table-register:
+ %register [format] [schema.table_name] [file_path] - Register a new table in the current Ballista Context
+
+Query:
+ %sql <query> - Execute single-line SQL query
+
+ %%sql [options] [var] - Execute multi-line SQL query
+ Options:
+ --no-display - Don't display results
+ --limit N - Limit displayed rows (default: 50)
+ var - Store result in variable
+
+History:
+ %ballista history - Show recent query history
+
+Examples:
+ %ballista connect df://localhost:50050
+ %ballista tables
+ %ballista schema orders
+
+ %sql SELECT COUNT(*) FROM orders
+
+ %%sql my_result
+ SELECT customer_id, SUM(amount) as total
+ FROM orders
+ GROUP BY customer_id
+ ORDER BY total DESC
+ LIMIT 10
+ """
+ for row in help_info.split("\n"):
+ print(row)
Review Comment:
Fixed
##########
python/README.md:
##########
@@ -79,6 +79,94 @@ pyarrow_batches = df.collect()
Check [DataFusion python](https://datafusion.apache.org/python/) provides more examples and manuals.
+## Jupyter Notebook Support
+
+PyBallista provides first-class Jupyter notebook support with SQL magic commands and rich HTML rendering.
+
Review Comment:
Fixed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: jupyter notebook support [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513#issuecomment-4212735251
I think it's ready for merge. any additional comments? @martin-g
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feature/jupyter notebook support [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513#issuecomment-4157153267
Fixed the pytest functions and ran through the ruff
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feature/jupyter notebook support [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on code in PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513#discussion_r2946182318
##########
python/python/ballista/jupyter.py:
##########
@@ -0,0 +1,493 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Ballista Jupyter Magic Commands.
+
+This module provides IPython magic commands for interacting with Ballista
+clusters directly from Jupyter notebooks.
+
+Usage:
+ %load_ext ballista.jupyter
+
+ # Connect to a Ballista cluster
+ %ballista connect df://localhost:50050
+
+ # Check connection status
+ %ballista status
+
+ # List registered tables
+ %ballista tables
+
+ # Show schema for a table
+ %ballista schema my_table
+
+ # Execute a simple SQL query (line magic)
+ %sql SELECT COUNT(*) FROM my_table
+
+ # Execute a complex SQL query (cell magic)
+ %%sql
+ SELECT
+ customer_id,
+ SUM(amount) as total
+ FROM orders
+ GROUP BY customer_id
+ ORDER BY total DESC
+ LIMIT 10
+"""
+
+from typing import Optional, List, Dict, Any
+import time
+import threading
+
+try:
+ from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic
+ from IPython.display import display, HTML
+
+ IPYTHON_AVAILABLE = True
+except ImportError:
+ IPYTHON_AVAILABLE = False
+
+ # Provide stub classes for when IPython is not available
+ class Magics:
+ def __init__(self, shell=None):
+ self.shell = shell
+
+ def magics_class(cls):
+ return cls
+
+ def line_magic(name_or_func):
+ """Stub line_magic decorator for when IPython is not available."""
+ # Handle both @line_magic and @line_magic("name") usage
+ if callable(name_or_func):
+ # Used as @line_magic without arguments
+ return name_or_func
+ else:
+ # Used as @line_magic("name") with arguments
+ def decorator(func):
+ return func
+ return decorator
+
+ def cell_magic(name_or_func):
+ """Stub cell_magic decorator for when IPython is not available."""
+ # Handle both @cell_magic and @cell_magic("name") usage
+ if callable(name_or_func):
+ return name_or_func
+ else:
+ def decorator(func):
+ return func
+ return decorator
+
+
+from .extension import BallistaSessionContext, DistributedDataFrame
+
+
+class BallistaConnectionError(Exception):
+ """Raised when not connected to a Ballista cluster."""
+
+ pass
+
+
+@magics_class
+class BallistaMagics(Magics):
+ """
+ IPython magic commands for Ballista.
+
+ Provides convenient commands for connecting to Ballista clusters,
+ executing SQL queries, and exploring table schemas.
+ """
+
+ def __init__(self, shell=None):
+ super().__init__(shell)
+ self._ctx: Optional[BallistaSessionContext] = None
+ self._address: Optional[str] = None
+ self._last_result: Optional[DistributedDataFrame] = None
+ self._query_history: List[Dict[str, Any]] = []
+
+ @property
+ def ctx(self) -> BallistaSessionContext:
+ """Get the current context, raising an error if not connected."""
+ if self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ return self._ctx
+
+ @property
+ def is_connected(self) -> bool:
+ """Check if connected to a Ballista cluster."""
+ return self._ctx is not None
+
+ @line_magic
+ def ballista(self, line: str) -> Optional[str]:
+ """
+ Ballista management commands.
+
+ Usage:
+ %ballista connect df://localhost:50050 - Connect to cluster
+ %ballista status - Show connection status
+ %ballista tables - List registered tables
+ %ballista schema <table> - Show table schema
+ %ballista disconnect - Disconnect from cluster
+ %ballista history - Show query history
+
+ Examples:
+ %ballista connect df://localhost:50050
+ %ballista tables
+ %ballista schema orders
+ """
+ parts = line.strip().split(maxsplit=1)
+ if not parts:
+ return self._show_help()
+
+ cmd = parts[0].lower()
+ args = parts[1] if len(parts) > 1 else ""
+
+ if cmd == "connect":
+ return self._connect(args)
+ elif cmd == "status":
+ return self._status()
+ elif cmd == "tables":
+ return self._tables()
+ elif cmd == "schema":
+ return self._schema(args)
+ elif cmd == "disconnect":
+ return self._disconnect()
+ elif cmd == "history":
+ return self._show_history()
+ elif cmd == "help":
+ return self._show_help()
+ else:
+ return f"Unknown command: {cmd}. Use '%ballista help' for available commands."
+
+ @line_magic
+ def register(self, line: str) -> Optional[str]:
+ """Register a new table"""
+ if not line:
+ return "You should provide file extension and table name to register"
+ elif self._ctx is None:
+ raise BallistaConnectionError(
+ "Not connected to a Ballista cluster. "
+ "Use: %ballista connect df://host:port"
+ )
+ else:
+ args = line.strip().split()
+ file_type = args[0]
+ if len(args) < 2:
+ return f"You should provide table name for this .{file_type} file"
+ table_name = args[1]
+ if len(args) < 3:
+ return f"You should provide path to your file"
+ file_name = args[2]
+
+ if file_type == "parquet":
+ self._ctx.register_parquet(table_name, file_name)
+ elif file_type == "csv":
+ self._ctx.register_csv(table_name, file_name)
+ else:
+ raise NotImplemented("Currently not supporting the inserted file format")
+
+ @line_cell_magic
+ def sql(self, line: str, cell=None) -> Optional[DistributedDataFrame]:
+ """
+ Execute a SQL query (both line and cell magic).
+
+ Two cases possible: with cell or without cell
+
+ Examples:
+ 1. Without a cell (line_magic)
+ %sql SELECT * FROM test_table
+ 2. With a cell (cell_magic)
+ %%sql --no-display
+ SELECT
+ id,
+ bool_col,
+ tinyint_col
+ FROM test_data_v1
+ WHERE id > 2
+ ORDER BY id
+ LIMIT 5
+ """
+ if not cell:
+ return self._execute_sql(line.strip()) if line.strip() else None
+ else:
+ args = line.strip().split()
+ display_results = True
+ limit = 50
+ var_name = None
+
+ i = 0
+ while i < len(args):
+ if args[i] == "--no-display":
+ display_results = False
+ elif args[i] == "--limit" and i + 1 < len(args):
+ try:
+ limit = int(args[i + 1])
+ i += 1
+ except ValueError:
+ pass
+ elif not args[i].startswith("--"):
+ var_name = args[i]
+ i += 1
+
+ query = cell.strip()
+ if not query:
+ return None
+
+ result = self._execute_sql(query, display_results=display_results, limit=limit)
+
+ # Store in user namespace if variable name provided
+ if var_name and self.shell is not None:
+ self.shell.user_ns[var_name] = result
+
+ return result
+
+ def _connect(self, address: str) -> Optional[str]:
+ """Connect to a Ballista cluster."""
+ if not address:
+ return "Usage: %ballista connect df://host:port"
+
+ # Normalize address
+ if not address.startswith("df://"):
+ address = f"df://{address}"
+
+ try:
+ self._ctx = BallistaSessionContext(address)
+ self._address = address
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Connected to Ballista cluster at {address}"))
+ else:
+ return f"✓ Connected to Ballista cluster at {address}"
+ except Exception as e:
+ self._ctx = None
+ self._address = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✗ Failed to connect to {address}: {e}"))
+ else:
+ return f"✗ Failed to connect to {address}: {e}"
+
+ def _disconnect(self) -> Optional[str]:
+ """Disconnect from the Ballista cluster."""
+ if not self.is_connected:
+ return "Not connected to any cluster."
+
+ address = self._address
+ self._ctx = None
+ self._address = None
+ self._last_result = None
+ if IPYTHON_AVAILABLE:
+ display(HTML(f"✓ Disconnected from {address}"))
+ else:
+ return f"✓ Disconnected from {address}"
+
+ def _status(self) -> Optional[str]:
+ """Show connection status."""
+ if not self.is_connected:
+ return "Status: Not connected\n\nUse '%ballista connect df://host:port' to connect."
+
+ status_lines = [
+ f"Status: Connected",
+ f"Address: {self._address}",
+ f"Session ID: {self._ctx.session_id}",
+ f"Queries executed: {len(self._query_history)}",
+ ]
+
+ if self._last_result is not None:
+ status_lines.append("Last result: Available (access via '_' or '_last_result')")
+
+ def _format_html_status_output(line: str) -> str:
+ name, value = line.split(":", 1)
+ return f"<p><strong>{name}:</strong> {value.strip()}</p>"
+
+ html = "".join(_format_html_status_output(line) for line in status_lines)
+ if IPYTHON_AVAILABLE:
+ display(HTML(html))
+ else:
+ print("\n".join(status_lines))
+
+ def _tables(self) -> Optional[str]:
+ """List all registered tables."""
+ try:
+ # Get table names from the catalog and their respective schemas
+ tables = self.ctx.get_tables()
+ 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())
+ # Build a nice table display (HTML-formatted if applicable)
+ lines = [
+ {"content": f"Total: {table_count} table(s) in {schema_count} schema(s)", "is_info": True},
+ {"content": "Registered tables:", "is_info": True},
+ *[{"content": f"Schema: {schema_name}. Tables: {", ".join(list(table_names))}", "is_info": False}
+ for schema_name, table_names in tables.items()]
+ ]
+
+ def _format_html_tables_output(line: str, is_info: bool = False) -> str:
+ if is_info:
+ return f"<p><strong>{line}<strong></p>"
+ else:
+ return f"<p><pre><i>{line}</i></pre></p>"
+
+ if IPYTHON_AVAILABLE:
+ display(
+ HTML("".join(_format_html_tables_output(val["content"], val["is_info"]) for val in lines))
+ )
+ else:
+ return "".join(val["content"] for val in lines)
+ except Exception as e:
+ return f"Error listing tables: {e}"
+
+ def _schema(self, table_name: str) -> Optional[str]:
+ """Show schema for a table."""
+ if not table_name:
+ return "Usage: %ballista schema <table_name>"
+
+ try:
+ # Query the table with LIMIT 0 to get schema without data
+ df = self.ctx.sql(f"SELECT * FROM {table_name} LIMIT 0")
Review Comment:
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.
##########
python/python/ballista/jupyter.py:
##########
@@ -0,0 +1,493 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Ballista Jupyter Magic Commands.
+
+This module provides IPython magic commands for interacting with Ballista
+clusters directly from Jupyter notebooks.
+
+Usage:
+ %load_ext ballista.jupyter
+
+ # Connect to a Ballista cluster
+ %ballista connect df://localhost:50050
+
+ # Check connection status
+ %ballista status
+
+ # List registered tables
+ %ballista tables
+
+ # Show schema for a table
+ %ballista schema my_table
+
+ # Execute a simple SQL query (line magic)
+ %sql SELECT COUNT(*) FROM my_table
+
+ # Execute a complex SQL query (cell magic)
+ %%sql
+ SELECT
+ customer_id,
+ SUM(amount) as total
+ FROM orders
+ GROUP BY customer_id
+ ORDER BY total DESC
+ LIMIT 10
+"""
+
+from typing import Optional, List, Dict, Any
+import time
+import threading
+
+try:
+ from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic
+ from IPython.display import display, HTML
+
+ IPYTHON_AVAILABLE = True
+except ImportError:
+ IPYTHON_AVAILABLE = False
+
+ # Provide stub classes for when IPython is not available
+ class Magics:
+ def __init__(self, shell=None):
+ self.shell = shell
+
+ def magics_class(cls):
+ return cls
+
+ def line_magic(name_or_func):
+ """Stub line_magic decorator for when IPython is not available."""
+ # Handle both @line_magic and @line_magic("name") usage
+ if callable(name_or_func):
+ # Used as @line_magic without arguments
+ return name_or_func
+ else:
+ # Used as @line_magic("name") with arguments
+ def decorator(func):
+ return func
+ return decorator
+
+ def cell_magic(name_or_func):
+ """Stub cell_magic decorator for when IPython is not available."""
+ # Handle both @cell_magic and @cell_magic("name") usage
+ if callable(name_or_func):
+ return name_or_func
+ else:
+ def decorator(func):
+ return func
+ return decorator
Review Comment:
Fixed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feature/jupyter notebook support [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on code in PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513#discussion_r2946186536
##########
python/python/ballista/extension.py:
##########
@@ -229,9 +231,319 @@ def write_parquet(
df = self._to_internal_df()
df.write_parquet(str(path), compression.value, compression_level)
+ def explain_visual(self, analyze: bool = False) -> "ExecutionPlanVisualization":
+ """
+ Generate a visual representation of the execution plan.
+
+ This method creates an SVG visualization of the query execution plan,
+ which can be displayed directly in Jupyter notebooks.
+
+ Args:
+ analyze: If True, includes runtime statistics from actual execution.
+
+ Returns:
+ ExecutionPlanVisualization: An object that renders as SVG in Jupyter.
+
+ Example:
+ >>> df = ctx.sql("SELECT * FROM orders WHERE amount > 100")
+ >>> df.explain_visual() # Displays SVG in notebook
+ >>> viz = df.explain_visual(analyze=True)
+ >>> viz.save("plan.svg") # Save to file
+ """
+ # Get the execution plan as a string representation
+ # Note: explain() prints but doesn't return a string, so we use logical_plan()
+ try:
+ plan = self.logical_plan()
+ plan_str = plan.display_indent()
+ except Exception:
+ # Fallback if logical_plan() fails
+ plan_str = "Unable to retrieve execution plan"
+ return ExecutionPlanVisualization(plan_str, analyze=analyze)
+
+ def collect_with_progress(
+ self,
+ callback: Optional[callable] = None,
+ poll_interval: float = 0.5,
+ ):
+ """
+ Collect results with progress indication.
+
+ For long-running queries, this method provides progress updates
+ through a callback function or displays a progress bar in Jupyter.
+
+ Args:
+ callback: Optional function to call with progress updates.
+ Signature: callback(status: str, progress: float)
+ poll_interval: How often to check progress (seconds).
+
+ Returns:
+ The collected result batches.
+
+ Example:
+ >>> def my_callback(status, progress):
+ ... print(f"{status}: {progress:.1%}")
+ >>> batches = df.collect_with_progress(callback=my_callback)
+ """
+ import threading
+ import time
+
+ result = [None]
+ error = [None]
+ done = threading.Event()
+
+ def execute():
+ try:
+ result[0] = self.collect()
+ except Exception as e:
+ error[0] = e
+ finally:
+ done.set()
+
+ thread = threading.Thread(target=execute)
+ thread.start()
+
+ # Check if we're in a Jupyter environment
+ try:
+ from IPython.display import display, clear_output
+ from IPython import get_ipython
+
+ in_jupyter = get_ipython() is not None
+ except (ImportError, AttributeError):
+ in_jupyter = False
+
+ start_time = time.time()
+
+ if in_jupyter and callback is None:
+ # Display a simple progress indicator
+ try:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ clear_output(wait=True)
+ print(f"⏳ Query executing... ({elapsed:.1f}s elapsed)")
+
+ clear_output(wait=True)
+ elapsed = time.time() - start_time
+ print(f"✓ Query completed in {elapsed:.1f}s")
+ except Exception:
+ pass # Ignore display errors
+ elif callback is not None:
+ while not done.wait(timeout=poll_interval):
+ elapsed = time.time() - start_time
+ callback(f"Executing ({elapsed:.1f}s)", -1.0) # -1 means indeterminate
+
+ elapsed = time.time() - start_time
+ callback(f"Completed in {elapsed:.1f}s", 1.0)
+ else:
+ done.wait()
+
+ thread.join()
+
+ if error[0] is not None:
+ raise error[0]
+
+ return result[0]
+
+
+class ExecutionPlanVisualization:
+ """
+ A wrapper for execution plan visualizations that can render as SVG in Jupyter.
+
+ This class takes the text representation of an execution plan and converts
+ it to a Graphviz DOT format, which is then rendered as SVG.
+ """
+
+ def __init__(self, plan_str: str, analyze: bool = False):
+ self.plan_str = plan_str
+ self.analyze = analyze
+ self._svg_cache: Optional[str] = None
+
+ def _parse_plan_to_dot(self) -> str:
+ """Convert the plan string to DOT format for Graphviz."""
+ lines = self.plan_str.strip().split("\n")
+
+ dot_lines = [
+ "digraph ExecutionPlan {",
+ ' rankdir=TB;',
+ ' node [shape=box, style="rounded,filled", fontname="Helvetica"];',
+ ' edge [fontname="Helvetica"];',
+ "",
+ ]
+
+ nodes = []
+ edges = []
+ node_id = 0
+ stack = [] # (indent_level, node_id)
+
+ for line in lines:
+ if not line.strip():
+ continue
+
+ # Calculate indent level
+ indent = len(line) - len(line.lstrip())
+ content = line.strip()
+
+ # Skip non-plan lines
+ if content.startswith("physical_plan") or content.startswith("logical_plan"):
+ continue
+
+ # Create a node for this plan element
+ current_id = node_id
+ node_id += 1
+
+ # Determine node color based on operation type
+ color = "#E3F2FD" # Default light blue
+ if "Scan" in content or "TableScan" in content:
+ color = "#E8F5E9" # Light green for scans
+ elif "Filter" in content:
+ color = "#FFF3E0" # Light orange for filters
+ elif "Aggregate" in content or "HashAggregate" in content:
+ color = "#F3E5F5" # Light purple for aggregations
+ elif "Join" in content:
+ color = "#FFEBEE" # Light red for joins
+ elif "Sort" in content:
+ color = "#E0F7FA" # Light cyan for sorts
+ elif "Projection" in content:
+ color = "#FFF8E1" # Light amber for projections
+
+ # Escape special characters for DOT format
+ label = content.replace('"', '\\"').replace("\n", "\\n")
+ if len(label) > 60:
+ # Wrap long labels
+ label = label[:57] + "..."
+
+ nodes.append(f' node{current_id} [label="{label}", fillcolor="{color}"];')
+
+ # Connect to parent based on indentation
+ while stack and stack[-1][0] >= indent:
+ stack.pop()
+
+ if stack:
+ parent_id = stack[-1][1]
+ edges.append(f" node{parent_id} -> node{current_id};")
+
+ stack.append((indent, current_id))
+
+ dot_lines.extend(nodes)
+ dot_lines.append("")
+ dot_lines.extend(edges)
+ dot_lines.append("}")
+
+ return "\n".join(dot_lines)
+
+ def to_dot(self) -> str:
+ """Get the DOT representation of the execution plan."""
+ return self._parse_plan_to_dot()
+
+ def to_svg(self) -> str:
+ """
+ Convert the execution plan to SVG format.
+
+ Requires graphviz to be installed. If graphviz is not available,
+ returns a simple HTML representation instead.
+ """
+ if self._svg_cache is not None:
+ return self._svg_cache
+
+ dot_source = self._parse_plan_to_dot()
+
+ try:
+ import subprocess
+
+ # Try to use graphviz's dot command
+ process = subprocess.run(
+ ["dot", "-Tsvg"],
+ input=dot_source.encode(),
+ capture_output=True,
+ timeout=30,
+ )
+
+ if process.returncode == 0:
+ self._svg_cache = process.stdout.decode()
+ return self._svg_cache
+ except (subprocess.SubprocessError, FileNotFoundError, subprocess.TimeoutExpired):
+ pass
+
+ # Fallback: return a pre-formatted HTML representation
+ escaped_plan = (
+ self.plan_str.replace("&", "&")
+ .replace("<", "<")
+ .replace(">", ">")
+ )
+ self._svg_cache = f"""
+ <div style="font-family: monospace; background: #f5f5f5; padding: 10px;
+ border-radius: 5px; overflow-x: auto;">
+ <div style="color: #666; margin-bottom: 5px;">
+ Execution Plan {'(with statistics)' if self.analyze else ''}
+ <br><small>Install graphviz for visual diagram: brew install graphviz</small>
+ </div>
+ <pre style="margin: 0;">{escaped_plan}</pre>
+ </div>
+ """
+ return self._svg_cache
+
+ def save(self, path: str) -> None:
+ """Save the visualization to a file (SVG or DOT format)."""
+ if path.endswith(".dot"):
+ content = self.to_dot()
+ else:
+ content = self.to_svg()
+
+ with open(path, "w") as f:
+ f.write(content)
+
+ def _repr_html_(self) -> str:
+ """HTML representation for Jupyter notebooks."""
+ return self.to_svg()
+
+ def _repr_svg_(self) -> str:
+ """SVG representation for Jupyter notebooks."""
+ svg = self.to_svg()
+ # Only return if it's actual SVG content
+ if svg.strip().startswith("<svg") or svg.strip().startswith("<?xml"):
+ return svg
+ return ""
+
+ def __repr__(self) -> str:
+ """String representation."""
+ return f"ExecutionPlanVisualization(analyze={self.analyze})\n{self.plan_str}"
+
class BallistaSessionContext(SessionContext, metaclass=RedefiningSessionContextMeta):
+ """
+ A session context for connecting to and querying a Ballista cluster.
+
+ This class extends DataFusion's SessionContext to work with distributed
+ Ballista clusters, automatically routing query execution to the cluster
+ while maintaining API compatibility with local DataFusion usage.
+
+ Example:
+ >>> from ballista import BallistaSessionContext
+ >>> ctx = BallistaSessionContext("df://localhost:50050")
+ >>> df = ctx.sql("SELECT * FROM my_table LIMIT 10")
+ >>> df.show()
+
+ For Jupyter notebook users:
+ >>> %load_ext ballista.jupyter
+ >>> %ballista connect df://localhost:50050
+ >>> %sql SELECT * FROM my_table
+ """
+
def __init__(self, address: str, config=None, runtime=None):
super().__init__(config, runtime)
self.address = address
self.session_id = self.session_id()
+
+ def get_tables(self) -> Optional[dict[str, str]]:
+ """Get tables and their respective schemas (in terms of database schema)."""
+ try:
+ catalog = self.catalog()
+ schema_names = list(catalog.schema_names())
+ if schema_names is not None:
Review Comment:
Fixed. Thank you
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feature/jupyter notebook support [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on code in PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513#discussion_r2946193437
##########
python/python/ballista/extension.py:
##########
@@ -229,9 +231,319 @@ def write_parquet(
df = self._to_internal_df()
df.write_parquet(str(path), compression.value, compression_level)
+ def explain_visual(self, analyze: bool = False) -> "ExecutionPlanVisualization":
+ """
+ Generate a visual representation of the execution plan.
+
+ This method creates an SVG visualization of the query execution plan,
+ which can be displayed directly in Jupyter notebooks.
+
+ Args:
+ analyze: If True, includes runtime statistics from actual execution.
+
+ Returns:
+ ExecutionPlanVisualization: An object that renders as SVG in Jupyter.
+
+ Example:
+ >>> df = ctx.sql("SELECT * FROM orders WHERE amount > 100")
+ >>> df.explain_visual() # Displays SVG in notebook
+ >>> viz = df.explain_visual(analyze=True)
+ >>> viz.save("plan.svg") # Save to file
+ """
+ # Get the execution plan as a string representation
+ # Note: explain() prints but doesn't return a string, so we use logical_plan()
+ try:
+ plan = self.logical_plan()
Review Comment:
Wasn't originally used as it works right now as intended. Might be prone to change in the future, I guess?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feature/jupyter notebook support [datafusion-ballista]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513#issuecomment-4172160728
thanks @sandugood will have a look asap
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: jupyter notebook support [datafusion-ballista]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm merged PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feature/jupyter notebook support [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on PR #1513:
URL: https://github.com/apache/datafusion-ballista/pull/1513#issuecomment-4209309180
thank you for your review @milenkovicm
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org