Polars — разделение коллекции LazyFrame на основе ключевого столбца ⇐ Python

[Anonymous]

Проблема
У меня довольно большой LazyFrame, который аварийно завершает работу, даже если я использую .collect(streaming=True). Чтобы разделить вычисление на более мелкие подмножества, которые можно обрабатывать в памяти, я хотел бы разделить таблицу на несколько таблиц на основе значения «ключевого столбца».

Пример

Код: Выделить всё

shape: (3, 213,231,322)
┌───────────────────────────────┬─────────┬────────────┐
│ Timestamp                     ┆ Temp    ┆ Location   │
│ ---                           ┆ ---     ┆ ---        │
│ pl.Datetime                   ┆ str     ┆ str        │
╞═══════════════════════════════╪═════════╪════════════╡
│ 2023-08-01 23:06:99.512383    ┆ 19      ┆ Dallas     │
│ 2023-08-01 23:21:01.818792    ┆ 20      ┆ Austin     │
│ ...                           ┆ ...     ┆ ...        │
│ 2023-08-30 23:23:00.238093    ┆ 21      ┆ New York   │
└───────────────────────────────┴─────────┴────────────┘

Чтобы выразить это конкретными словами. Глядя на таблицу выше, я вижу массивный LazyFrame с точными данными о погоде для множества городов. Поскольку один столбец — «Местоположение», я бы хотел разделить этот большой столбец на N меньших, по одному для каждого города. Тогда я бы вместо этого запустил для них .collect_all(streaming=True). Ключи заранее не известны. Поэтому мы не можем жестко запрограммировать цикл for...

Что у меня есть
Сейчас я делаю это так: запускаю исходную коллекцию, которая тщательно фильтруется для получения набора уникальных ключей.
Затем я запускаю .collect_all() для списка ленивых фреймов, отфильтрованных по ключу. Это минимально жизнеспособный (вроде как) работоспособный пример того, как именно работает мой код прямо сейчас:

Код: Выделить всё

# Load initial master data, use .cache() to speed up individual symbols' collection
master_lf = pl.scan_csv(...).cache()

# Get unique keys:
keys = master_lf.select(pl.col('key_col')).unique(subset='key_col').collect()['key_col']

# Create Symbol LazyFrames
symbol_lfs = [df.filter(pl.col('key_col') == k) for k in keys]

# Collect Them All
symbol_dfs = pl.collect_all(symbol_lfs, streaming=True)

Чего бы мне хотелось
Этот первоначальный сбор раздражает. Кажется, мы можем оптимизировать еще больше, нет? Есть ли способ избежать этого?

Изменить:
Я нашел что-то под названием parition_by, которое делает именно это, но для df. Похоже, Polars нужна потоковая версия.

Подробнее здесь: https://stackoverflow.com/questions/772 ... key-column