feat(rust): Expose IRFunctionExpr::DynamicPred in the python visitor

[Matt711]

Follows up #26495

Also, I added myself as a code owner for the Python node visitor here, but I haven't gotten notifications of changes (which I need to keep the IR in the GPU engine up-to-date). I may need write-access for GitHub to actually trigger the notifications?

[wence-]

question: where does the predicate that one must evaluate actually live? It seems like it's in pred.pred but that is not exposed anywhere. I'm also not sure it can be because it's an Arc<dyn PredicateExpr>?

[Matt711]

The predicate is simple in this case. For example, col < threshold (top_k) which we know from the sort order.

[ritchie46]

Is this something that should be exposed? I don't think the predicate can be called. In this case the predicate is simple, but that doesn't seem like a guarantee.

[Matt711]

I think we could reconstruct the filter (even for complex predicates) because the unique id lets us match each dynamic_pred to its parent sort-slice. But I would want to avoid this manual reconstruction if possible.

[wence-]

I read through @ritchie46's PR but I think I didn't fully understand the structure of how these things are matched up.

IIUC, we have something like:

    df = pl.DataFrame({"x": [1], "y": [1]})
    plan = df.lazy().with_columns(pl.col.x * pl.col.x).sort("y").head(3)

That produces:

SORT BY [slice: (0, 3, dynamic_pred: id-1)] [col("y")]
   WITH_COLUMNS:
   [[(col("x")) * (col("x"))]] 
    FILTER col("y").dynamic_predicate() # id-1
    FROM
      DF ["x", "y"]; PROJECT */2 COLUMNS

And so the idea here is that you're going to read df in chunks and apply the filter based on things you've already seen. So you read the first chunk, the predicate initially return true until you've "filled up" your slice. Then the next time the predicate runs on the next chunk, it delivers values if they are less than the max in the filled up slice, and so forth.

OK, in that scenario I can see how we can have an expression for the predicate.

But what about (it's not implemented yet) if there was a transformation like:

df = pl.DataFrame({"x": [1], "y": [1]})
plan = df.lazy().with_columns(pl.col.x * pl.col.x).unique("y")

That produced:

UNIQUE BY [col("y"), dynamic_pred: id-1]
  WITH_COLUMNS:
    [[(col("x")) * (col("x"))]]
      FILTER col("y").dynamic_predicate(id-1)
      FROM
        ...

Where in this case the dynamic_predicate is set membership of the already seen values.

[ritchie46]

Indeed. That's how it might be used as well. Another thing we plan to use it for is members of the hash-table in a join. But this is something that will dynamically at runtime be determined. Not something we can statically make a predicate for. Otherwise we would have done that already. :)

[Matt711]

OK so the id is sufficient today to reconstruct the predicate. But as you all said we could set membership or members for the hash table for a join. We could (TODO) tag the predicate with a description (eg. TopK, UniqueMembership, JoinMembership)?

[ritchie46]

I don't think it should be tagged. The id should be sufficient. You would find a dynamicpredicate with id=x in a scan and then the dynamicpredicate setting with id=x in a different IR node. The IR node would indicate what to do. How that is done is left to the implementation.

[Matt711]

Ok cool. I bumped the IR version in 8fcf646. Also gentle reminder about adding me as a code owner for the visitor in the PR description.

[codecov]

Codecov Report

❌ Patch coverage is 0% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.62%. Comparing base (4ed8f61) to head (45c734b).

Files with missing lines	Patch %	Lines
...rates/polars-python/src/lazyframe/visitor/nodes.rs	0.00%	3 Missing ⚠️
.../polars-python/src/lazyframe/visitor/expr_nodes.rs	0.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #27616      +/-   ##
==========================================
+ Coverage   79.63%   81.62%   +1.98%     
==========================================
  Files        1846     1846              
  Lines      256061   256064       +3     
  Branches     3180     3180              
==========================================
+ Hits       203907   209000    +5093     
+ Misses      51325    46235    -5090     
  Partials      829      829              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.