Array Processing with Polars

import polars as pl
from datetime import date

df = pl.DataFrame(
    {
        "name": [
            "Alice Archer",
            "Ben Brown",
            "Chloe Cooper",
            "Daniel Donovan",
        ],
        "birthdate": [
            date(1997, 1, 10),
            date(1985, 2, 15),
            date(1983, 3, 22),
            date(1981, 4, 30),
        ],
        # in kilos
        "weight": [57.9, 72.5, 53.6, 83.1],
        # in metres
        "height": [1.56, 1.77, 1.65, 1.75],
    }
)

print(df)

shape: (4, 4)
┌────────────────┬────────────┬────────┬────────┐
│ name           ┆ birthdate  ┆ weight ┆ height │
│ ---            ┆ ---        ┆ ---    ┆ ---    │
│ str            ┆ date       ┆ f64    ┆ f64    │
╞════════════════╪════════════╪════════╪════════╡
│ Alice Archer   ┆ 1997-01-10 ┆ 57.9   ┆ 1.56   │
│ Ben Brown      ┆ 1985-02-15 ┆ 72.5   ┆ 1.77   │
│ Chloe Cooper   ┆ 1983-03-22 ┆ 53.6   ┆ 1.65   │
│ Daniel Donovan ┆ 1981-04-30 ┆ 83.1   ┆ 1.75   │
└────────────────┴────────────┴────────┴────────┘

csv_file = "/tmp/sample.csv"
df.write_csv(csv_file)
df_csv = pl.read_csv(csv_file, try_parse_dates=True)
print(df_csv)

shape: (4, 4)
┌────────────────┬────────────┬────────┬────────┐
│ name           ┆ birthdate  ┆ weight ┆ height │
│ ---            ┆ ---        ┆ ---    ┆ ---    │
│ str            ┆ date       ┆ f64    ┆ f64    │
╞════════════════╪════════════╪════════╪════════╡
│ Alice Archer   ┆ 1997-01-10 ┆ 57.9   ┆ 1.56   │
│ Ben Brown      ┆ 1985-02-15 ┆ 72.5   ┆ 1.77   │
│ Chloe Cooper   ┆ 1983-03-22 ┆ 53.6   ┆ 1.65   │
│ Daniel Donovan ┆ 1981-04-30 ┆ 83.1   ┆ 1.75   │
└────────────────┴────────────┴────────┴────────┘

Note that nothing about the data types is written to the CSV files so it may be that information is lost in this process. We can see this if we tell the reader not to parse the dates.

pl.read_csv(csv_file)

shape: (4, 4)
┌────────────────┬────────────┬────────┬────────┐
│ name           ┆ birthdate  ┆ weight ┆ height │
│ ---            ┆ ---        ┆ ---    ┆ ---    │
│ str            ┆ str        ┆ f64    ┆ f64    │
╞════════════════╪════════════╪════════╪════════╡
│ Alice Archer   ┆ 1997-01-10 ┆ 57.9   ┆ 1.56   │
│ Ben Brown      ┆ 1985-02-15 ┆ 72.5   ┆ 1.77   │
│ Chloe Cooper   ┆ 1983-03-22 ┆ 53.6   ┆ 1.65   │
│ Daniel Donovan ┆ 1981-04-30 ┆ 83.1   ┆ 1.75   │
└────────────────┴────────────┴────────┴────────┘

We can refer to columns of a data frame

pl.col("weight")

<Expr ['col("weight")'] at 0x7FF09ACE4150>

and we can build expressions from them

pl.col("weight") / (pl.col("height") ** 2)

<Expr ['[(col("weight")) / (col("heigh…'] at 0x7FF09ACE4250>

You might recognise this expression as the Body mass index, used as a "convenient rule of thumb used to broadly categorize a person as underweight, normal or overweight".

We can use such expressions in Polars contexts, for example to select columns by expression.

result = df.select(
    pl.col("name"),
    pl.col("birthdate").dt.year().alias("birth_year"),
    (pl.col("weight") / (pl.col("height") ** 2)).alias("bmi"),
)
print(result)

shape: (4, 3)
┌────────────────┬────────────┬───────────┐
│ name           ┆ birth_year ┆ bmi       │
│ ---            ┆ ---        ┆ ---       │
│ str            ┆ i32        ┆ f64       │
╞════════════════╪════════════╪═══════════╡
│ Alice Archer   ┆ 1997       ┆ 23.791913 │
│ Ben Brown      ┆ 1985       ┆ 23.141498 │
│ Chloe Cooper   ┆ 1983       ┆ 19.687787 │
│ Daniel Donovan ┆ 1981       ┆ 27.134694 │
└────────────────┴────────────┴───────────┘

• Note how we have selected two columns and one expression.
• We have transformed birthdate to extract the year from the date and renamed the column to birth_year
• We have given the expression calculated from the columns weight and height the name bmi
• Note that the original dataframe remains unchanged.

Polars allows one expression to be used for multiple expression. This is called expression expansion.

result = df.select(
    pl.col("name"),
    (pl.col("weight", "height") * 0.95).round(2).name.suffix("-5%"),
)
print(result)

shape: (4, 3)
┌────────────────┬───────────┬───────────┐
│ name           ┆ weight-5% ┆ height-5% │
│ ---            ┆ ---       ┆ ---       │
│ str            ┆ f64       ┆ f64       │
╞════════════════╪═══════════╪═══════════╡
│ Alice Archer   ┆ 55.0      ┆ 1.48      │
│ Ben Brown      ┆ 68.88     ┆ 1.68      │
│ Chloe Cooper   ┆ 50.92     ┆ 1.57      │
│ Daniel Donovan ┆ 78.94     ┆ 1.66      │
└────────────────┴───────────┴───────────┘

• One expression applies to both columns, weight and height
• Each column is
  1. first multiplied by 0.95
  2. then rounded to 2 decimal places
  3. then the column is renamed by suffixing its name with -5%

Sometimes we want to add columns to our date frame rather than select columns from it. For this we use with_columns

result = df.with_columns(
    birth_year=pl.col("birthdate").dt.year(),
    bmi=pl.col("weight") / (pl.col("height") ** 2),
)
print(result)

shape: (4, 6)
┌────────────────┬────────────┬────────┬────────┬────────────┬───────────┐
│ name           ┆ birthdate  ┆ weight ┆ height ┆ birth_year ┆ bmi       │
│ ---            ┆ ---        ┆ ---    ┆ ---    ┆ ---        ┆ ---       │
│ str            ┆ date       ┆ f64    ┆ f64    ┆ i32        ┆ f64       │
╞════════════════╪════════════╪════════╪════════╪════════════╪═══════════╡
│ Alice Archer   ┆ 1997-01-10 ┆ 57.9   ┆ 1.56   ┆ 1997       ┆ 23.791913 │
│ Ben Brown      ┆ 1985-02-15 ┆ 72.5   ┆ 1.77   ┆ 1985       ┆ 23.141498 │
│ Chloe Cooper   ┆ 1983-03-22 ┆ 53.6   ┆ 1.65   ┆ 1983       ┆ 19.687787 │
│ Daniel Donovan ┆ 1981-04-30 ┆ 83.1   ┆ 1.75   ┆ 1981       ┆ 27.134694 │
└────────────────┴────────────┴────────┴────────┴────────────┴───────────┘

• Here the original columns remain unchanged and we add two columns.
• Note, again, that the original data frame remains unchanged.
• Note how the new column birth_year is of type i32. This takes half the space of an i64 and is more that sufficient to store calendar years.
• We use a more convenient "assignment form" to make the new columns, rather than alias. If the new column name contains, for example, a space then we would need to use alias.

Sometimes we need to filter the rows of a data frame. For example

result = df.filter(pl.col("birthdate").dt.year() < 1990)
print(result)

shape: (3, 4)
┌────────────────┬────────────┬────────┬────────┐
│ name           ┆ birthdate  ┆ weight ┆ height │
│ ---            ┆ ---        ┆ ---    ┆ ---    │
│ str            ┆ date       ┆ f64    ┆ f64    │
╞════════════════╪════════════╪════════╪════════╡
│ Ben Brown      ┆ 1985-02-15 ┆ 72.5   ┆ 1.77   │
│ Chloe Cooper   ┆ 1983-03-22 ┆ 53.6   ┆ 1.65   │
│ Daniel Donovan ┆ 1981-04-30 ┆ 83.1   ┆ 1.75   │
└────────────────┴────────────┴────────┴────────┘

• Note how filter takes a predicate on the column birthdate.
• Such filters can be more complex Boolean expressions over multiple columns.

• Rather than writing Boolean ∧, we can write multiple filters.

  result = df.filter(
      pl.col("birthdate").is_between(date(1982, 12, 31), date(1996, 1, 1)),
      pl.col("height") > 1.7,
  )
  print(result)
  

  shape: (1, 4)
  ┌───────────┬────────────┬────────┬────────┐
  │ name      ┆ birthdate  ┆ weight ┆ height │
  │ ---       ┆ ---        ┆ ---    ┆ ---    │
  │ str       ┆ date       ┆ f64    ┆ f64    │
  ╞═══════════╪════════════╪════════╪════════╡
  │ Ben Brown ┆ 1985-02-15 ┆ 72.5   ┆ 1.77   │
  └───────────┴────────────┴────────┴────────┘
  

  • Note that you can indicate to is_between whether to include either endpoint or both by using the closed optional argument.

We have had selecting columns, adding columns and filtering rows. Another commonly needed operation on data frames is grouping. For this we use group_by

Suppose we would like to see how many rows in the data frame have a birth date in each decade. We can do it like this.

result = df.group_by(
    decade = (pl.col("birthdate").dt.year() // 10 * 10),
    maintain_order=True,
).len()
print(result)

shape: (2, 2)
┌────────┬─────┐
│ decade ┆ len │
│ ---    ┆ --- │
│ i32    ┆ u32 │
╞════════╪═════╡
│ 1990   ┆ 1   │
│ 1980   ┆ 3   │
└────────┴─────┘

1. We take the birthdate column
2. We extract its year
3. We divide its year by 10 using integer division (//). This drops the fractional part.
4. Now we multiply by 10 again. We are left with the decade.
5. We name it decade.
6. We use maintain_order to keep the order of the decades consistent with the dates in the data frame. This is slower but more convenient to read.
7. We take the length, i.e. the number of rows in each decade.

Now that we have grouped the rows into decades, we can take aggregations over each group.

result = df.group_by(
    decade=(pl.col("birthdate").dt.year() // 10 * 10),
    maintain_order=True,
).agg(
    pl.len().alias("sample_size"),
    pl.col("weight").mean().round(2).alias("avg_weight"),
    pl.col("height").max().alias("tallest"),
)
print(result)

shape: (2, 4)
┌────────┬─────────────┬────────────┬─────────┐
│ decade ┆ sample_size ┆ avg_weight ┆ tallest │
│ ---    ┆ ---         ┆ ---        ┆ ---     │
│ i32    ┆ u32         ┆ f64        ┆ f64     │
╞════════╪═════════════╪════════════╪═════════╡
│ 1990   ┆ 1           ┆ 57.9       ┆ 1.56    │
│ 1980   ┆ 3           ┆ 69.73      ┆ 1.77    │
└────────┴─────────────┴────────────┴─────────┘

• Here we take
  • the "length" of each group and name it sample_size
  • the weight which we average (mean), round to 2 decimal places and name avg_weight
  • the height which we take the largest of (max) and name tallest
• Note that we can not use the assignment form for agg; we must use alias.

We can build more complex queries by chaining expressions and contexts.

result = (
    df.with_columns(
        pl.col("name").str.split(by=" ").list.first(),
        decade=(pl.col("birthdate").dt.year() // 10 * 10),
    )
    .select(pl.all().exclude("birthdate"))
    .group_by(pl.col("decade"), maintain_order=True)
    .agg(
        pl.col("name"),
        pl.col("weight", "height").mean().round(2).name.prefix("avg_"),
    )
)
print(result)

shape: (2, 4)
┌────────┬────────────────────────────┬────────────┬────────────┐
│ decade ┆ name                       ┆ avg_weight ┆ avg_height │
│ ---    ┆ ---                        ┆ ---        ┆ ---        │
│ i32    ┆ list[str]                  ┆ f64        ┆ f64        │
╞════════╪════════════════════════════╪════════════╪════════════╡
│ 1990   ┆ ["Alice"]                  ┆ 57.9       ┆ 1.56       │
│ 1980   ┆ ["Ben", "Chloe", "Daniel"] ┆ 69.73      ┆ 1.72       │
└────────┴────────────────────────────┴────────────┴────────────┘

• We extract the first name of each name by string manipulation.
• The name column is aggregated into a list of names.

To conclude, we will look at two ways of combining data frames.

First of all we make a new data frame and then join it with the first one.

df2 = pl.DataFrame(
    {
        "name": [
            "Ben Brown",
            "Daniel Donovan",
            "Alice Archer",
            "Chloe Cooper",
        ],
        "parent": [True, False, False, False],
        "siblings": [1, 2, 3, 4],
    }
)
print(df2)

shape: (4, 3)
┌────────────────┬────────┬──────────┐
│ name           ┆ parent ┆ siblings │
│ ---            ┆ ---    ┆ ---      │
│ str            ┆ bool   ┆ i64      │
╞════════════════╪════════╪══════════╡
│ Ben Brown      ┆ true   ┆ 1        │
│ Daniel Donovan ┆ false  ┆ 2        │
│ Alice Archer   ┆ false  ┆ 3        │
│ Chloe Cooper   ┆ false  ┆ 4        │
└────────────────┴────────┴──────────┘

We see the same four names with additional columns. Now let's join this data frame to the earlier one.

print(df.join(df2, on="name", how="left"))

shape: (4, 6)
┌────────────────┬────────────┬────────┬────────┬────────┬──────────┐
│ name           ┆ birthdate  ┆ weight ┆ height ┆ parent ┆ siblings │
│ ---            ┆ ---        ┆ ---    ┆ ---    ┆ ---    ┆ ---      │
│ str            ┆ date       ┆ f64    ┆ f64    ┆ bool   ┆ i64      │
╞════════════════╪════════════╪════════╪════════╪════════╪══════════╡
│ Alice Archer   ┆ 1997-01-10 ┆ 57.9   ┆ 1.56   ┆ false  ┆ 3        │
│ Ben Brown      ┆ 1985-02-15 ┆ 72.5   ┆ 1.77   ┆ true   ┆ 1        │
│ Chloe Cooper   ┆ 1983-03-22 ┆ 53.6   ┆ 1.65   ┆ false  ┆ 4        │
│ Daniel Donovan ┆ 1981-04-30 ┆ 83.1   ┆ 1.75   ┆ false  ┆ 2        │
└────────────────┴────────────┴────────┴────────┴────────┴──────────┘

We get a combination of the two data frames.

There are many ways to join for dealing with, e.g cases where a name is present in one data frame but not in the other or a name occurs multiple times in one or other data frame. It is worth taking the time to experiment with them. Here is a link to the Monopoly data used in the examples.

More straightforwardly than joining, we can just concatenate data frames, horizontally or vertically.

df3 = pl.DataFrame(
    {
        "name": [
            "Ethan Edwards",
            "Fiona Foster",
            "Grace Gibson",
            "Henry Harris",
        ],
        "birthdate": [
            date(1977, 5, 10),
            date(1975, 6, 23),
            date(1973, 7, 22),
            date(1971, 8, 3),
        ],
        "weight": [67.9, 72.5, 57.6, 93.1],
        "height": [1.76, 1.6, 1.66, 1.8],
    }
)

print(pl.concat([df, df3], how="vertical"))

shape: (8, 4)
┌────────────────┬────────────┬────────┬────────┐
│ name           ┆ birthdate  ┆ weight ┆ height │
│ ---            ┆ ---        ┆ ---    ┆ ---    │
│ str            ┆ date       ┆ f64    ┆ f64    │
╞════════════════╪════════════╪════════╪════════╡
│ Alice Archer   ┆ 1997-01-10 ┆ 57.9   ┆ 1.56   │
│ Ben Brown      ┆ 1985-02-15 ┆ 72.5   ┆ 1.77   │
│ Chloe Cooper   ┆ 1983-03-22 ┆ 53.6   ┆ 1.65   │
│ Daniel Donovan ┆ 1981-04-30 ┆ 83.1   ┆ 1.75   │
│ Ethan Edwards  ┆ 1977-05-10 ┆ 67.9   ┆ 1.76   │
│ Fiona Foster   ┆ 1975-06-23 ┆ 72.5   ┆ 1.6    │
│ Grace Gibson   ┆ 1973-07-22 ┆ 57.6   ┆ 1.66   │
│ Henry Harris   ┆ 1971-08-03 ┆ 93.1   ┆ 1.8    │
└────────────────┴────────────┴────────┴────────┘