Почему сбор LazyFrame перед объединением в Polars решает мою проблему с несоответствиями индексов? ⇐ Python

[Anonymous]

РЕДАКТИРОВАТЬ
Вот пример, который можно запустить и который демонстрирует проблему. Начальный LazyFrame включает попарные расстояния между точками на плоскости. Поскольку расстояние A->B равно расстоянию B->A, я сохраняю только уникальные пары, а затем создаю метки LazyFrame, чтобы сохранять только уникальные метки. Я соединяю исходный LazyFrame с метками LazyFrame, чтобы получить соответствующие индексы для col_1 и col_2. . Например, в результирующем кадре данных расстояние от значения индекса до того же значения индекса (например, A-A (0->0), B-B (1->1), CC (2->2) ) должно быть равно нулю, но это не так, когда LazyFrame не собирается перед объединением.

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

import polars as pl
import numpy as np

data = {
"col_1": ["A", "A", "A", "B", "B", "B", "C", "C", "C"],
"col_2": ["B", "C", "A", "A", "C", "B", "A", "B", "C"],
"col_3": [1.0, 2.0, 0.0, 1.0, 1.5, 0.0, 2.0, 1.5, 0.0]
}

# Create LazyFrame
pairwise_distances = pl.LazyFrame(data)

# Function to concatenate and sort two elements and return a string with '||'
def concat_and_sort(a, b):
return "||".join(sorted([a, b]))

# Add pairs column and get unique pairs
pairwise_distances = pairwise_distances.with_columns(
pl.struct(["col_1", "col_2"]).map_elements(
lambda x: concat_and_sort(x["col_1"], x["col_2"]),
return_dtype=pl.String
).alias("pairs")
).unique(subset="pairs").select(pl.col("col_1"), pl.col("col_2"), pl.col("col_3"))

# Concatenate col_1 and col_2, and get unique labels with index
labels = (
pl.concat([
pairwise_distances.select(pl.col("col_1").alias("label")),
pairwise_distances.select(pl.col("col_2").alias("label"))
])
.unique(keep="first")
.with_row_count(name="index")
)

# Collect labels LazyFrame to DataFrame
labels_df = labels.collect()
pairwise_distances_df = pairwise_distances.collect()

# Join to get index_1 (without collecting labels)
data_joined_direct = pairwise_distances.join(
labels,
left_on="col_1",
right_on="label",
how="left"
).rename({"index": "index_1"})

# Join to get index_2 (without collecting labels)
data_joined_direct = data_joined_direct.join(
labels,
left_on="col_2",
right_on="label",
how="left"
).rename({"index": "index_2"})

# Select relevant columns (without collecting labels)
result_direct = data_joined_direct.select(["index_1", "index_2", "col_3"])

# Collect the result to execute the lazy operations (without collecting labels first)
# Result without collecting labels first
result_direct_df = result_direct.collect()

# Join to get index_1 (with collecting labels first)
data_joined_collected = pairwise_distances_df.join(
labels_df,
left_on="col_1",
right_on="label",
).rename({"index": "index_1"})

# Join to get index_2 (with collecting labels first)
data_joined_collected = data_joined_collected.join(
labels_df,
left_on="col_2",
right_on="label",
).rename({"index": "index_2"})

# Select relevant columns (with collecting labels first)
result_collected = data_joined_collected.select(["index_1", "index_2", "col_3"])

Исходное сообщение
Я использую Polars для обработки набора данных, где мне нужно создать уникальные метки из двух столбцов, а затем выполнить соединения для получения индексов для этих ярлыков. Однако я заметил, что если я выполняю соединения непосредственно в LazyFrame, индексы кажутся неправильными. Когда я собираю LazyFrame в DataFrame перед выполнением соединений, индексы верны.
Вот соответствующие фрагменты кода:

• Создание ярлыков LazyFrame:

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

import polars as pl

# Assume data is a LazyFrame
data = pl.scan_csv(
source=source_filepath,
separator='\t',
has_header=False,
)

# Concatenate col_1 and col_2, and get unique labels with index
labels = (
pl.concat([
data.select(pl.col("col_1").alias("label")),
data.select(pl.col("col_2").alias("label"))
])
.unique(keep="first")
.with_row_count(name="label_index")
)

• Присоединение без сбора (при этом даются неправильные индексы):

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

# Join to get index_1
data = data.join(
labels,
left_on="col_1",
right_on="label",
).rename({"label_index": "index_1"})

# Join to get index_2
data = data.join(
labels,
left_on="col_2",
right_on="label",
).rename({"label_index": "index_2"})

result = data.select(["index_1", "index_2", "col_3"])
result_df = result.collect()

• Присоединение после сбора меток (это дает правильные значения index_1 и index_2):

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

# Collect labels and data LazyFrame to DataFrame
labels_df = labels.collect()
data_df = data.collect()

# Join to get index_1
data_df = data_df.join(
labels_df,
left_on="col_1",
right_on="label",
how="left"
).rename({"label_index": "index_1"})

# Join to get index_2
data_df = data_df.join(
labels_df,
left_on="col_2",
right_on="label",
how="left"
).rename({"label_index": "index_2"})

result_df = data_df.select(["index_1", "index_2", "col_3"])

Почему существует такое несоответствие между использованием LazyFrame непосредственно для соединений и его сбором перед выполнением соединений? Как я могу обеспечить правильное поведение без необходимости преждевременного сбора LazyFrame?
Будем очень благодарны за любую информацию о том, почему это происходит и как это решить!

Подробнее здесь: https://stackoverflow.com/questions/786 ... with-index