Перейти к содержанию

Фильтрация и выборка

IndexedParquetDataset предоставляет несколько способов ограничения и перемешивания данных. Помните, что все методы фильтрации возвращают новый объект датасета (lazy), не копируя сами данные.

Случайное перемешивание (shuffle)

Метод .shuffle() перемешивает порядок индексов. Вы можете зафиксировать результат, передав seed.

# Перемешать всё
dataset = dataset.shuffle(seed=42)

Фильтрация по условиям (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 недостаточно, вы можете передать функцию-предикат. Она будет вызвана для каждой строки. Это медленнее, но гибче.

dataset = dataset.filter(predicate=lambda x: len(x["text"]) > 100 and x["label"] in [1, 5])

3. Фильтрация по файлам

Вы можете оставить данные только из определенных файлов, используя glob-патерны или список путей.

# Только файлы из папки 2023 года
dataset = dataset.filter(path_filter="**/2023/*.parquet")

Выборка и ограничение

Ограничение количества (limit)

# Взять только первые 1000 строк
dataset = dataset.limit(1000)

Случайная выборка (sample)

# Выбрать 500 случайных строк без повторений
dataset = dataset.sample(500, seed=123)

Выбор по индексам (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] Стратификация требует прочтения указанной колонки для всех строк датасета в момент вызова метода. На огромных датасетах это может занять некоторое время.