Вот пример, который можно запустить и который демонстрирует проблему. Начальный 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 перед выполнением соединений, индексы верны.

Вот соответствующие фрагменты кода:

1. Создание меток 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")
)

2. Присоединение без сбора (это дает неверные индексы):

# 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()

3. Присоединение после сбора меток (это дает правильные значения 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?

Будем очень признательны за любую информацию о том, почему это происходит и как это решить!