First draft - Querying with SQL

polars/07-querying-with-sql.py

@@ -0,0 +1,328 @@
+# /// script
+# requires-python = ">=3.12"
+# dependencies = [
+#     "duckdb==1.4.3",
+#     "kagglehub==0.3.13",
+#     "polars==1.36.1",
+#     "pyarrow==22.0.0",
+#     "sqlalchemy==2.0.45",
+#     "sqlglot==28.3.0",
+# ]
+# ///
+
+import marimo
+
+__generated_with = "0.18.4"
+app = marimo.App(width="medium")
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""
+    ## SQL Features in Marimo and Polars
+
+    _By [etrotta](https://github.com/etrotta)_
+
+    For this Notebook, we'll be using a [hotel booking analytics](https://www.kaggle.com/datasets/alperenmyung/international-hotel-booking-analytics) dataset.
+
+    We will see many ways in which you can use SQL inside of marimo and how each feature interacts with polars, including:
+    - How to read data from a SQLite file (or any Database connection)
+    - What are SQL Cells in Marimo
+    - How to load an SQL query into a DataFrame
+    - How to query DataFrames using SQL
+    """)
+    return
+
+
+@app.cell
+def _(mo, reviews, sqlite_engine):
+    _df = mo.sql(
+        f"""
+        SELECT * FROM reviews LIMIT 100
+        """,
+        engine=sqlite_engine,
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""
+    We will start by using `kagglehub` to download a `.sqlite` file, then create an `SQLAlchemy` engine to let marimo know about the database.
+    """)
+    return
+
+
+@app.cell
+def _(kagglehub):
+    dataset_id = "alperenmyung/international-hotel-booking-analytics"
+    cached_file = kagglehub.dataset_download(dataset_id, "booking_db.sqlite")
+    return (cached_file,)
+
+
+@app.cell
+def _(cached_file):
+    cached_file
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""
+    ### Using Marimo's SQL Cells
+    """)
+    return
+
+
+@app.cell
+def _(cached_file, sqlalchemy):
+    sqlite_engine = sqlalchemy.create_engine("sqlite:///" + cached_file)
+    return (sqlite_engine,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""
+    After creating the Engine, you should be able to see it in the **Data Sources** pannel in the sidebar. Whenever you create an SQLAlchemy engine as a global variable, Marimo picks up on it and makes it available for use in SQL Cells
+
+    You can use it to consult all tables and their columns, as well as click "Add table to notebook" to get the code to use it, creating our first SQL Cell:
+    """)
+    return
+
+
+@app.cell
+def _(hotels, mo, sqlite_engine):
+    _df = mo.sql(
+        f"""
+        SELECT * FROM hotels LIMIT 10
+        """,
+        engine=sqlite_engine,
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""
+    The `Output variable:` can be used to save the output as a polars DataFrame you can access later
+
+    For example, fetching all scores then performing a group by in polars
+    """)
+    return
+
+
+@app.cell
+def _(mo, reviews, sqlite_engine, users):
+    polars_age_groups = mo.sql(
+        f"""
+        SELECT reviews.*, age_group FROM reviews JOIN users ON reviews.user_id = users.user_id LIMIT 1000
+        """,
+        engine=sqlite_engine,
+    )
+    return (polars_age_groups,)
+
+
+@app.cell
+def _(pl, polars_age_groups):
+    _mean_scores = pl.col("^score_.*$").mean()
+    _age_group_start = pl.col("age_group").str.slice(0, 2).cast(int)
+    polars_age_groups.group_by("age_group").agg(_mean_scores).sort(_age_group_start)
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""
+    Although you could also calculate it directly in SQL, this gives you the flexibility to use polars for operations that are harder to describe in SQL
+    """)
+    return
+
+
+@app.cell
+def _(mo, reviews, sqlite_engine, users):
+    _df = mo.sql(
+        f"""
+        SELECT age_group, AVG(reviews.score_overall) FROM reviews JOIN users ON reviews.user_id = users.user_id GROUP BY age_group
+        """,
+        engine=sqlite_engine,
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""
+    You can also use SQL Cells to query DataFrames via DuckDB, but remember to change the Engine from the SQLite engine into the DuckDB Memory engine when doing so
+    """)
+    return
+
+
+@app.cell
+def _(mo, polars_age_groups):
+    _df = mo.sql(
+        f"""
+        SELECT * FROM polars_age_groups LIMIT 10
+        """
+    )
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""
+    ### Using Polars directly
+    """)
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""
+    Polars also offers some methods to interact with databases and query DataFrames using SQL directly, which you can use inside or outside of marimo the same.
+    """)
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""
+    Reading data from Databases:
+    """)
+    return
+
+
+@app.cell
+def _(pl, sqlite_engine):
+    hotels = pl.read_database("SELECT * FROM hotels LIMIT 10", sqlite_engine)
+    hotels
+    return (hotels,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""
+    Querying DataFrames with SQL:
+    """)
+    return
+
+
+@app.cell
+def _(hotels):
+    hotels.sql("SELECT * from self ORDER BY cleanliness_base DESC LIMIT 5")
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""
+    ### Using DuckDB
+    """)
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""
+    While marimo's SQL Cells are very practical and polars's direct methods are about as portable as it gets using polars, you can also use other libraries that integrate with polars via Arrow tables or input plugins.
+
+    One example of such integrations is DuckDB, which can be used with polars's Lazy mode as of 1.4.0
+    """)
+    return
+
+
+@app.cell
+def _(cached_file, duckdb):
+    duckdb_conn = duckdb.connect(cached_file)
+    return (duckdb_conn,)
+
+
+@app.cell
+def _(duckdb_conn):
+    # Loading into a normal DataFrame:
+    duckdb_conn.sql("SELECT * FROM hotels LIMIT 10").pl()
+    return
+
+
+@app.cell
+def _(duckdb_conn):
+    # Loading into a LazyFrame:
+    duckdb_conn.sql("SELECT * FROM hotels").pl(lazy=True).limit(10).collect()
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""
+    Note that this is very similar to SQL cells backed by DuckDB, with the biggest difference being that you can control how the result is consumed as opposed to it always being loaded into memory.
+
+    Many features such as querying from DataFrames work the same using DuckDB directly as they do in DuckDB-backed SQL Cells, and vice-versa
+    """)
+    return
+
+
+@app.cell
+def _(duckdb, hotels):
+    duckdb.sql("SELECT * FROM hotels").pl(lazy=True).sort("cleanliness_base", descending=True).limit(5).collect()
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md(r"""
+    ### Utilities
+    """)
+    return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    delete_file_button = mo.ui.run_button(label="Delete cached file", kind="warn")
+    mo.vstack([mo.md("If you want to delete the downloaded file from your cache"), delete_file_button])
+    return (delete_file_button,)
+
+
+@app.cell
+def _(cached_file, delete_file_button, pathlib):
+    if delete_file_button.value:
+        pathlib.Path(cached_file).unlink()
+    return
+
+
+@app.cell
+def _():
+    import marimo as mo
+    return (mo,)
+
+
+@app.cell
+def _():
+    import polars as pl
+    return (pl,)
+
+
+@app.cell
+def _():
+    import duckdb
+    return (duckdb,)
+
+
+@app.cell
+def _():
+    import sqlalchemy
+    return (sqlalchemy,)
+
+
+@app.cell
+def _():
+    import kagglehub
+    return (kagglehub,)
+
+
+@app.cell
+def _():
+    import pathlib
+    return (pathlib,)
+
+
+if __name__ == "__main__":
+    app.run()