Установка и первые шаги

DataFusion предоставляет полноценные Python-биндинги, позволяющие использовать тот же движок запросов из Python-кода. Вы получаете производительность Rust с удобством Python.

Установка

Пакет datafusion доступен на PyPI. Требуется Python 3.10 или выше:

pip install datafusion

Проверка установки:

import datafusion
print(datafusion.__version__)  # 53.0.0

SessionContext в Python

Точка входа в DataFusion Python --- SessionContext. Это Python-обёртка над Rust SessionContext, предоставляющая те же возможности:

from datafusion import SessionContext

ctx = SessionContext()

# Выполнить SQL
df = ctx.sql("SELECT 1 + 2 AS result")
df.show()
# +--------+
# | result |
# +--------+
# | 3      |
# +--------+

Чтение файлов

SessionContext предоставляет методы для чтения основных форматов:

CSV

ctx = SessionContext()
df = ctx.read_csv("data/orders.csv")
df.show()

По умолчанию DataFusion определяет типы колонок автоматически. Для явного указания схемы:

import pyarrow as pa

schema = pa.schema([
    ("order_id", pa.int64()),
    ("region", pa.utf8()),
    ("amount", pa.float64()),
])

df = ctx.read_csv("data/orders.csv", schema=schema)

Parquet

df = ctx.read_parquet("data/orders.parquet")
df.show()

Parquet-файлы содержат схему внутри, поэтому указание типов не требуется. DataFusion использует predicate pushdown и projection pruning при чтении Parquet.

JSON (NDJSON)

df = ctx.read_json("data/events.json")
df.show()

Регистрация таблиц

Для повторного использования данных в SQL-запросах удобно регистрировать файлы как именованные таблицы:

ctx = SessionContext()

ctx.register_csv("orders", "data/orders.csv")
ctx.register_parquet("events", "data/events.parquet")

# Теперь доступны в SQL
df = ctx.sql("SELECT * FROM orders WHERE amount > 100")
df.show()

Создание данных из Python-структур

Когда данные уже в Python, можно создать DataFrame без файлов:

from_pydict

ctx = SessionContext()

data = {
    "name": ["Alice", "Bob", "Carol"],
    "score": [95, 82, 91],
}

df = ctx.from_pydict(data, name="students")
df.show()
# +-------+-------+
# | name  | score |
# +-------+-------+
# | Alice | 95    |
# | Bob   | 82    |
# | Carol | 91    |
# +-------+-------+

from_pylist

rows = [
    {"name": "Alice", "score": 95},
    {"name": "Bob", "score": 82},
    {"name": "Carol", "score": 91},
]

df = ctx.from_pylist(rows, name="students")

Глобальные функции контекста

Начиная с версии 46.0, DataFusion Python предоставляет глобальные функции для быстрого старта без явного создания SessionContext:

import datafusion

# Глобальный контекст создаётся автоматически
df = datafusion.read_csv("data/orders.csv")
df.show()

df2 = datafusion.read_parquet("data/events.parquet")

Глобальные функции удобны для скриптов и notebook-ов, но для production-кода рекомендуется явно создавать SessionContext --- это даёт контроль над конфигурацией и изоляцию между сессиями.

Полный пример: загрузка и анализ

from datafusion import SessionContext

ctx = SessionContext()

# Регистрация данных
ctx.register_csv("orders", "data/orders.csv")
ctx.register_parquet("products", "data/products.parquet")

# Комбинирование SQL и DataFrame
df = ctx.sql("""
    SELECT o.region, p.category, SUM(o.amount) as total
    FROM orders o
    JOIN products p ON o.product_id = p.id
    GROUP BY o.region, p.category
    ORDER BY total DESC
    LIMIT 10
""")

# Результат --- DataFrame, можно продолжить обработку
result = df.collect()
print(f"Rows: {sum(batch.num_rows for batch in result)}")

Итоги

• pip install datafusion --- единственная зависимость для начала работы
• SessionContext --- центральный объект для чтения данных и выполнения запросов
• read_csv, read_parquet, read_json --- чтение файлов в DataFrame
• from_pydict, from_pylist --- создание DataFrame из Python-структур
• register_csv, register_parquet --- регистрация файлов как SQL-таблиц
• Глобальные функции (с v46.0) --- быстрый старт без явного SessionContext