[Python] How to count consecutive occurrences in a Polars DataFrame column?

I have a polars dataframe like the following:

shape: (200_583_873, 4)
┌──────┬─────┬──────────┬───────────┐
│ x ┆ y ┆ polarity ┆ timestamp │
│ --- ┆ --- ┆ --- ┆ --- │
│ u16 ┆ u16 ┆ i16 ┆ i64 │
╞══════╪═════╪══════════╪═══════════╡
│ 0 ┆ 0 ┆ 0 ┆ 4139422 │
│ 0 ┆ 0 ┆ 1 ┆ 4166683 │
│ 0 ┆ 0 ┆ 0 ┆ 4169904 │
│ 0 ┆ 0 ┆ 0 ┆ 8931478 │
│ 0 ┆ 0 ┆ 1 ┆ 9546786 │
│ … ┆ … ┆ … ┆ … │
│ 1279 ┆ 719 ┆ 0 ┆ 13067858 │
│ 1279 ┆ 719 ┆ 1 ┆ 13275637 │
│ 1279 ┆ 719 ┆ 0 ┆ 13277169 │
│ 1279 ┆ 719 ┆ 1 ┆ 13293135 │
│ 1279 ┆ 719 ┆ 0 ┆ 13294628 │
└──────┴─────┴──────────┴───────────┘


I want to count the consecutive occurances (by timestamp) of the same event polarity, per pixel. To illustrate what I mean, if look at just pixel 0,0, we have

shape: (12, 4)
┌─────┬─────┬──────────┬───────────┐
│ x ┆ y ┆ polarity ┆ timestamp │
│ --- ┆ --- ┆ --- ┆ --- │
│ u16 ┆ u16 ┆ i16 ┆ i64 │
╞═════╪═════╪══════════╪═══════════╡
│ 0 ┆ 0 ┆ 0 ┆ 4139422 │
│ 0 ┆ 0 ┆ 1 ┆ 4166683 │
│ 0 ┆ 0 ┆ 0 ┆ 4169904 │
│ 0 ┆ 0 ┆ 0 ┆ 8931478 │
│ 0 ┆ 0 ┆ 1 ┆ 9546786 │
│ 0 ┆ 0 ┆ 0 ┆ 9570021 │
│ 0 ┆ 0 ┆ 0 ┆ 14775777 │
│ 0 ┆ 0 ┆ 1 ┆ 14802742 │
│ 0 ┆ 0 ┆ 0 ┆ 14804840 │
│ 0 ┆ 0 ┆ 0 ┆ 21824841 │
│ 0 ┆ 0 ┆ 1 ┆ 21848296 │
│ 0 ┆ 0 ┆ 0 ┆ 21851722 │
└─────┴─────┴──────────┴───────────┘


And so we count 8 occurances of length 1 for 0 polarity and 4 occurances of length 1 for 1 polarity. I've attempted to accomplish this with the following function:

def count_occurances_and_length(df):
# **Step 1: Identify consecutive timestamp runs**
df = df.with_columns(
(pl.col("timestamp").diff().fill_null(1) > 1).cum_sum().alias("group_id")
)

# **Step 2: Compute actual run lengths**
run_lengths = (
df.group_by(["x", "y", "polarity", "group_id"])
.agg(pl.count().alias("run_length")) # Count rows per group (actual run length)
.drop("group_id")
)

# **Step 3: Count occurrences of each run length**
run_counts = (
run_lengths.group_by(["x", "y", "polarity", "run_length"])
.count()
.rename({"count": "occurrences"})
)

# **Step 4: Aggregate results into a structured format**
result = run_counts.group_by(["x", "y", "polarity"]).agg(
pl.struct(["run_length", "occurrences"]).alias("run_counts")
)

result=result.sort('x','y', 'polarity')
print(result)


Which almost works. The output correctly returns the occurances and the lengths for most pixels, which you can see by comparison for the 0,0 pixel:

shape: (1_666_917, 4)
┌──────┬─────┬──────────┬─────────────────┐
│ x ┆ y ┆ polarity ┆ run_counts │
│ --- ┆ --- ┆ --- ┆ --- │
│ u16 ┆ u16 ┆ i16 ┆ list[struct[2]] │
╞══════╪═════╪══════════╪═════════════════╡
│ 0 ┆ 0 ┆ 0 ┆ [{1,8}] │
│ 0 ┆ 0 ┆ 1 ┆ [{1,4}] │
│ 0 ┆ 1 ┆ 0 ┆ [{1,37}] │
│ 0 ┆ 1 ┆ 1 ┆ [{1,30}] │
│ 0 ┆ 2 ┆ 0 ┆ [{1,10}] │
│ … ┆ … ┆ … ┆ … │
│ 1279 ┆ 716 ┆ 1 ┆ [{1,9}] │
│ 1279 ┆ 718 ┆ 0 ┆ [{1,3}] │
│ 1279 ┆ 718 ┆ 1 ┆ [{1,2}] │
│ 1279 ┆ 719 ┆ 0 ┆ [{1,13}] │
│ 1279 ┆ 719 ┆ 1 ┆ [{1,11}] │
└──────┴─────┴──────────┴─────────────────┘


However, for the edge case that a pixel has more than one event per timestamp, it does not work correctly. For example, if we isolate just pixel 1256,519, we have

shape: (88_481_183, 4)
┌──────┬─────┬──────────┬───────────┐
│ x ┆ y ┆ polarity ┆ timestamp │
│ --- ┆ --- ┆ --- ┆ --- │
│ u16 ┆ u16 ┆ i16 ┆ i64 │
╞══════╪═════╪══════════╪═══════════╡
│ 1256 ┆ 519 ┆ 0 ┆ 7616 │
│ 1256 ┆ 519 ┆ 0 ┆ 7616 │
│ 1256 ┆ 519 ┆ 0 ┆ 7616 │
│ 1256 ┆ 519 ┆ 0 ┆ 7616 │
│ 1256 ┆ 519 ┆ 0 ┆ 7617 │
│ … ┆ … ┆ … ┆ … │
│ 1256 ┆ 519 ┆ 0 ┆ 21979614 │
│ 1256 ┆ 519 ┆ 0 ┆ 21979615 │
│ 1256 ┆ 519 ┆ 0 ┆ 21979615 │
│ 1256 ┆ 519 ┆ 0 ┆ 21979615 │
│ 1256 ┆ 519 ┆ 0 ┆ 21979615 │
└──────┴─────┴──────────┴───────────┘


I can manually compute the consecutive occurances and length for this pixel with

df1256=df.filter((pl.col('x').eq(1256)) & (pl.col('y').eq(519)))
non_consecs = df1256.with_columns([
# Create a group id for consecutive timestamps
(pl.col("timestamp").diff() != 1).cum_sum().alias("group")
]).with_columns([
# Create a row number within each group
pl.col("timestamp").rank().over("group").alias("rank")
])

# Count the occurrences and lengths
non_consecs = non_consecs.group_by("group").agg([
pl.col("rank").max().alias("length")
]).group_by("length").agg([
pl.count().alias("occurrences")
])
print(f'non_consecutive_count: {non_consecs}')


Which returns

non_consecutive_count: shape: (9, 2)
┌────────┬─────────────┐
│ length ┆ occurrences │
│ --- ┆ --- │
│ f64 ┆ u32 │
╞════════╪═════════════╡
│ 3.0 ┆ 2941 │
│ 5.0 ┆ 27 │
│ 2.0 ┆ 21964691 │
│ 8.0 ┆ 1 │
│ 6.0 ┆ 13 │
│ 1.0 ┆ 44542445 │
│ 4.0 ┆ 72 │
│ 10.0 ┆ 1 │
│ 7.0 ┆ 2 │
└────────┴─────────────┘


This matches results for that pixel from a much more tedious approach using Numpy. However, my function does not return the correct results for this pixel:

┌──────┬─────┬──────────┬─────────────────────────────────┐
│ x ┆ y ┆ polarity ┆ run_counts │
│ --- ┆ --- ┆ --- ┆ --- │
│ u16 ┆ u16 ┆ i16 ┆ list[struct[2]] │
╞══════╪═════╪══════════╪═════════════════════════════════╡
│ 1256 ┆ 519 ┆ 0 ┆ [{917014,1}, {2657944,1}, … {1… │
└──────┴─────┴──────────┴─────────────────────────────────┘


You can see the counts are very different. I cannot figure out what to change about my function to handle these edge cases correctly, while also still correctly counting for the rest of the pixels. Any help would be greatly appreciated.

Continue reading...