JupySQL и chDB

JupySQL — это библиотека Python, которая позволяет выполнять SQL в среде Jupyter Notebook и оболочке IPython. В этом руководстве мы узнаем, как выполнять запросы к данным с помощью chDB и JupySQL.

Подготовка

Сначала создадим виртуальное окружение:

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

Затем установим JupySQL, IPython и Jupyter Lab:

pip install jupysql ipython jupyterlab

Мы можем использовать JupySQL в IPython, который запускается командой:

ipython

Или в Jupyter Lab, выполнив команду:

jupyter lab

Скачивание набора данных

Мы будем использовать набор данных о такси Нью-Йорка, содержащий около 3 миллионов поездок, а также сведения о стоимости, чаевых и районе посадки для каждой из них. Данные о поездках распределены по нескольким файлам в формате TSV, поэтому сначала скачаем их:

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

Настройка chDB и JupySQL

Далее импортируем модуль dbapi для chDB:

from chdb import dbapi

И создадим подключение к chDB. Все сохранённые данные будут записаны в каталог taxi.chdb:

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

Теперь загрузим магию sql и создадим подключение к chDB:

%load_ext sql
%sql conn --alias chdb

Затем зададим лимит вывода, чтобы результаты запросов не обрезались:

%config SqlMagic.displaylimit = None

Запросы к данным в файлах в формате TSV

Мы скачали несколько файлов с префиксом trips_. Давайте используем оператор DESCRIBE, чтобы узнать схему:

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

Мы также можем выполнить SELECT-запрос напрямую к этим файлам, чтобы посмотреть, как выглядят данные:

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

Если вернуться к схеме, видно, что несколько столбцов с денежными значениями — trip_distance, fare_amount и tip_amount — были определены как String, а не как числовой тип. Мы исправим это при импорте данных в таблицу.

Импорт файлов в формате TSV в chDB

Теперь сохраним данные из этих файлов в формате TSV в таблице. База данных по умолчанию не хранит данные на диске, поэтому сначала необходимо создать другую базу данных:

%sql CREATE DATABASE taxi

Теперь создадим таблицу trips, схема которой будет определена на основе структуры данных в файлах в формате TSV. С помощью условия REPLACE приведём столбцы с денежными значениями к типу Float64, а с помощью функции transform преобразуем числовой столбец pickup_borocode в понятное человеку название боро:

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

Давайте быстро проверим данные в нашей таблице:

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

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

Чуть более 3 миллионов поездок — добавим ещё одну таблицу. Комиссия по такси и лимузинам Нью-Йорка делит город на зоны такси, а файл соответствий связывает каждую зону с соответствующим боро. Скачаем этот файл:

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

Затем создайте таблицу zones на основе содержимого 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

После завершения выполнения можно посмотреть на принятые данные:

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

Запросы к chDB

Ингестия данных завершена — теперь перейдём к самому интересному: запросам к данным!

Каждое боро включает разное количество зон такси. Напишем запрос, который соединит две таблицы, чтобы узнать, сколько поездок началось в каждом боро и сколько поездок приходится на каждую зону такси:

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

В Манхэттене и Куинсе одинаковое количество зон такси, но в Манхэттене совершается более чем в 14 раз больше посадок.

Сохранение запросов

Запросы можно сохранять с помощью параметра --save, указанного в той же строке, что и магическая команда %%sql. Параметр --no-execute позволяет пропустить выполнение запроса.

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

При запуске сохранённый запрос перед выполнением преобразуется в общее табличное выражение (CTE). В следующем запросе мы определяем районы с наибольшим средним размером чаевых:

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

В верхних строках — районы, где было всего несколько поездок, поэтому одна щедро оплаченная поездка искажает среднее значение. Отфильтруем их.

Выполнение запросов с параметрами

В запросах также можно использовать параметры. Параметры — это обычные переменные:

min_trips = 10000

Затем в запросе можно использовать синтаксис {{variable}}. Следующий запрос находит районы с наибольшим средним размером чаевых среди районов, где совершено более 10 000 поездок:

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

За поездки из аэропорта оставляют значительно больше чаевых — длительные поездки в город обходятся недёшево.

Построение гистограмм

JupySQL также предоставляет ограниченные возможности для построения диаграмм. Мы можем строить ящичные диаграммы или гистограммы.

Мы собираемся построить гистограмму, но сначала давайте напишем (и сохраним) запрос, который возвращает расстояние для каждой поездки длиной менее 20 миль. Это позволит нам построить гистограмму, показывающую, сколько поездок попадает в каждый бакет расстояний:

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

Затем можно построить гистограмму следующим образом:

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

Большинство поездок — это короткие маршруты длиной от одной до трёх миль, а длинный хвост распределения приходится на поездки в аэропорт.

Введение в JupySQL с chDB (YouTube)