You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ag...@apache.org on 2026/05/23 22:24:42 UTC
(datafusion-ballista) 02/04: fix(python): clear latent lint debt and gate wheel jobs on lint/tests (#1749)
This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch branch-53
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
commit 3667d97480b0507db5ad7dc203bc986ed07e7d3a
Author: Andy Grove <ag...@apache.org>
AuthorDate: Sat May 23 14:02:23 2026 -0600
fix(python): clear latent lint debt and gate wheel jobs on lint/tests (#1749)
---
.github/workflows/build.yml | 14 +++++--
python/pyproject.toml | 6 +++
python/python/ballista/extension.py | 53 +++++++++++++++-----------
python/python/ballista/jupyter.py | 75 +++++++++++++++++++++++++++----------
python/python/tests/test_context.py | 3 +-
5 files changed, 104 insertions(+), 47 deletions(-)
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 6f24716f0..c93668ce7 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -100,6 +100,7 @@ jobs:
uv run --no-project ruff format --check .
generate-license:
+ if: github.event_name != 'pull_request'
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6.0.2
@@ -114,7 +115,8 @@ jobs:
path: LICENSE.txt
build-python-mac-win:
- needs: [ generate-license ]
+ if: github.event_name != 'pull_request'
+ needs: [ generate-license, build, test-release ]
name: Mac/Win
runs-on: ${{ matrix.os }}
strategy:
@@ -178,7 +180,8 @@ jobs:
path: python/target/wheels/*
build-manylinux-x86_64:
- needs: [ generate-license ]
+ if: github.event_name != 'pull_request'
+ needs: [ generate-license, build, test-release ]
name: Manylinux x86_64
runs-on: ubuntu-latest
steps:
@@ -219,7 +222,8 @@ jobs:
path: python/target/wheels/*
build-manylinux-aarch64:
- needs: [ generate-license ]
+ if: github.event_name != 'pull_request'
+ needs: [ generate-license, build, test-release ]
name: Manylinux arm64
runs-on: ubuntu-latest
steps:
@@ -260,7 +264,8 @@ jobs:
path: python/target/wheels/*
build-sdist:
- needs: [ generate-license ]
+ if: github.event_name != 'pull_request'
+ needs: [ generate-license, build, test-release ]
name: Source distribution
runs-on: ubuntu-latest
steps:
@@ -305,6 +310,7 @@ jobs:
shell: bash
merge-build-artifacts:
+ if: github.event_name != 'pull_request'
runs-on: ubuntu-latest
needs:
- build-python-mac-win
diff --git a/python/pyproject.toml b/python/pyproject.toml
index 7c0ff7f62..a97b8a7a2 100644
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -59,6 +59,12 @@ repository = "https://github.com/apache/datafusion-ballista"
[tool.isort]
profile = "black"
+[tool.ruff.format]
+# Jupyter notebooks are committed with the JSON indentation that the Jupyter
+# UI writes; ruff format would rewrite that on every run, producing large
+# whitespace-only diffs. Lint (ruff check) still applies to notebook cells.
+exclude = ["*.ipynb"]
+
[tool.maturin]
python-source = "python"
module-name = "ballista._internal_ballista"
diff --git a/python/python/ballista/extension.py b/python/python/ballista/extension.py
index 519b1d94a..90a3651bc 100644
--- a/python/python/ballista/extension.py
+++ b/python/python/ballista/extension.py
@@ -20,16 +20,10 @@ from datafusion import (
SessionContext,
DataFrame,
ParquetWriterOptions,
- DataFrameWriteOptions
+ DataFrameWriteOptions,
)
from datafusion.dataframe import Compression
-from typing import (
- Dict,
- List,
- Union,
- Optional,
- Callable
-)
+from typing import Dict, List, Union, Optional, Callable
import warnings
from ._internal_ballista import create_ballista_data_frame
@@ -151,6 +145,7 @@ class DistributedDataFrame(DataFrame, metaclass=RedefiningDataFrameMeta):
self.address = address
self.session_id = session_id
self.cluster_config = cluster_config
+
#
# this will create a ballista dataframe, which has ballista
# session context, and ballista planner.
@@ -164,12 +159,19 @@ class DistributedDataFrame(DataFrame, metaclass=RedefiningDataFrameMeta):
self.cluster_config,
)
return df
-
- def write_json(self, path, write_options: Union[DataFrameWriteOptions, None] = None):
+
+ def write_json(
+ self, path, write_options: Union[DataFrameWriteOptions, None] = None
+ ):
df = self._to_internal_df()
df.write_json(path, None)
- def write_csv(self, path, with_header=False, write_options: Union[DataFrameWriteOptions, None] = None,):
+ def write_csv(
+ self,
+ path,
+ with_header=False,
+ write_options: Union[DataFrameWriteOptions, None] = None,
+ ):
df = self._to_internal_df()
df.write_csv(path, with_header, None)
@@ -221,10 +223,7 @@ class DistributedDataFrame(DataFrame, metaclass=RedefiningDataFrameMeta):
df = self._to_internal_df()
df.write_parquet_with_options(
- str(path),
- options_internal,
- column_specific_options_internal,
- None
+ str(path), options_internal, column_specific_options_internal, None
)
def write_parquet(
@@ -384,7 +383,7 @@ class ExecutionPlanVisualization:
dot_lines = [
"digraph ExecutionPlan {",
- ' rankdir=TB;',
+ " rankdir=TB;",
' node [shape=box, style="rounded,filled", fontname="Helvetica"];',
' edge [fontname="Helvetica"];',
"",
@@ -404,7 +403,9 @@ class ExecutionPlanVisualization:
content = line.strip()
# Skip non-plan lines
- if content.startswith("physical_plan") or content.startswith("logical_plan"):
+ if content.startswith("physical_plan") or content.startswith(
+ "logical_plan"
+ ):
continue
# Create a node for this plan element
@@ -432,7 +433,9 @@ class ExecutionPlanVisualization:
# Wrap long labels
label = label[:57] + "..."
- nodes.append(f' node{current_id} [label="{label}", fillcolor="{color}"];')
+ nodes.append(
+ f' node{current_id} [label="{label}", fillcolor="{color}"];'
+ )
# Connect to parent based on indentation
while stack and stack[-1][0] >= indent:
@@ -481,7 +484,11 @@ class ExecutionPlanVisualization:
if process.returncode == 0:
self._svg_cache = process.stdout.decode()
return self._svg_cache
- except (subprocess.SubprocessError, FileNotFoundError, subprocess.TimeoutExpired) as e:
+ except (
+ subprocess.SubprocessError,
+ FileNotFoundError,
+ subprocess.TimeoutExpired,
+ ) as e:
warnings.warn(f"Could not convert the execution plan to SVG format: {e}")
pass
@@ -495,7 +502,7 @@ class ExecutionPlanVisualization:
<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 ''}
+ 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>
@@ -599,9 +606,11 @@ class BallistaSessionContext(SessionContext, metaclass=RedefiningSessionContextM
if schema_names:
tables_info = {}
for schema_name in schema_names:
- tables_info[schema_name] = list(catalog.schema(name=schema_name).table_names())
+ tables_info[schema_name] = list(
+ catalog.schema(name=schema_name).table_names()
+ )
return tables_info
except (AttributeError, NotImplementedError) as e:
warnings.warn(f"Could not retrieve tables from catalog: {e}")
pass
- return {}
\ No newline at end of file
+ return {}
diff --git a/python/python/ballista/jupyter.py b/python/python/ballista/jupyter.py
index c6e2b83de..d1d8c38ed 100644
--- a/python/python/ballista/jupyter.py
+++ b/python/python/ballista/jupyter.py
@@ -80,6 +80,7 @@ except ImportError:
# Used as @line_magic("name") with arguments
def decorator(func):
return func
+
return decorator
def cell_magic(name_or_func):
@@ -88,18 +89,22 @@ except ImportError:
if callable(name_or_func):
return name_or_func
else:
+
def decorator(func):
return func
+
return decorator
-
+
def line_cell_magic(name_or_func):
"""Stub line_cell_magic decorator for when IPython is not available"""
# Handle both @line_cell_magic and @line_cell_magic("name") usage
if callable(name_or_func):
return name_or_func
else:
+
def decorator(func):
return func
+
return decorator
@@ -183,8 +188,10 @@ class BallistaMagics(Magics):
elif cmd == "help":
return self._show_help()
else:
- return f"Unknown command: {cmd}. Use '%ballista help' for available commands."
-
+ return (
+ f"Unknown command: {cmd}. Use '%ballista help' for available commands."
+ )
+
@line_magic
def register(self, line: str) -> Optional[str]:
"""Register a new table"""
@@ -210,11 +217,13 @@ class BallistaMagics(Magics):
elif file_type == "csv":
self._ctx.register_csv(table_name, file_name)
else:
- raise NotImplementedError("Currently not supporting the inserted file format")
+ raise NotImplementedError(
+ "Currently not supporting the inserted file format"
+ )
- @line_cell_magic
+ @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
@@ -231,7 +240,7 @@ class BallistaMagics(Magics):
FROM test_data_v1
WHERE id > 2
ORDER BY id
- LIMIT 5
+ LIMIT 5
`my_result` will store the result of the SQL-query
"""
@@ -242,7 +251,7 @@ class BallistaMagics(Magics):
query = cell.strip()
if not query:
return None
-
+
args = line.strip().split()
i = 0
while i < len(args):
@@ -293,7 +302,7 @@ class BallistaMagics(Magics):
if IPYTHON_AVAILABLE:
display(HTML(f"✓ Disconnected from {address}"))
else:
- print(f"✓ Disconnected from {address}")
+ print(f"✓ Disconnected from {address}")
def _status(self) -> Optional[str]:
"""Show connection status."""
@@ -308,7 +317,9 @@ class BallistaMagics(Magics):
]
if self._last_result is not None:
- status_lines.append("Last result: Available (access via '_' or '_last_result')")
+ status_lines.append(
+ "Last result: Available (access via '_' or '_last_result')"
+ )
def _format_html_status_output(line: str) -> str:
name, value = line.split(":", 1)
@@ -319,7 +330,7 @@ class BallistaMagics(Magics):
display(HTML(html))
else:
print("\n".join(status_lines))
-
+
def _tables(self) -> Optional[str]:
"""List all registered tables."""
try:
@@ -331,10 +342,18 @@ class BallistaMagics(Magics):
table_count = sum(len(v) for v in 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": 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(table_names)}", "is_info": False}
- for schema_name, table_names in tables.items()]
+ *[
+ {
+ "content": f"Schema: {schema_name}. Tables: {', '.join(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:
@@ -342,10 +361,15 @@ class BallistaMagics(Magics):
return f"<pre><strong>{line}</strong></pre>"
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))
+ HTML(
+ "".join(
+ _format_html_tables_output(val["content"], val["is_info"])
+ for val in lines
+ )
+ )
)
else:
print("".join(val["content"] for val in lines))
@@ -386,7 +410,7 @@ class BallistaMagics(Magics):
# Store result
self._last_result = result
- if self.shell is not None and hasattr(self.shell, 'user_ns'):
+ if self.shell is not None and hasattr(self.shell, "user_ns"):
self.shell.user_ns["_last_result"] = result
# Record in history
@@ -405,7 +429,11 @@ class BallistaMagics(Magics):
error_msg = f"Query failed: {e}"
if IPYTHON_AVAILABLE:
try:
- display(HTML(f'<div style="color: red; font-weight: bold;">{error_msg}</div>'))
+ display(
+ HTML(
+ f'<div style="color: red; font-weight: bold;">{error_msg}</div>'
+ )
+ )
except Exception:
print(error_msg)
else:
@@ -419,9 +447,15 @@ class BallistaMagics(Magics):
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 = (
+ 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"{i}. [{entry['timestamp']}] ({entry['elapsed_seconds']:.2f}s)"
+ )
lines.append(f" {query_preview}")
lines.append("-" * 60)
output = "\n".join(lines)
@@ -479,6 +513,7 @@ Examples:
else:
print(help_info)
+
def load_ipython_extension(ipython):
"""
Load the Ballista IPython extension.
diff --git a/python/python/tests/test_context.py b/python/python/tests/test_context.py
index 3c38c63eb..bb92eb24c 100644
--- a/python/python/tests/test_context.py
+++ b/python/python/tests/test_context.py
@@ -16,7 +16,7 @@
# under the License.
from ballista import BallistaSessionContext, setup_test_cluster
-from datafusion import col, lit, InsertOp, DataFrameWriteOptions
+from datafusion import col, lit
import pytest
import pyarrow as pa
@@ -115,6 +115,7 @@ def test_cluster_config_accepts_ballista_namespaced_keys():
assert df.cluster_config == overrides
assert len(df.collect()) == 1
+
def test_write_csv(ctx, tmp_path):
df = ctx.read_csv("testdata/test.csv", has_header=True)
out_dir = str(tmp_path / "out")
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org