Repository navigation
feat: Automatically load interactive extras on import with opt-out controls - #199
Merged
Merged
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces Jupyter Notebook Extras to automatically load interactive conveniences like explore_dataframe(), %dpip, and %%sparksql when importing the package inside an IPython kernel. It includes configuration options to opt out of these extras, updates dependencies, and adds comprehensive tests. The review feedback highlights a potential memory leak in _ipython.py due to strong references to the IPython shell in _SHELL_STATES, recommending the use of weakref.WeakKeyDictionary. Additionally, a minor grammatical redundancy was pointed out in the README.md documentation.
dborowitz
force-pushed
the
colabsqlviz
branch
from
September 21, 2026 18:44
95f3ebd to
62a606f
Compare
medb
requested changes
Sep 23, 2026
dborowitz
changed the title
…ntrols
Automatically initialize interactive extras when importing
google.cloud.managed_spark_connect inside an IPython kernel:
- Inject colabsqlviz's explore_dataframe() into IPython user_ns
- Load the %dpip line magic extension (google.cloud.managed_spark_magics)
- Package interactive dependencies (google-colabsqlviz>=0.3.0, ipython>=8.0,
traitlets>=5.1) under the optional [interactive] extra
(pip install 'google-cloud-spark-connect[interactive]') so non-interactive
installs incur zero extra dependencies
Conditional runtime initialization:
- _ipython is only imported when traitlets is installed, and each extra
is initialized only if its dependency is installed (importlib.util.find_spec).
Uninstalled extras are skipped silently; only installed-but-broken extras
produce a warning.
- Because the active set of extras depends on what is installed in the
environment, print a single summary line listing whichever extras were newly
loaded:
Loaded interactive extras: explore_dataframe(), %dpip.
To disable: %config ManagedSparkConnect.enable_extras = False
Extras that were already present are not listed.
Add opt-out and runtime configuration controls:
- Environment variable: MANAGED_SPARK_CONNECT_ENABLE_EXTRAS=false
- IPython traitlet: ManagedSparkConnect.enable_extras = False (supports both
persistent file-based config and runtime toggling via %config, tracking and
only undoing changes that managed_spark_connect itself performed).
dborowitz
force-pushed
the
colabsqlviz
branch
from
October 14, 2026 15:37
096ffec to
f2c32d4
Compare
Contributor
Author
|
/gemini review |
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces interactive extras for Managed Spark Connect, enabling automatic setup of IPython conveniences such as explore_dataframe, %dpip, and %%sparksql magics, along with corresponding configuration options, documentation, and tests. The review feedback suggests improving memory management by using weak references (weakref) to store the IPython shell and shell states, thereby preventing potential memory leaks. Additionally, it is recommended to make the brace-matching logic in the Spark SQL magic more robust by directly checking string values instead of relying on fragile sqlparse token classifications.
Running Spark SQL from a notebook has so far meant installing the
third-party sparksql-magic package, which predates Spark Connect and offers
no safe way to pass Python values into a query. This adds a clean-room
%%sparksql to managed_spark_magics, written only against public pyspark
APIs, that works with any SparkSession, classic or Spark Connect.
%%sparksql [-s SESSION] [-l LIMIT] [-v VIEW] [-c] [-e] [-q]
[--explain [MODE]] [variable]
df = %sparksql [-s SESSION] [-v VIEW] [-c] [-e] [--explain [MODE]]
SELECT ...
The cell form displays the result of the last statement (and optionally
stores the DataFrame in variable); the line form is always --quiet and
returns the DataFrame directly. Options and the SparkSql config name match
the old package's, so existing %config SparkSql.limit lines keep working.
The default limit stays 20 rows.
- Parameters: :name binds a notebook variable as a typed SQL literal via
spark.sql(args=...); IDENTIFIER(:name) covers table and column names.
{name} references a Spark DataFrame to query as a table, without allowing
arbitrary Python expressions or non-DataFrame interpolation. Both are found
with sqlparse, so text inside strings and comments is left alone.
- Multiple statements: a cell is split with sqlparse and run in order, only
the last result is shown or stored. Every statement's parameters are
bound before any of them runs, so a missing variable fails the cell
without side effects. A cell starting with BEGIN is passed through whole
as a SQL script.
- Rendering: values are formatted by Spark with CAST(... AS STRING), so they
look the way Spark prints them rather than however they convert to Python
(toPandas() turns arrays with nulls into floats and shows binary as
b'...'). limit+1 rows are fetched to report truncation exactly. The HTML
table shows column types, right-aligns numbers and marks NULLs, with a
show()-style text/plain fallback. DDL shows OK.
- Errors: a PySparkException is shown as Spark's condition, message and
SQLSTATE, which statement failed, and a caret under the offending SQL in
the cell, without the Python and JVM stack traces. The caret is left out
when Spark's position is not into the cell's own text: after {name}
rewriting, or for an error inside a view.
%config SparkSql.full_traceback = True re-raises the original.
- Session: -s NAME, else the active session, else the only SparkSession in
the notebook.
The magic is structured as a self-contained IPython extension in
managed_spark_magics.sparksql (with its own SparkSqlMagics class and
load_ipython_extension / unload_ipython_extension hooks) rather than being
bundled onto ManagedSparkMagics. managed_spark_magics delegates to it, so
%load_ext google.cloud.managed_spark_magics continues to load both %dpip and
%%sparksql, while %load_ext google.cloud.managed_spark_magics.sparksql loads
only %%sparksql. This keeps the implementation decoupled from
managed_spark_connect so it can eventually be split out into a standalone
project, leaving a thin re-export shim behind.
sqlparse is a new dependency of the [interactive] extra. It and pyspark are
imported lazily, so loading the extension, and %dpip, needs neither.
sparksql-magic is dropped from the dev requirements, along with the README
and DEVELOPING sections about it.
Since managed_spark_magics now provides more than %dpip, the interactive
extras loader tracks every magic an extension adds. A same-named magic that
was already registered (e.g. %%sparksql from sparksql-magic) is restored
after loading and survives unloading, and the summary lists a line+cell
magic once, as %%sparksql.
tests/local runs the magic end to end against local classic and Spark
Connect sessions, and CI now runs it using the runner's preinstalled JDK 17.
The README had grown a long %%sparksql reference, while %dpip had no documentation beyond a one-line mention. Move each bundled magic's reference into its own page, docs/dpip.md and docs/sparksql.md, and leave the README with a short Bundled Magics section: one example per magic plus a link to its page. The links use absolute GitHub URLs, since the README is also the PyPI long description, where relative links don't resolve.
Upgrade google-colabsqlviz to >=0.4.0, which supports chart_type='table' on explore_dataframe(), and route %%sparksql's result display through it by default so results open as an interactive table with pagination and chart switching. - Add -d / --display (colabsqlviz, viz, or html) and %config SparkSql.display (defaulting to 'colabsqlviz') to choose between colabsqlviz and the built-in HTML/text table renderer. - When colabsqlviz is the default and google-colabsqlviz is not installed, fall back to the HTML renderer; if -d colabsqlviz / -d viz was explicitly passed on the cell, raise a UsageError pointing to the [interactive] extra. - Default SparkSql.limit to None so colabsqlviz paginates all rows unless capped via -l N or %config SparkSql.limit = N (which passes df.limit(N)), while the HTML renderer still defaults to 20 rows.
dborowitz
force-pushed
the
colabsqlviz
branch
from
October 14, 2026 17:47
f2c32d4 to
5e06bad
Compare
medb
approved these changes
Oct 15, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Automatically initialize interactive extras when importing
google.cloud.managed_spark_connectinside an IPython kernel:colabsqlviz'sexplore_dataframe()into IPythonuser_ns%dpipline magic (google.cloud.managed_spark_magics)%%sparksqlcell and%sparksqlline magic (google.cloud.managed_spark_magics.sparksql)google-colabsqlviz>=0.4.0,ipython>=8.0,sqlparse>=0.5.0,traitlets>=5.1) under the optional[interactive]extra (pip install 'google-cloud-spark-connect[interactive]') so non-interactive installs incur zero extra dependenciesBuilt-in
%%sparksqlMagicRather than relying on the third-party
sparksql-magicpackage,google.cloud.managed_spark_magicsnow bundles a clean-room, full-featured Spark SQL magic written against public PySpark APIs that works with anySparkSession(classic or Spark Connect)::namebinds Python variables as typed SQL literals viaspark.sql(args=...);IDENTIFIER(:name)binds catalog/table/column names;{df}queries DataFrame variables directly. Handled viasqlparseso literals inside strings and comments are left intact.google-colabsqlviz>=0.4.0(falling back to HTML/text tables if not installed). Configurable per-cell via-d/--displayor globally via%config SparkSql.display.BEGIN ... ENDcompound scripts.--explain [MODE], lazy caching (-c), and eager checkpointing (-e).%config SparkSql.full_traceback = Trueto re-raise).%sparksql SELECT ...returns the resultingDataFramedirectly.google.cloud.managed_spark_magics.sparksqlso it can be loaded independently (%load_ext google.cloud.managed_spark_magics.sparksql) or bundled viagoogle.cloud.managed_spark_magics. Full documentation is available indocs/sparksql.md.Conditional Runtime Initialization & Opt-out Controls
_ipythonis only imported whentraitletsis installed, and each extra is initialized only if its dependency is installed (importlib.util.find_spec). Uninstalled extras are skipped silently; only installed-but-broken extras produce a warning.Loaded interactive extras: explore_dataframe(), %dpip, %%sparksql. To disable: %config ManagedSparkConnect.enable_extras = False(Extras that were already loaded are not listed).
MANAGED_SPARK_CONNECT_ENABLE_EXTRAS=falseManagedSparkConnect.enable_extras = False(supports persistent file-based configuration and runtime toggling via%config, tracking and cleanly undoing changes made bymanaged_spark_connect).Dependency Footprint & Import Latency
Because interactive dependencies live behind the optional
[interactive]extra:google-cloud-spark-connect): 0 new dependencies and 0 runtime overhead outside IPython (find_spec("traitlets")short-circuits whentraitletsis absent, andget_ipython()returnsNoneoutside a kernel).[interactive]installed:All core data dependencies (
pyspark,pandas>=2.0.0,pyarrow>=10.0.1,protobuf>=4.24.0,packaging>=20.0) are already satisfied bypyspark[connect]andgoogle-api-core. Installing[interactive]causes 0 upgrades or downgrades to existing packages.sqlparseHas Zero Extra Transitive Dependencies:sqlparse(~44 KiB wheel, ~450 KiB installed with.pyc) is a lightweight, pure-Python parser with no dependencies of its own, replacing the externalsparksql-magicpackage.In local notebook kernels (e.g., VS Code or Jupyter, which require
ipykernel), theipythonstack is already present, so[interactive]adds 1.08–3.59 MiB of wheel downloads:.pyc, e.g.uv).pyc(pip)pip)[interactive])ipykernel+ipywidgets)google-colabsqlviz,sqlparse,anywidget,psygnal)ipykernelonly)ipywidgets,jupyterlab-widgets,widgetsnbextension)ipythonstack)(Note: Pure Python code across all 7 packages added in the minimal VS Code kernel scenario is < 1 MiB. The remaining extracted space is static frontend JS bundles/source maps in
widgetsnbextensionandanywidget[~11.5 MiB],psygnal's compiled mypyc.sobinary [~1.3 MiB], andpip's.pycbytecode cache duplicating embedded JS strings incolabsqlviz[~1.0 MiB].)Initializing all three extras inside IPython takes ~35 ms hot / ~45–56 ms cold page cache (~28 ms of which is importing
anywidget→ipywidgetsviagoogle.colabsqlviz).