Поляры — взять окно из N строк, окружающих строку, выполняющую условие

Поляры — взять окно из N строк, окружающих строку, выполняющую условие ⇐ Python

Ответить

1 сообщение • Страница 1 из 1

Anonymous

Поляры — взять окно из N строк, окружающих строку, выполняющую условие

Цитата

Сообщение Anonymous » 02 ноя 2025, 19:43

Рассмотрим следующий фрейм данных:

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

df = pl.DataFrame({
"letters": ["A", "B", "C", "D", "E", "F", "G", "H"],
"values": ["aa", "bb", "cc", "dd", "ee", "ff", "gg", "hh"]
})

print(df)
shape: (8, 2)
┌─────────┬────────┐
│ letters ┆ values │
│ ---     ┆ ---    │
│ str     ┆ str    │
╞═════════╪════════╡
│ A       ┆ aa     │
│ B       ┆ bb     │
│ C       ┆ cc     │
│ D       ┆ dd     │
│ E       ┆ ee     │
│ F       ┆ ff     │
│ G       ┆ gg     │
│ H       ┆ hh     │
└─────────┴────────┘

Как мне создать окно размера +/- N вокруг любой строки, удовлетворяющей заданному условию? Например, условием является pl.col("letters").contains("D|F") и N = 2. Тогда результат должен быть:

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

┌─────────┬────────────────────────────────┐
│ letters ┆ output                         │
│ ---     ┆ ---                            │
│ str     ┆ list[str]                      │
╞═════════╪════════════════════════════════╡
│ D       ┆ ["bb", "cc", "dd", "ee", "ff"] │
│ F       ┆ ["dd", "ee", "ff", "gg", "hh"] │
└─────────┴────────────────────────────────┘

Обратите внимание, что в этом случае окна перекрываются (окно F также содержит dd, а окна D также содержат ff). Также обратите внимание, что здесь для простоты N = 2, но на самом деле оно будет больше (~ 10 - 20). И набор данных относительно большой, поэтому я хотел бы сделать это максимально эффективно, не увеличивая использование памяти.

РЕДАКТИРОВАТЬ: Чтобы сделать вопрос более явным, вот запрос в синтаксисе SQL DuckDB, который дает правильный ответ (и я хотел бы знать, как перевести его в Polars):

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

df_table = df.to_arrow()
con = duckdb.connect()
query = """
SELECT
letters,
list(values) OVER (
ROWS BETWEEN 2 PRECEDING
AND 2 FOLLOWING
) as combined
FROM df_table
QUALIFY letters in ('D', 'F')
"""
print(pl.from_arrow(con.execute(query).arrow()))

shape: (2, 2)
┌─────────┬────────────────────────┐
│ letters ┆ combined               │
│ ---     ┆ ---                    │
│ str     ┆ list[str]              │
╞═════════╪════════════════════════╡
│ D       ┆ ["bb", "cc", ... "ff"] │
│ F       ┆ ["dd", "ee", ... "hh"] │
└─────────┴────────────────────────┘

Эталоны предлагаемых решений
Я запустил предлагаемые решения в блокноте Jupyter на одном из компьютеров Amazon ml.c5.xlarge. Пока ноутбук работал, я также держал htop открытым в терминале, чтобы наблюдать за использованием процессора и памяти. В наборе данных было более 12 миллионов строк.
Я запускал оба решения с помощью как активного, так и ленивого API. На всякий случай я также попробовал использовать простой цикл for Python для извлечения срезов после идентификации интересующих строк, а также DuckDB.
Сводная таблица
Polars продемонстрировал действительно высокую производительность и разумное использование памяти (с помощью метода @jqurious') благодаря умной реализации .shift() без копирования. Удивительно, но хорошо продуманный цикл for в Python справился с такой же задачей. DuckDB показала довольно плохие результаты как по скорости, так и по использованию памяти.
Ни Polars, ни DuckDB не используют для своей работы более одного ядра. Не уверен, связано ли это с отсутствием оптимизации или эта проблема просто поддается распараллеливанию. Я полагаю, что мы фильтруем только один столбец, а затем берем фрагменты этого же столбца, поэтому несколько потоков не могут сделать многого.

метод
использование процессора
использование памяти
время

ΩΠΟΚΕΚΡΥΜΜΕΝΟΣ
одноядерный
взрыв

jqurious
одноядерный
2,53–2,53 ГБ
4,63 с

(smart) для цикла
одноядерный
2,53G до 2,58 ГБ
4,91 с

DuckDB
одноядерный
от 1,62 ГБ до 6,13 ГБ
38,6 с

Использование процессора показывает, были ли задействованы несколько ядер во время операции.
Использование памяти показывает, сколько памяти использовалось до операции, а также максимальное использование памяти во время операции.

@ΩΠΟΚΕΚΡΥΜΜΕΝΟΣ Решение:

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

preceding = 2
following = 2

look_around = [pl.col("body").shift(-i)
for i in range(-preceding, following + 1)]

(
df
.with_columns(
pl.when(pl.col('body').str.contains(regex))
.then(pl.concat_list(look_around))
.alias('combined')
)
.filter(pl.col('combined').is_not_null())
)

К сожалению, в моем довольно большом наборе данных это решение привело к резкому увеличению использования памяти и сбою ядра как при активном, так и при ленивом API.
Решение @jqurious

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

preceding = 2
following = 2

look_around = [
pl.col("body").shift(-i).alias(f"lag_{i}") for i in range(-preceding, following + 1)
]

(
df
.with_columns(
look_around
)
.filter(pl.col("body").str.contains(regex))
.select(
pl.col("body"),
pl.concat_list([f"lag_{i}" for i in range(-2, 3)]).alias("output")
)
)

готовность:

использование процессора: одноядерное
использование памяти: 2,53 ГБ -> 2,53 ГБ
время: 4,63 с ± 6,6 мс на цикл (среднее значение ± стандартное dev. из 7 запусков, по 1 циклу каждый)

[*]ленивый:

использование процессора: одноядерное
использование памяти: 2,53 ГБ -> 2,53 ГБ
время: 4,63 с ± 3,85 мс на цикл (среднее ± стандартное отклонение для 7 запусков, по 1 циклу каждый)

(Smart) Python для цикла

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

preceding = 2
following = 2

output = []

indices = df.with_row_index().select(
pl.col("index").filter(pl.col("body").str.contains(regex))
)["index"]

for idx, x in enumerate(indices):
offset = max(0, x - preceding)
length = preceding + following + 1
output.append(df["body"].slice(offset, length))

Использование процессора: одноядерное
Использование памяти: 2,53 ГБ -> 2,58 ГБ
время: 4,91 с ± 24,5 мс на цикл (среднее ± стандартное отклонение для 7 запусков по 1 циклу каждый)

DuckDB
Обратите внимание, что я сначала преобразовал df в Arrow.Table, прежде чем запускать запрос, чтобы DuckDB мог напрямую воздействовать на него. Кроме того, я не уверен, что преобразование результата обратно в Arrow требует огромного количества вычислений и несправедливо по отношению к нему.

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

preceding = 2
following = 2

query = f"""
SELECT
body,
list(body) OVER (
ROWS BETWEEN {preceding} PRECEDING
AND {following} FOLLOWING
) as combined
FROM df_table
QUALIFY regexp_matches(body, '{regex}')
"""

result = con.execute(query).arrow()

При использовании DuckDB моя первая попытка выполнить вычисления потерпела неудачу. Мне пришлось повторить попытку, прочитав таблицу стрелок напрямую, без использования Polars (это сэкономило около 1 ГБ памяти), чтобы дать DuckDB больше памяти.

первая попытка:

процессор: одноядерный
память: 2,53 ГБ -> 6,93 ГБ -> сбой!
время: нет

[*]вторая попытка:

процессор: одноядерный
память: 1,62 ГБ -> 6,13 ГБ
время: 38,6 с ± 311 мс на цикл (среднее ± стандартное отклонение для 7 циклов, по 1 циклу каждый)

Подробнее здесь: https://stackoverflow.com/questions/745 ... -condition

1762101824

Anonymous

Рассмотрим следующий фрейм данных:
[code]df = pl.DataFrame({
"letters": ["A", "B", "C", "D", "E", "F", "G", "H"],
"values": ["aa", "bb", "cc", "dd", "ee", "ff", "gg", "hh"]
})

print(df)
shape: (8, 2)
┌─────────┬────────┐
│ letters ┆ values │
│ ---     ┆ ---    │
│ str     ┆ str    │
╞═════════╪════════╡
│ A       ┆ aa     │
│ B       ┆ bb     │
│ C       ┆ cc     │
│ D       ┆ dd     │
│ E       ┆ ee     │
│ F       ┆ ff     │
│ G       ┆ gg     │
│ H       ┆ hh     │
└─────────┴────────┘
[/code]
Как мне создать окно размера +/- N вокруг любой строки, удовлетворяющей заданному условию? Например, условием является pl.col("letters").contains("D|F") и N = 2. Тогда результат должен быть:
[code]┌─────────┬────────────────────────────────┐
│ letters ┆ output                         │
│ ---     ┆ ---                            │
│ str     ┆ list[str]                      │
╞═════════╪════════════════════════════════╡
│ D       ┆ ["bb", "cc", "dd", "ee", "ff"] │
│ F       ┆ ["dd", "ee", "ff", "gg", "hh"] │
└─────────┴────────────────────────────────┘
[/code]
Обратите внимание, что в этом случае окна перекрываются (окно F также содержит dd, а окна D также содержат ff). Также обратите внимание, что здесь для простоты N = 2, но на самом деле оно будет больше (~ 10 - 20). И набор данных относительно большой, поэтому я хотел бы сделать это максимально эффективно, не увеличивая использование памяти.

[b]РЕДАКТИРОВАТЬ:[/b] Чтобы сделать вопрос более явным, вот запрос в синтаксисе SQL DuckDB, который дает правильный ответ (и я хотел бы знать, как перевести его в Polars):
[code]df_table = df.to_arrow()
con = duckdb.connect()
query = """
SELECT
letters,
list(values) OVER (
ROWS BETWEEN 2 PRECEDING
AND 2 FOLLOWING
) as combined
FROM df_table
QUALIFY letters in ('D', 'F')
"""
print(pl.from_arrow(con.execute(query).arrow()))

shape: (2, 2)
┌─────────┬────────────────────────┐
│ letters ┆ combined               │
│ ---     ┆ ---                    │
│ str     ┆ list[str]              │
╞═════════╪════════════════════════╡
│ D       ┆ ["bb", "cc", ... "ff"] │
│ F       ┆ ["dd", "ee", ... "hh"] │
└─────────┴────────────────────────┘
[/code]

Эталоны предлагаемых решений
Я запустил предлагаемые решения в блокноте Jupyter на одном из компьютеров Amazon ml.c5.xlarge. Пока ноутбук работал, я также держал htop открытым в терминале, чтобы наблюдать за использованием процессора и памяти. В наборе данных было более 12 миллионов строк.
Я запускал оба решения с помощью как активного, так и ленивого API. На всякий случай я также попробовал использовать простой цикл for Python для извлечения срезов после идентификации интересующих строк, а также DuckDB.
Сводная таблица
Polars продемонстрировал действительно высокую производительность и разумное использование памяти (с помощью метода @jqurious') благодаря умной реализации .shift() без копирования. Удивительно, но хорошо продуманный цикл for в Python справился с такой же задачей. DuckDB показала довольно плохие результаты как по скорости, так и по использованию памяти.
Ни Polars, ни DuckDB не используют для своей работы более одного ядра. Не уверен, связано ли это с отсутствием оптимизации или эта проблема просто поддается распараллеливанию.  Я полагаю, что мы фильтруем только один столбец, а затем берем фрагменты этого же столбца, поэтому несколько потоков не могут сделать многого.



метод
использование процессора
использование памяти
время




ΩΠΟΚΕΚΡΥΜΜΕΝΟΣ
одноядерный
взрыв



jqurious
одноядерный
2,53–2,53 ГБ
4,63 с


(smart) для цикла
одноядерный
2,53G до 2,58 ГБ
4,91 с


DuckDB
одноядерный
от 1,62 ГБ до 6,13 ГБ
38,6 с



[list]
[*]Использование процессора показывает, были ли задействованы несколько ядер во время операции.
[*]Использование памяти показывает, сколько памяти использовалось до операции, а также максимальное использование памяти во время операции.
[/list]
@ΩΠΟΚΕΚΡΥΜΜΕΝΟΣ Решение:
[code]preceding = 2
following = 2

look_around = [pl.col("body").shift(-i)
for i in range(-preceding, following + 1)]

(
df
.with_columns(
pl.when(pl.col('body').str.contains(regex))
.then(pl.concat_list(look_around))
.alias('combined')
)
.filter(pl.col('combined').is_not_null())
)
[/code]
К сожалению, в моем довольно большом наборе данных это решение привело к резкому увеличению использования памяти и сбою ядра как при активном, так и при ленивом API.
Решение @jqurious
[code]preceding = 2
following = 2

look_around = [
pl.col("body").shift(-i).alias(f"lag_{i}") for i in range(-preceding, following + 1)
]

(
df
.with_columns(
look_around
)
.filter(pl.col("body").str.contains(regex))
.select(
pl.col("body"),
pl.concat_list([f"lag_{i}" for i in range(-2, 3)]).alias("output")
)
)
[/code]
[list]
[*][b]готовность:[/b]

[b]использование процессора:[/b] одноядерное
[*][b]использование памяти:[/b] 2,53 ГБ -> 2,53 ГБ
[*][b]время:[/b] 4,63 с ± 6,6 мс на цикл (среднее значение ± стандартное dev. из 7 запусков, по 1 циклу каждый)
[/list]

[*][b]ленивый:[/b]
[list]
[b]использование процессора:[/b] одноядерное
[*][b]использование памяти:[/b] 2,53 ГБ -> 2,53 ГБ
[*][b]время:[/b] 4,63 с ± 3,85 мс на цикл (среднее ± стандартное отклонение для 7 запусков, по 1 циклу каждый)
[/list]


(Smart) Python для цикла
[code]preceding = 2
following = 2

output = []

indices = df.with_row_index().select(
pl.col("index").filter(pl.col("body").str.contains(regex))
)["index"]

for idx, x in enumerate(indices):
offset = max(0, x - preceding)
length = preceding + following + 1
output.append(df["body"].slice(offset, length))
[/code]
[list]
[*][b]Использование процессора:[/b] одноядерное
[*][b]Использование памяти:[/b] 2,53 ГБ -> 2,58 ГБ
[*][b]время:[/b] 4,91 с ± 24,5 мс на цикл (среднее ± стандартное отклонение для 7 запусков по 1 циклу каждый)
[/list]
DuckDB
Обратите внимание, что я сначала преобразовал df в Arrow.Table, прежде чем запускать запрос, чтобы DuckDB мог напрямую воздействовать на него. Кроме того, я не уверен, что преобразование результата обратно в Arrow требует огромного количества вычислений и несправедливо по отношению к нему.
[code]preceding = 2
following = 2

query = f"""
SELECT
body,
list(body) OVER (
ROWS BETWEEN {preceding} PRECEDING
AND {following} FOLLOWING
) as combined
FROM df_table
QUALIFY regexp_matches(body, '{regex}')
"""

result = con.execute(query).arrow()
[/code]
При использовании DuckDB моя первая попытка выполнить вычисления потерпела неудачу.  Мне пришлось повторить попытку, прочитав таблицу стрелок напрямую, без использования Polars (это сэкономило около 1 ГБ памяти), чтобы дать DuckDB больше памяти.
[list]
[*][b]первая попытка:[/b]

[b]процессор:[/b] одноядерный
[*][b]память:[/b] 2,53 ГБ -> 6,93 ГБ -> сбой!
[*][b]время:[/b] нет
[/list]

[*][b]вторая попытка:[/b]
[list]
[b]процессор:[/b] одноядерный
[*][b]память:[/b] 1,62 ГБ -> 6,13 ГБ
[*][b]время:[/b] 38,6 с ± 311 мс на цикл (среднее ± стандартное отклонение для 7 циклов, по 1 циклу каждый)
[/list]

 

Подробнее здесь: [url]https://stackoverflow.com/questions/74513494/polars-take-a-window-of-n-rows-surrounding-a-row-fulfilling-a-condition[/url]