Declare expression types

In Polars, a function that takes an expression and returns one is a common way to reuse code. A declaration is the type of such an expression. It states the columns that the expression reads, the column that it gives, and whether it works row by row, combines rows, or is a fixed value. You write a declaration once, give it a name, and use the name as the type of a parameter, of a result or of a variable.

A plain expression loses its type

pl.Expr accepts any expression, and it keeps nothing about it. A frame that reads the result then has no known schema, so the checker gives the plain frame and stops tracking its columns:

pricing.pymypy: no errors

1class Trades(Columns):

2 symbol: pl.String

3 price: pl.Float64

4

5

6def discounted(price: pl.Expr) -> pl.Expr:

7 return price * 0.9

8

9

10trades: TypedDataFrame[Trades]

Resulttrades.select(discounted(pl.col("price")))

DataFrame

Declare one column

column(name, dtype) declares an expression that reads one column and gives it back under its own name, row by row. Below, Price is the type of an expression that reads price as a Float64 and gives a Float64 named price. It is the type of the parameter and of the result of discounted:

pricing.pymypy: 3 errors

1class Trades(Columns):

2 symbol: pl.String

3 price: pl.Float64

4 quantity: pl.Int64

5

6

7Price = column("price", pl.Float64)

8

9

10def discounted(price: Price) -> Price:

11 return price * 0.9

12

13

14trades: TypedDataFrame[Trades]

15discounted(pl.col("quantity"))

↑error: Argument "price" does not match its declared type: it uses column "quantity", which the declared type does not allow [expression-fit]

16discounted(pl.col("price").cast(pl.String))

↑error: Argument "price" does not match its declared type: it makes a String column, but the declared type is Float64 [expression-fit]

17discounted(pl.col("price").sum())

↑error: Argument "price" does not match its declared type: it combines rows into one value, but the declared type works row by row [expression-fit]

Resulttrades.select(discounted(pl.col("price")))declared

priceFloat64

Revealed type

TypedDataFrame[TypedDict({'price': Float64}, closed=True), None]

At each call, the checker checks the argument against the declaration:

• pl.col("quantity") reads a column that the declaration does not name.
• The cast gives a String, and the declaration gives a Float64.
• sum() combines the rows into one value, and the declaration works row by row.

The checker also checks the body against the declaration: price * 0.9 must give a Float64 named price.

The frame checks the data types

A call does not know which frame will read its result. So the frame checks that it has the columns that the declaration reads, with their data types. Below, the frame has price as an Int32, and the declaration reads a Float64:

pricing.pymypy: 1 error

1class Sizes(Columns):

2 price: pl.Int32

3

4

5Price = column("price", pl.Float64)

6

7

8def discounted(price: Price) -> Price:

9 return price * 0.9

10

11

12sizes: TypedDataFrame[Sizes]

13sizes.select(discounted(pl.col("price")))

↑error: "price" has no result type at this frame's column types: the declaration gives no result at the data types of this frame [unsupported-dtype-op]

Give a new column

typed_expr(*reads, gives=, named=) declares an expression that reads some columns and gives a new one. Each read is a column(...), often a name that holds one:

flags.pymypy: no errors

1class Trades(Columns):

2 symbol: pl.String

3 quantity: pl.Int64

4

5

6Quantity = column("quantity", pl.Int64)

7BuyFlag = typed_expr(Quantity, gives=pl.Boolean, named="is_buy")

8

9

10def buy_flag(quantity: Quantity) -> BuyFlag:

11 return (quantity > 0).alias("is_buy")

12

13

14trades: TypedDataFrame[Trades]

Resulttrades.with_columns(buy_flag(pl.col("quantity")))declaredadded

symbolStringquantityInt64is_buyBoolean

Revealed type

TypedDataFrame[TypedDict({'symbol': String, 'quantity': Int64, 'is_buy': Boolean}, closed=True), None]

named= gives the name of the result. With one read and no named=, the result has the name of that read. With two reads or more, named= is required.

Accept a family of data types

The data type of a read can be a Polars family, such as NumericType, or a set of data types, such as {pl.Int32, pl.Int64}. The checker then checks the body at each data type of the family, and a frame can have any of them:

flags.pymypy: 1 error

1class Small(Columns):

2 quantity: pl.Int8

3

4

5class Texts(Columns):

6 quantity: pl.String

7

8

9Quantity = column("quantity", NumericType)

10BuyFlag = typed_expr(Quantity, gives=pl.Boolean, named="is_buy")

11

12

13def buy_flag(quantity: Quantity) -> BuyFlag:

14 return (quantity > 0).alias("is_buy")

15

16

17small: TypedDataFrame[Small]

18texts: TypedDataFrame[Texts]

19texts.select(buy_flag(pl.col("quantity")))

↑error: "is_buy" has no result type at this frame's column types: the declaration gives no result at the data types of this frame [unsupported-dtype-op]

Resultsmall.select(buy_flag(pl.col("quantity")))declared

is_buyBoolean

Revealed type

TypedDataFrame[TypedDict({'is_buy': Boolean}, closed=True), None]

quantity > 0 gives a Boolean at each numeric data type, so the body passes. A String is not a number, so the select on texts is an error.

A read of Unchecked, as in column("quantity", Unchecked), accepts every data type. The checker still checks the name of the column and how the expression treats rows. It does not check the data type, so a wrong data type fails only when the code runs.

Combine rows

typed_agg(...) declares an expression that gives one value for each group of rows, as sum() does. In select, the result is one row. In agg, it is one value for each group:

totals.pymypy: no errors

1class Trades(Columns):

2 symbol: pl.String

3 quantity: pl.Int64

4

5

6Quantity = column("quantity", pl.Int64)

7TotalQuantity = typed_agg(Quantity, gives=pl.Int64)

8

9

10def total_quantity(quantity: Quantity) -> TotalQuantity:

11 return quantity.sum()

12

13

14trades: TypedDataFrame[Trades]

Resulttrades.select(total_quantity(pl.col("quantity")))declared

quantityInt64

Revealed type

TypedDataFrame[TypedDict({'quantity': Int64}, closed=True), None]

Resulttrades.group_by("symbol").agg(total_quantity(pl.col("quantity")))declared

symbolStringquantityInt64

Revealed type

TypedDataFrame[TypedDict({'symbol': String, 'quantity': Int64}, closed=True), None]

typed_lit(gives=) declares a fixed value that reads no column, for example pl.lit(1, dtype=pl.Int32). An argument must treat rows as its declaration says: a fixed value or an aggregation does not fit a declaration that works row by row.

A result that depends on the data type

Polars often gives a result whose data type depends on the input. sum() gives an Int64 for an Int8 or an Int16, and keeps the data type of a wider integer. One declaration can give a result for each group of data types: row(...) gives one result for some data types. same is the data type of the column that the expression reads:

totals.pymypy: 1 error

1class Small(Columns):

2 quantity: pl.Int8

3

4

5class Wide(Columns):

6 quantity: pl.Int32

7

8

9Quantity = column("quantity", {pl.Int8, pl.Int16, pl.Int32, pl.Int64})

10SameType = typed_agg(Quantity, gives=same)

11Total = typed_agg(

12 row(column("quantity", {pl.Int8, pl.Int16}), gives=pl.Int64),

13 row(column("quantity", {pl.Int32, pl.Int64}), gives=same),

14)

15

16

17def wrong_total(quantity: Quantity) -> SameType:

18 return quantity.sum()

↑error: This expression does not match its declared type: it makes an Int64 column, but the declared type is Int16 [expression-fit]

19

20

21def total(quantity: Quantity) -> Total:

22 return quantity.sum()

23

24

25small: TypedDataFrame[Small]

26wide: TypedDataFrame[Wide]

Resultsmall.select(total(pl.col("quantity")))declared

quantityInt64

Revealed type

TypedDataFrame[TypedDict({'quantity': Int64}, closed=True), None]

Resultwide.select(total(pl.col("quantity")))declared

quantityInt32

Revealed type

TypedDataFrame[TypedDict(Wide, {'quantity': Int32}, closed=True), None]

SameType promises the input data type back at each integer type, and sum() gives an Int64 at an Int16, so wrong_total is an error. Total states both results. Each row of one declaration reads the same columns, and no two rows cover the same data type.

Accept any column

some_column(region) reads one column whose name each call chooses. With gives=same, the result is that column, with its own data type:

scores.pymypy: 1 error

1class Trades(Columns):

2 symbol: pl.String

3 price: pl.Float64

4

5

6class Readings(Columns):

7 temp: pl.Float32

8

9

10Standardised = typed_expr(some_column(FloatType), gives=same)

11

12

13def zscore(e: Standardised) -> Standardised:

14 return (e - e.mean()) / e.std()

15

16

17trades: TypedDataFrame[Trades]

18readings: TypedDataFrame[Readings]

19zscore(pl.col("price").alias("z"))

↑error: Argument "e" does not match its declared type: it must give back the one column it uses, under the same name [expression-fit]

Resulttrades.select(zscore(pl.col("price")))declared

priceFloat64

Revealed type

TypedDataFrame[TypedDict({'price': Float64}, closed=True), None]

Resultreadings.select(zscore(pl.col("temp")))declared

tempFloat32

Revealed type

TypedDataFrame[TypedDict(Readings, {'temp': Float32}, closed=True), None]

The call gives the name: the result of zscore(pl.col("price")) is named price. An argument must give its column back under the same name, so .alias("z") does not fit. Rename the result after the call instead. In the body, the name is not known yet, so the body cannot give the expression to a frame.

Take a rule from the caller

A schema class can be the reads of a declaration. The argument can then read any column of the class. Below, flagged accepts any frame with the columns of Trades, and any rule over those columns that gives a Boolean named flag:

flags.pymypy: 2 errors

1class Trades(Columns):

2 symbol: pl.String

3 price: pl.Float64

4 quantity: pl.Int64

5

6

7class Flagged(Trades):

8 flag: pl.Boolean

9

10

11Rest = other_columns()

12FlagRule = typed_expr(Trades, gives=pl.Boolean, named="flag")

13

14

15def flagged(

16 df: TypedDataFrame[Trades, Rest],

17 rule: FlagRule,

18) -> TypedDataFrame[Flagged, InferredRemainder]:

19 return infer(df.with_columns(rule))

20

21

22class Blotter(Trades):

23 venue: pl.String

24

25

26blotter: TypedDataFrame[Blotter]

27large = (pl.col("price") * pl.col("quantity") > 1e6).alias("flag")

28result = flagged(blotter, large)

29flagged(blotter, (pl.col("venue") == "X").alias("flag"))

↑error: Argument "rule" does not match its declared type: it uses column "venue", which the declared type does not allow [expression-fit]

30flagged(blotter, pl.lit(True).alias("flag"))

↑error: Argument "rule" does not match its declared type: it is a fixed value, but the declared type works row by row [expression-fit]

Resultresultdeclaredaddedfrom the remainder

symbolStringpriceFloat64quantityInt64flagBooleanvenueString

Revealed type

TypedDataFrame[TypedDict({'symbol': String, 'price': Float64, 'quantity': Int64, 'flag': Boolean, 'venue': String}, closed=True), None]

• The first rule reads price and quantity, which Trades declares. The result keeps the column venue of the caller.
• The second rule reads venue. A caller can pass a frame without venue, so the declaration does not allow it.
• pl.lit(True) is a fixed value, and the declaration works row by row.

Variables and plain expressions

The checker checks an annotated variable in the same way as an argument. A plain pl.Expr fits no declaration, because the checker does not know its columns:

flags.pymypy: 2 errors

1Price = column("price", pl.Float64)

2Expensive = typed_expr(Price, gives=pl.Boolean, named="expensive")

3

4rule: Expensive = (pl.col("price") > 100.0).alias("expensive")

5wrong: Expensive = pl.col("price").alias("expensive")

↑error: This expression does not match its declared type: it makes a Float64 column, but the declared type is Boolean [expression-fit]

6

7

8def expensive(price: Price) -> Expensive:

9 return (price > 100.0).alias("expensive")

10

11

12plain: pl.Expr

13expensive(plain)

↑error: Argument "price" does not match its declared type: this value is not a typed expression, so its columns and types are unknown [expression-fit]

↑note: set allow_unproven = "asserted" in [tool.astrolabe] to allow this, or set [tool.astrolabe.allow_unproven_overrides] fit = "asserted"

The error on a plain expression is an unproven result: it passes at the level asserted of the setting allow_unproven. See Configuration.

Let the body decide the result

InferredExpr as a return type lets the checker compute the type of the result from the body. See Inferred return types. The result keeps the declarations of the parameters, and the frame that reads it checks its data types against them. A parameter has no inferred form, because the body must work for each argument that a caller can pass.

Read the columns of a schema

column_of(Trades, "price") reads price at the data type that Trades gives it. The declaration then follows the class when that data type changes. subset_of(Trades, "price", "quantity") reads several columns.

Data types with parameters

A read covers kinds: data types with their parameters taken away. column("t", pl.Datetime("ms")) reads a Datetime column of any time unit and any time zone, and the checker checks the argument at each of them. A fixed result keeps its parameters: gives=pl.Decimal(38, 2) gives that one data type. wrap(pl.List, same) gives a List of the data type of the column, as implode() does.

Rules of a declaration

The checker reports a declaration that breaks a rule of form, where you write it:

features.pymypy: 2 errors

1import polars as pl

2

3from astrolabe import column, typed_expr

4

5NoRead = typed_expr(gives=pl.Int64, named="n")

↑error: Malformed declaration: no column is read. Read one, name a typed_agg, or use typed_lit [malformed-declaration]

6Pair = typed_expr(column("price", pl.Float64), column("quantity", pl.Int64), gives=pl.Float64)

↑error: Malformed declaration: named= is required with two or more reads [malformed-declaration]

To apply one expression to each column that a selector picks, declare a selection: see Selections.