Abstract polars function expression nodes to ensure they are serializable

[pentschev]

Description

Use Enums to define Python types as references to polars.polars._expr_nodes.*Function as to ensure cudf_polars.dsl.expressions specializations of Expr are serializable.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[wence-]

Two suggestions

[wence-]

suggestion: use IntEnum.

Also, let's nest this class inside the BooleanFunction class below, and just call it Name.

[pentschev]

Done in 3f69d98 .

[wence-]

I think this doesn't return a polars object, but one of ours. Also, because we're calling getattr, that suggests we should make this a classmethod

@classmethod
def from_polars(cls, obj: pl_expr.BooleanFunction) -> Self:
    function, name = str(obj).split(".", maxsplit=1)
    if function != "BooleanFunction":
        raise ValueError("BooleanFunction required")
    return getattr(cls, name)

[pentschev]

Yeah, I totally mix everything up here, thanks for catching. Done in be7fa52 .

[wence-]

nit: Because we have enums, we should now use is for comparison.

[pentschev]

Done in 7d009f9 .

[wence-]

Thanks Peter, I think this makes sense. Should we add a test that these things are now pickleable, or does that need other bits and pieces?

[pentschev]

Thanks Peter, I think this makes sense. Should we add a test that these things are now pickleable, or does that need other bits and pieces?

Thanks for the review Lawrence. I agree, tests are a good idea. I've made a small improvement in 9b54437 that I discovered could be better while writing some tests in 5a04207. Those tests check that serialization of the Name classes as well as the from_polars method indeed can be done correctly, but they do not check for the full serialization pipeline with Translator. Adding standalone tests for the full Translator pipeline seems significantly more complex, the existing tests seem to all use a GPUEngine for that which is, AFAIU, the only realistic way of using Translator, such cases will then be addressed in #17364 where I'll test the entire suite with the Distributed scheduler which implies serialization of the objects handled in this PR. With that said, do you think addressing the full pipeline in #17364 via the Distributed scheduler is sufficient or do you think adding some sort of mock testing covering the full Translator pipeline with serialization is necessary?

[wence-]

Some small further suggestions

[wence-]

Let's abstract this into a reusable fixture:

@pytest.fixture(params=["BooleanFunction", "StringFunction", "TemporalFunction"])
def function(request):
    return request.param

[pentschev]

Good idea, done in fcf820f . I've kept the module references instead of using strings and resolved them in the tests with __name__, let me know if you have a strong preference for strings instead.

[wence-]

No, that's fine.

[wence-]

/merge