Python Polars时间戳同步延迟计算

我想使用Polars LazyFrames将两个数量众多的时间戳数组彼此同步.

让我们假设我有两个使用LazyFrames存储的时间戳的Numty数组:

import polars as pl


timestamps = pl.LazyFrame(
    np.array(
        [
            np.datetime64("1970-01-01T00:00:00.500000000"),
            np.datetime64("1970-01-01T00:00:01.500000000"),
            np.datetime64("1970-01-01T00:00:02.600000000"),
            np.datetime64("1970-01-01T00:00:03.400000000"),
            np.datetime64("1970-01-01T00:00:04.500000000"),
            np.datetime64("1970-01-01T00:00:05.300000000"),
            np.datetime64("1970-01-01T00:00:06.200000000"),
            np.datetime64("1970-01-01T00:00:07.400000000"),
            np.datetime64("1970-01-01T00:00:08.500000000"),
        ]
    ),
    schema={"values": pl.Datetime}
)

other_timestamps = pl.LazyFrame(
    np.array(
        [
            np.datetime64("1970-01-01T00:00:01.500000000"),
            np.datetime64("1970-01-01T00:00:02.000000000"),
            np.datetime64("1970-01-01T00:00:02.500000000"),
            np.datetime64("1970-01-01T00:00:04.500000000"),
            np.datetime64("1970-01-01T00:00:06.000000000"),
            np.datetime64("1970-01-01T00:00:06.500000000"),
        ]
    ),
    schema={"values": pl.Datetime}
)

我还在numpy中实现了预期的功能:

import numpy as np
import numpy.typing as npt


def _np_sync_to(
    timestamps: npt.ArrayLike[np.datetime64],
    other: npt.ArrayLike[np.datetime64],
    tolerance: str,
):
    outer_diffs = np.abs(np.subtract.outer(other, timestamps))
    closest_timestamps_indices = outer_diffs.argmin(0)
    closest_timestamps = other[closest_timestamps_indices]
    diffs = np.abs(closest_timestamps - timestamps)
    tolerance = parse_timedelta(tolerance)
    within_tolerance = diffs <= tolerance

    ts1_synced = timestamps[within_tolerance]
    ts2_synced = closest_timestamps[within_tolerance]

    return ts1_synced, ts2_synced


np_ts1_synced, np_ts2_synced = _np_sync_to(
    timestamps=np.squeeze(timestamps.collect().to_numpy()),
    other=np.squeeze(other_timestamps.collect().to_numpy()),
    tolerance="500ms",
)

预期结果为:

np_ts1_synced = np.array([
    np.datetime64('1970-01-01T00:00:01.500000000'), 
    np.datetime64('1970-01-01T00:00:02.600000000'), 
    np.datetime64('1970-01-01T00:00:04.500000000'), 
    np.datetime64('1970-01-01T00:00:06.200000000')
])

np_ts2_synced = np.array([
    np.datetime64('1970-01-01T00:00:01.500000000'), 
    np.datetime64('1970-01-01T00:00:02.500000000'), 
    np.datetime64('1970-01-01T00:00:04.500000000'), 
    np.datetime64('1970-01-01T00:00:06.000000000')
])

因此，同步的时间戳基本上是指定容差内最近的时间戳. 现在我想使用Polars LazyFrames实现相同的功能来处理 Big Data .

我试着用Polars等效地实现它，但是外部减法的维数不正确，我想有一种更好的方法来进行一般的计算:

# inspired by https://stackoverflow.com/questions/77748729/polars-equivalent-to-np-outer
def sync_to(timestamps, other, tolerance: str):
    def _outer(
        a: pl.DataFrame | pl.LazyFrame, b: pl.DataFrame | pl.LazyFrame
    ):
        # I guess the following line is incorrect
        nrows = pl.len().sqrt().cast(pl.Int32)
        return (
            a.select("values")
            .join(b.select("values"), how="cross")
            .select(
                computed=(pl.col("values") - pl.col("values_right")).abs()
            )
            .group_by(pl.arange(0, pl.len()) // nrows, maintain_order=True)
            .agg("computed")
            .select(pl.col("computed").list.to_struct())
            .unnest("computed")
        )

    outer_diffs = timestamps.pipe(_outer, other)

我的另一个 idea 是:

ts1 = timestamps.sort("values").join_asof(
    other.sort("values"),
    on="values",
    strategy="nearest",
    tolerance=tolerance,
)

但yields 并不是我想要的.