JupySQL e chDB

JupySQL é uma biblioteca Python que permite executar SQL em notebooks Jupyter e no shell do IPython. Neste guia, vamos aprender a consultar dados usando o chDB e o JupySQL.

Configuração

Vamos criar primeiro um ambiente virtual:

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

E então, vamos instalar o JupySQL, o IPython e o Jupyter Lab:

pip install jupysql ipython jupyterlab

Podemos usar o JupySQL no IPython, que pode ser iniciado executando:

ipython

Ou, no Jupyter Lab, executando:

jupyter lab

Baixando um conjunto de dados

Vamos usar o conjunto de dados de táxis da cidade de Nova York, que contém cerca de 3 milhões de corridas, incluindo a tarifa, a gorjeta e o bairro de embarque de cada uma. As corridas estão divididas em vários arquivos TSV, então vamos começar fazendo o download deles:

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",
  )

Configurando o chDB e o JupySQL

Em seguida, vamos importar o módulo dbapi do chDB:

from chdb import dbapi

E vamos criar uma conexão com o chDB. Todos os dados que persistirmos serão salvos no diretório taxi.chdb:

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

Agora, vamos carregar a magic sql e criar uma conexão com o chDB:

%load_ext sql
%sql conn --alias chdb

Em seguida, vamos mostrar o limite de exibição para que os resultados das consultas não sejam truncados:

%config SqlMagic.displaylimit = None

Consultando dados em arquivos TSV

Baixamos vários arquivos com o prefixo trips_. Vamos usar a cláusula DESCRIBE para entender o esquema:

%%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)

Também podemos executar uma consulta SELECT diretamente nesses arquivos para ver como os dados se apresentam:

%%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      |
+------------+---------------------+----------------------------------------+---------------+-------------+------------+

Ao examinarmos novamente o esquema, algumas colunas relacionadas a valores monetários — trip_distance, fare_amount e tip_amount — foram inferidas como String, em vez de um tipo numérico. Vamos corrigir isso ao importar os dados para uma tabela.

Importando arquivos TSV para o chDB

Agora, vamos armazenar os dados desses arquivos TSV em uma tabela. O banco de dados default não persiste dados em disco, portanto, primeiro precisamos criar outro banco de dados:

%sql CREATE DATABASE taxi

Agora, vamos criar uma tabela chamada trips, cujo esquema será derivado da estrutura dos dados nos arquivos TSV. Usaremos a cláusula REPLACE para converter as colunas de valores monetários em Float64 e a função transform para converter a coluna numérica pickup_borocode em um nome de distrito legível:

%%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

Vamos verificar rapidamente os dados na nossa tabela:

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

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

Pouco mais de 3 milhões de viagens — vamos também importar uma segunda tabela. A Taxi & Limousine Commission da cidade de Nova York divide a cidade em zonas de táxi, e um arquivo de consulta associa cada zona ao respectivo distrito. Vamos baixar esse arquivo:

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

Em seguida, crie uma tabela chamada zones com base no conteúdo do arquivo CSV:

%%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

Quando a execução terminar, podemos conferir os dados que ingerimos:

%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   |
+------------+---------------+-------------------------+--------------+

Consultando o chDB

A ingestão de dados foi concluída; agora chegou a hora da parte divertida: consultar os dados!

Cada distrito é dividido em um número diferente de zonas de táxi. Vamos escrever uma consulta que faz a junção das duas tabelas para descobrir quantas viagens começaram em cada distrito e quantas viagens isso representa por zona de táxi:

%%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 e Queens têm o mesmo número de zonas de táxi, mas Manhattan registra mais de 14 vezes mais corridas.

Salvando consultas

É possível salvar consultas usando o parâmetro --save na mesma linha do magic %%sql. O parâmetro --no-execute faz com que a execução da consulta seja ignorada.

%%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

Quando executamos uma consulta salva, ela é convertida em uma expressão de tabela comum (CTE) antes de ser executada. Na consulta a seguir, calculamos os bairros com a maior média de gorjetas:

%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  |
+-----------------------------------+-------+---------+

As primeiras entradas são bairros com poucas corridas, portanto uma única corrida com valor alto distorce a média. Vamos filtrá-los.

Consultas com parâmetros

Também podemos usar parâmetros em nossas consultas. Parâmetros são apenas variáveis comuns:

min_trips = 10000

Em seguida, podemos usar a sintaxe {{variable}} em nossa consulta. A consulta a seguir encontra os bairros com a maior média de gorjetas entre aqueles com mais de 10.000 viagens:

%%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  |
+----------------------------------------+--------+---------+

As corridas de aeroporto recebem, de longe, as maiores gorjetas — essas longas viagens até a cidade pesam no bolso.

Plotando histogramas

O JupySQL também tem recursos limitados de criação de gráficos. Podemos criar box plots ou histogramas.

Vamos criar um histograma, mas primeiro vamos escrever (e salvar) uma consulta que retorna a distância de cada viagem com menos de 20 milhas. Poderemos usar isso para criar um histograma que conta quantas viagens se enquadram em cada bucket de distância:

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

Em seguida, podemos criar um histograma da seguinte forma:

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)
)

A maioria das viagens é curta, de uma a três milhas, com uma longa cauda que se estende até as viagens para o aeroporto.

Introdução ao JupySQL com chDB (YouTube)