So wiederholen und kürzen Sie Polarlistenelemente auf eine feste Länge ⇐ Python

[Anonymous]

Ich habe Daten, die wie folgt aussehen:

Code: Select all

import polars as pl

lf = pl.LazyFrame(
{
"points": [
[
[1.0, 2.0],
],
[
[3.0, 4.0],
[5.0, 6.0],
],
[
[7.0, 8.0],
[9.0, 10.0],
[11.0, 12.0],
],
],
"other": ["foo", "bar", "baz"],
},
schema={
"points": pl.List(pl.Array(pl.Float32, 2)),
"other": pl.String,
},
)

Und ich möchte dafür sorgen, dass alle Listen die gleiche Anzahl an Elementen haben.
Wenn sie derzeit mehr enthält, als ich benötige, sollte sie abgeschnitten werden.
Wenn sie weniger enthält, als ich benötige, sollte sie sich der Reihe nach wiederholen, bis genug vorhanden ist.
Ich habe es geschafft, es zum Laufen zu bringen, aber ich habe das Gefühl, dass ich durch die Luft springe. Gibt es eine sauberere Möglichkeit, dies zu tun? Vielleicht mit Gather?

Code: Select all

target_length = 3

result = (
lf.with_columns(
needed=pl.lit(target_length).truediv(pl.col("points").list.len()).ceil()
)
.with_columns(
pl.col("points")
.repeat_by("needed")
.list.eval(pl.element().explode())
.list.head(target_length)
)
.drop("needed")
)

BEARBEITEN
Die obige Methode funktioniert für Spielzeugbeispiele, aber wenn ich versuche, sie in meinem realen Datensatz zu verwenden, schlägt sie mit Folgendem fehl:

Code: Select all

pyo3_runtime.PanicException: Polars' maximum length reached. Consider installing 'polars-u64-idx'.

Ich konnte dafür kein MRE erstellen, aber meine Daten haben 4 Millionen Zeilen und die „Punkte“-Liste in jeder Zeile enthält zwischen 1 und 8000 Elemente (und ich versuche, auf 800 Elemente aufzufüllen/zu kürzen). Diese scheinen alle ziemlich klein zu sein, ich sehe nicht, wie eine maximale u32-Länge erreicht werden kann.
Ich freue mich über alle alternativen Ansätze, die ich ausprobieren kann.
Der nächste, den ich habe (was keine Panik auslöst), ist:
Aber das füllt die Liste nicht der Reihe nach auf. Es füllt nur auf und wiederholt das letzte Element.

Code: Select all

target_length = 3

result = (
lf.with_columns(
pl.col("points")
.list.gather(
pl.int_range(target_length),
null_on_oob=True,
)
.list.eval(pl.element().forward_fill())
)
.drop("needed")
)