JupySQL and chDB

JupySQL is a Python library that lets you run SQL in Jupyter notebooks and the IPython shell. In this guide, we’re going to learn how to query data using chDB and JupySQL.

Setup

Let’s first create a virtual environment:

python -m venv .venv
source .venv/bin/activate

And then, we’ll install JupySQL, IPython, and Jupyter Lab:

pip install jupysql ipython jupyterlab

We can use JupySQL in IPython, which we can launch by running:

ipython

Or in Jupyter Lab, by running:

jupyter lab

If you’re using Jupyter Lab, you’ll need to create a notebook before following the rest of the guide.

Downloading a dataset

We’re going to use the New York City taxi dataset, which contains around 3 million taxi rides along with the fare, tip, and pickup neighborhood of each one. The trips are split across several TSV files, so let’s start by downloading those:

from urllib.request import urlretrieve

base = "https://datasets-documentation.s3.eu-west-3.amazonaws.com/nyc-taxi"
for n in range(3):
  _ = urlretrieve(
    f"{base}/trips_{n}.gz",
    f"trips_{n}.gz",
  )

Configuring chDB and JupySQL

Next, let’s import the dbapi module for chDB:

from chdb import dbapi

And we’ll create a chDB connection. Any data that we persist will be saved to the taxi.chdb directory:

conn = dbapi.connect(path="taxi.chdb")

Let’s now load the sql magic and create a connection to chDB:

%load_ext sql
%sql conn --alias chdb

Next, we’ll display the display limit so that results of queries won’t be truncated:

%config SqlMagic.displaylimit = None

Querying data in TSV files

We’ve downloaded a bunch of files with the trips_ prefix. Let’s use the DESCRIBE clause to understand the schema:

%%sql
DESCRIBE file('trips_*.gz')
SETTINGS describe_compact_output=1,
         schema_inference_make_columns_nullable=0

+--------------------+----------+
|        name        |   type   |
+--------------------+----------+
|      trip_id       |  Int64   |
|     vendor_id      |  Int64   |
|    pickup_date     |   Date   |
|  pickup_datetime   | DateTime |
|    dropoff_date    |   Date   |
|  dropoff_datetime  | DateTime |
| store_and_fwd_flag |  Int64   |
|    rate_code_id    |  Int64   |
+--------------------+----------+
(40 more rows)

We can also write a SELECT query directly against these files to see what the data looks like:

%%sql
SELECT trip_id, pickup_datetime, pickup_ntaname,
       trip_distance, fare_amount, tip_amount
FROM file('trips_*.gz')
LIMIT 3
SETTINGS schema_inference_make_columns_nullable=0

+------------+---------------------+----------------------------------------+---------------+-------------+------------+
|  trip_id   |   pickup_datetime   |             pickup_ntaname             | trip_distance | fare_amount | tip_amount |
+------------+---------------------+----------------------------------------+---------------+-------------+------------+
| 1199999902 | 2015-07-07 19:45:07 |      Lenox Hill-Roosevelt Island       |      2.59     |     14.5    |    3.26    |
| 1199999919 | 2015-07-07 20:26:29 |                Airport                 |      2.4      |      9      |     0      |
| 1199999944 | 2015-07-07 21:25:09 | SoHo-TriBeCa-Civic Center-Little Italy |      5.13     |      20     |     3      |
+------------+---------------------+----------------------------------------+---------------+-------------+------------+

If we look back at the schema, a few of the money-related columns — trip_distance, fare_amount, and tip_amount — were inferred as String rather than a numeric type. We’ll clean those up when we import the data into a table.

Importing TSV files into chDB

Now we’re going to store the data from these TSV files in a table. The default database doesn’t persist data on disk, so we need to create another database first:

%sql CREATE DATABASE taxi

And now we’re going to create a table called trips whose schema will be derived from the structure of the data in the TSV files. We’ll use the REPLACE clause to cast the money-related columns to Float64, and the transform function to turn the numeric pickup_borocode column into a human-readable borough name:

%%sql
CREATE TABLE taxi.trips
ENGINE = MergeTree
ORDER BY pickup_datetime AS
SELECT * REPLACE (
    toFloat64OrZero(trip_distance) AS trip_distance,
    toFloat64OrZero(fare_amount) AS fare_amount,
    toFloat64OrZero(tip_amount) AS tip_amount,
    toFloat64OrZero(total_amount) AS total_amount
  ),
  transform(pickup_borocode, [1, 2, 3, 4, 5],
            ['Manhattan', 'Bronx', 'Brooklyn', 'Queens', 'Staten Island'],
            'Unknown') AS pickup_borough
FROM file('trips_*.gz')
SETTINGS schema_inference_make_columns_nullable=0

Let’s do a quick check on the data in our table:

%sql SELECT count() AS trips FROM taxi.trips

+---------+
|  trips  |
+---------+
| 3000317 |
+---------+

Just over 3 million trips — let’s also bring in a second table. New York City’s Taxi & Limousine Commission divides the city into taxi zones, and a lookup file maps each zone to its borough. Let’s download that file:

_ = urlretrieve(
    f"{base}/taxi_zone_lookup.csv",
    "taxi_zone_lookup.csv",
)

And then create a table called zones based on the content of the CSV file:

%%sql
CREATE TABLE taxi.zones
ENGINE = MergeTree
ORDER BY LocationID AS
SELECT * FROM file('taxi_zone_lookup.csv')
SETTINGS schema_inference_make_columns_nullable=0

Once that’s finished running, we can have a look at the data we’ve ingested:

%sql SELECT * FROM taxi.zones LIMIT 5

+------------+---------------+-------------------------+--------------+
| LocationID |    Borough    |           Zone          | service_zone |
+------------+---------------+-------------------------+--------------+
|     1      |      EWR      |      Newark Airport     |     EWR      |
|     2      |     Queens    |       Jamaica Bay       |  Boro Zone   |
|     3      |     Bronx     | Allerton/Pelham Gardens |  Boro Zone   |
|     4      |   Manhattan   |      Alphabet City      | Yellow Zone  |
|     5      | Staten Island |      Arden Heights      |  Boro Zone   |
+------------+---------------+-------------------------+--------------+

Querying chDB

Data ingestion is done, now it’s time for the fun part - querying the data! Each borough is divided into a different number of taxi zones. We’re going to write a query that joins the two tables to find out how many trips were picked up in each borough, and how many trips that works out to per taxi zone:

%%sql
SELECT pickup_borough AS borough,
       zone_count,
       count() AS trips,
       round(count() / zone_count) AS trips_per_zone
FROM taxi.trips
JOIN (
    SELECT Borough, count() AS zone_count
    FROM taxi.zones
    GROUP BY Borough
) AS zones ON pickup_borough = zones.Borough
GROUP BY borough, zone_count
ORDER BY trips DESC

+---------------+------------+---------+----------------+
|    borough    | zone_count |  trips  | trips_per_zone |
+---------------+------------+---------+----------------+
|   Manhattan   |     69     | 2713990 |    39333.0     |
|     Queens    |     69     |  187737 |     2721.0     |
|    Brooklyn   |     61     |  52445  |     860.0      |
|    Unknown    |     2      |  43802  |    21901.0     |
|     Bronx     |     43     |   2300  |      53.0      |
| Staten Island |     20     |    43   |      2.0       |
+---------------+------------+---------+----------------+

Manhattan and Queens have the same number of taxi zones, but Manhattan generates over 14 times as many pickups.

Saving queries

We can save queries using the --save parameter on the same line as the %%sql magic. The --no-execute parameter means that query execution will be skipped.

%%sql --save tips_by_neighborhood --no-execute
SELECT pickup_ntaname AS neighborhood,
       count() AS trips,
       round(avg(tip_amount), 2) AS avg_tip
FROM taxi.trips
WHERE fare_amount > 0 AND pickup_ntaname != ''
GROUP BY neighborhood
ORDER BY avg_tip DESC

When we run a saved query it will be converted into a Common Table Expression (CTE) before executing. In the following query we compute the neighborhoods with the highest average tip:

%sql SELECT * FROM tips_by_neighborhood ORDER BY avg_tip DESC LIMIT 5

+-----------------------------------+-------+---------+
|            neighborhood           | trips | avg_tip |
+-----------------------------------+-------+---------+
| New Springville-Bloomfield-Travis |   2   |   35.0  |
|       New Dorp-Midland Beach      |   2   |  23.74  |
|      New Brighton-Silver Lake     |   3   |  16.67  |
|           Newark Airport          |  201  |  11.89  |
|   Grymes Hill-Clifton-Fox Hills   |   1   |   11.3  |
+-----------------------------------+-------+---------+

The top entries are neighborhoods with only a handful of trips, so a single generous ride skews the average. Let’s filter those out.

Querying with parameters

We can also use parameters in our queries. Parameters are just normal variables:

min_trips = 10000

And then we can use the {{variable}} syntax in our query. The following query finds the neighborhoods with the highest average tip among those with more than 10,000 trips:

%%sql
SELECT * FROM tips_by_neighborhood
WHERE trips >= {{min_trips}}
ORDER BY avg_tip DESC
LIMIT 10

+----------------------------------------+--------+---------+
|              neighborhood              | trips  | avg_tip |
+----------------------------------------+--------+---------+
|                Airport                 | 151171 |   4.92  |
|   Battery Park City-Lower Manhattan    | 89110  |   2.16  |
|         North Side-South Side          | 11152  |   1.79  |
| SoHo-TriBeCa-Civic Center-Little Italy | 144887 |   1.65  |
|               Chinatown                | 54780  |   1.65  |
|            Lower East Side             | 15753  |   1.64  |
|              East Village              | 99881  |   1.61  |
|  Hunters Point-Sunnyside-West Maspeth  | 10054  |   1.58  |
|        Turtle Bay-East Midtown         | 197035 |   1.57  |
|              West Village              | 210369 |   1.54  |
+----------------------------------------+--------+---------+

Airport pickups tip the most by a wide margin — those long rides into the city add up.

Plotting histograms

JupySQL also has limited charting functionality. We can create box plots or histograms. We’re going to create a histogram, but first let’s write (and save) a query that returns the distance of each trip under 20 miles. We’ll be able to use this to create a histogram that counts how many trips fall into each distance bucket:

%%sql --save trip_distances --no-execute
SELECT trip_distance
FROM taxi.trips
WHERE trip_distance > 0 AND trip_distance < 20

We can then create a histogram by running the following:

from sql.ggplot import ggplot, geom_histogram, aes

plot = (
  ggplot(
    table="trip_distances",
    with_="trip_distances",
    mapping=aes(x="trip_distance", fill="#69f0ae", color="#fff"),
  ) + geom_histogram(bins=50)
)

Most trips are short hops of one to three miles, with a long tail stretching out towards the airport runs.

Intro to JupySQL with chDB (YouTube)

Get started

Language integrations

DataStore

Configuration

Debugging

Developer guides

Technical reference

Setup

Downloading a dataset

Configuring chDB and JupySQL

Querying data in TSV files

Importing TSV files into chDB

Querying chDB

Saving queries

Querying with parameters

Plotting histograms

​Setup

​Downloading a dataset

​Configuring chDB and JupySQL

​Querying data in TSV files

​Importing TSV files into chDB

​Querying chDB

​Saving queries

​Querying with parameters

​Plotting histograms

​Related

Setup

Downloading a dataset

Configuring chDB and JupySQL

Querying data in TSV files

Importing TSV files into chDB

Querying chDB

Saving queries

Querying with parameters

Plotting histograms

Related