Фильтрация и выборка
IndexedParquetDataset предоставляет несколько способов ограничения и перемешивания данных. Помните, что все методы фильтрации возвращают новый объект датасета (lazy), не копируя сами данные.
Случайное перемешивание (shuffle)
Метод .shuffle() перемешивает порядок индексов. Вы можете зафиксировать результат, передав seed.
Фильтрация по условиям (filter)
Метод .filter() — самый мощный инструмент отбора данных. Он поддерживает три режима:
1. Серверная фильтрация (PyArrow-side)
Этот режим самый быстрый, так как выполняется на уровне C++ через PyArrow до того, как данные попадут в Python.
# Фильтрация по точному значению
dataset = dataset.filter(column_conditions={"status": "active"})
# Фильтрация по диапазону (используя кортеж)
dataset = dataset.filter(column_conditions={
"score": (">", 0.8),
"age": ("<=", 30)
})
2. Фильтрация через предикат (Python-side)
Если условий в column_conditions недостаточно, вы можете передать функцию-предикат. Она будет вызвана для каждой строки. Это медленнее, но гибче.
3. Фильтрация по файлам
Вы можете оставить данные только из определенных файлов, используя glob-патерны или список путей.
Выборка и ограничение
Ограничение количества (limit)
Случайная выборка (sample)
Выбор по индексам (select)
# Оставить только строки с конкретными индексами
dataset = dataset.select([0, 10, 50, 100])
# Можно использовать слайсы
dataset = dataset.select(slice(0, 500, 2)) # Каждая вторая строка из первых 500
Разделение на выборки (Train/Test Split)
Метод train_test_split — это "швейцарский нож" для подготовки обучения. Он поддерживает стратификацию!
# Обычный сплит 80/20
train_ds, val_ds = dataset.train_test_split(test_size=0.2, seed=42)
# Сплит со стратификацией по колонке 'category'
# (сохранит пропорции категорий в обеих выборках)
train_ds, val_ds = dataset.train_test_split(
test_size=0.2,
stratify_by="category"
)
[!NOTE] Стратификация требует прочтения указанной колонки для всех строк датасета в момент вызова метода. На огромных датасетах это может занять некоторое время.