Used TypeDict for CachingVisitor.state

[TomAugspurger]

This changes CachingVisitor to used a TypedDict for state. Previously, we used a Mapping[str, Any], which had two problems:

1. Typos in the key names can cause unexpected KeyErrors
2. The Any type for the values of state means that, at least by default, you won't get any type checking on things using values looked up from state.

Unit tests should catch all these, but this can provide some earlier feedback on any issues.

Now, we'll declare a State TypedDict with some specific fields for each transformer that uses CachingVisitor.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[TomAugspurger]

FYI @rjzamora this will have some overlap with #19130. Any new keys we add to state will need a field in CachingVisitorState.

[TomAugspurger]

Hmm, failing on 3.10: https://github.com/rapidsai/cudf/actions/runs/15594634303/job/43923137155?pr=19135#step:10:7551

 │ │ TypeError: cannot inherit from both a TypedDict type and a non-TypedDict base class

This might need to wait until we drop Python 3.10.

[TomAugspurger]

Just seeing if CI passes with TypedDict from typing_extensions. The docs say this should work. If it does, we can discuss adding that as a dependency for python=3.10. Maybe not worth it.

[TomAugspurger]

I'd like to add a dependency on typing-extensions for just python 3.10. For 3.11 and newer, we don't need typing-extensions. Here's what I think is the desired output:

diff --git a/conda/recipes/cudf-polars/recipe.yaml b/conda/recipes/cudf-polars/recipe.yaml
index 85740e5930..8133c7e8a5 100644
--- a/conda/recipes/cudf-polars/recipe.yaml
+++ b/conda/recipes/cudf-polars/recipe.yaml
@@ -50,6 +50,7 @@ requirements:
     - python
     - pylibcudf =${{ version }}
     - polars >=1.24,<1.31
+    - typing-extensions  # py==310
     - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }}
   ignore_run_exports:
     by_name:
diff --git a/python/cudf_polars/pyproject.toml b/python/cudf_polars/pyproject.toml
index b3d28e027d..9589a8a3f9 100644
--- a/python/cudf_polars/pyproject.toml
+++ b/python/cudf_polars/pyproject.toml
@@ -21,6 +21,7 @@ requires-python = ">=3.10"
 dependencies = [
     "polars>=1.25,<1.31",
     "pylibcudf==25.8.*,>=0.0.0a0",
+    "typing-extensions; python_version < '3.11'",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 classifiers = [
     "Intended Audience :: Developers",

I'd hoped that 4ed5168 would do the trick, but that gives the same output as what's on 25.08 (no change).

[gforsyth]

Hey @TomAugspurger -- I'm not sure if the dependency file generator knows how to handle version-specific dependencies in a pyproject.toml.

I think it's reasonable to do something like:

@@ -834,13 +834,11 @@ dependencies:
         packages:
           - polars>=1.25,<1.31
     specific:
-      - output_types: [conda, requirements, pyproject]
+      - output_types: pyproject
         matrices:
-          - matrix: {py: "3.10"}
-            packages:
-              - typing-extensions
           - matrix: null
-            packages: null
+            packages:
+              - "typing-extensions; python_version < '3.11'"
   run_cudf_polars_experimental:
     common:
       - output_types: [conda, requirements, pyproject]

The other spot where typing_extensions would need to show up is in the conda environment files, but it seems to already be there.

As for the rattler-build recipes, those aren't covered (yet) by dependencies.yaml, so you can go ahead and make the necessary changes directly.

[gforsyth]

Packaging and CI looks good!

[TomAugspurger]

@rjzamora do you have any concerns with this? This adds a bit more friction to putting something in state (you have to declare it in CachingVisitorState), but I think the benefit is worth it.

[rjzamora]

@rjzamora do you have any concerns with this? This adds a bit more friction to putting something in state (you have to declare it in CachingVisitorState), but I think the benefit is worth it.

Sorry I missed this earlier. This PR probably moves things in the right direction, so I think the changes are fine.

With that said, I do think it's pretty confusing to use the same CachingVisitorState type for different kinds of visitors that may need to track very different kinds of state. I think the approach taken here is to just include "everything" in a single CachingVisitorState class. I suppose it would be tricky to have separate DecomposeExprState(CachingVisitorState) and LowerIRState(CachingVisitorState) classes, for example?

[TomAugspurger]

I felt that awkwardness too. It's identical to what we have today, it's just that declaring the types makes it obvious that we're putting all kinds of things in there. I'll see what I can do.

[TomAugspurger]

I pushed some changes that seem to do what we want. It's a bit complicated... but I think worth it. Now we have a LowerIRTransformer type:

LowerIRTransformer: TypeAlias = GenericTransformer[
    "ir.IR", "tuple[ir.IR, MutableMapping[ir.IR, PartitionInfo]]", LowerIRState
]

and we know that inside lower_ir_node we always have a config_options available in the state, and nothing else.

If folks are OK with this overall I'll clean up a couple things (docs, where we define these Transformer subtypes) and ping for another round of reviews.

[wence-]

I'm pretty negative on this change. I would be OK if one could define the protocol for the state near the point of use and one was just generic (I think invariantly) over the parameter. But this is not great. I need to now have globally unique names for the state dictionary, and I feel like you were probably struggling for good ones.

[TomAugspurger]

I feel like you were probably struggling for good ones

What gave that away? :)

Just to make sure we're on the same page about what we want:

1. Visitor-specific state is a good thing. It lets users of the state know precisely what's available to them.
2. We'd like to define the state variables locally to where they're inserted into the CachingVisitor (TBD whether that type needs to be visible to the user of the state as well, but these are often close to each other).
3. Anything that requires enumerating the list of states somewhere is a non-starter.

I think this ugly union came about from trying to get GenericTransformer to type check successfully. I'll look into this a bit more (I haven't ever worked with covariance / contravariance before)

[TomAugspurger]

729b0b8 is hopefully closer to what we want. At a minimum, it's gotten rid of this Union type.

If that looks OK, then I can move things to where we want them. We just need to figure out where to define the transformers (ExprTransformer, LowerIRTransformer, ...) and their associated states. LMK if you have a preference on where things end up at (in cudf_polars.typing / where they're used / somewhere else).

[wence-]

I'll look into this a bit more (I haven't ever worked with covariance / contravariance before)

I wrote a thing (#17016 (comment)), I should probably make a web page...

[TomAugspurger]

Thanks, reading through that now

73fae84 has the proposed reorganization. The basic rules are to 1.) put the GenericTransformer subtype next to where they're used, 2.) put the state declaration next to the CachingVistor / GenericTransformer subtype, and 3.) prefer the name State for the state dict, unless you need to avoid a name collision (which happens in to_ast.py)

[rjzamora]

If folks are OK with this overall I'll clean up a couple things (docs, where we define these Transformer subtypes) and ping for another round of reviews.

I'll defer to you on whether the extra effort is worth it, and I'll be happy to review.

[wence-]

I think if we have StateT_co (and the type parameter for GenericTransformer should be covariant), then we don't need StateT (all the specialisations should use StateT_co), otherwise if you annotate something with StateT you can provide a GenericTransformer that has a subtype of the bound StateT.

[TomAugspurger]

Done in a19936a. I thought I had tried that, but apparently not. And I really am reading through #17016 (comment) now to understand the difference :)