Skip to content

Expr.map_key builds its key literal without regard to the map's key type #295

Description

@nielspardon

Expr.map_key(key) turns the Python value into a literal with Expr._coerce, so an int always becomes i64 and a str always becomes string. The literal is built without regard to the key type of the map it indexes.

#294 makes infer_expression_type check the literal against the map's key type, which matches substrait-go's MapKeyRef.GetType. With that check in place, a lookup on a non-i64 integer-keyed map raises:

import substrait.dataframe as sub
from substrait.builders.type import map, i32, string, named_struct, struct

ns = named_struct(names=["m"], struct=struct(types=[map(i32(), string())], nullable=False))
sub.read_named_table("t", ns).select(sub.col("m").map_key(7).alias("v")).to_plan()
# ValueError: Map key literal type does not match map key type

On main the call doesn't raise, but the plan carries an i64 key literal for an i32 key, which substrait-go rejects. map_key should build the literal using the parent map's key type.

Measured on the #294 head at 77b2a02 (raises) and on main (no error; the reference is typed as map).

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions