cereal data

cereal = pd.read_csv("https://sta663-sp26.github.io/slides/data/cereal.csv"); cereal

groupby

Groups can be created within a DataFrame via groupby() - this returns a DataFrameGroupBy object that contains information about the groups and how to access them.

cereal.groupby("type")

<pandas.api.typing.DataFrameGroupBy object at 0x133472f90>

cereal.groupby("type").groups

{'Cold': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76], 'Hot': [20, 43, 57]}

cereal.groupby("mfr").groups

{'General Mills': [5, 7, 11, 12, 13, 14, 18, 22, 31, 36, 40, 42, 47, 51, 59, 69, 70, 71, 72, 73, 75, 76], 'Kellogg's': [2, 3, 6, 16, 17, 19, 21, 24, 25, 26, 28, 38, 39, 46, 48, 49, 50, 53, 58, 60, 62, 66, 67], 'Maltex': [43], 'Nabisco': [0, 20, 63, 64, 65, 68], 'Post': [9, 27, 29, 30, 32, 33, 34, 37, 52], 'Quaker Oats': [1, 10, 35, 41, 54, 55, 56, 57], 'Ralston Purina': [4, 8, 15, 23, 44, 45, 61, 74]}

groupby and aggregation methods

These groups can be used by the standard aggregation methods (e.g. sum(), mean(), std(), etc.)

cereal.groupby("type").mean()

TypeError: dtype 'str' does not support operation 'mean'

( cereal
  .groupby("type")
  .mean(numeric_only=True)
)

cereal.groupby("mfr").size()

mfr
General Mills     22
Kellogg's         23
Maltex             1
Nabisco            6
Post               9
Quaker Oats        8
Ralston Purina     8
dtype: int64

Selecting groups

Groups can be accessed via get_group()

cereal.groupby("type").get_group("Hot")

cereal.groupby("mfr").get_group("Post")

Iterating over groups

DataFrameGroupBy objects can also be iterated over:

for name, group in cereal.groupby("type"):
  print(f"# {name}\n{group}\n\n")

# Cold
                         name             mfr  ... sugars     rating
0                   100% Bran         Nabisco  ...      6  68.402973
1           100% Natural Bran     Quaker Oats  ...      8  33.983679
2                    All-Bran       Kellogg's  ...      5  59.425505
3   All-Bran with Extra Fiber       Kellogg's  ...      0  93.704912
4              Almond Delight  Ralston Purina  ...      8  34.384843
..                        ...             ...  ...    ...        ...
72                    Triples   General Mills  ...      3  39.106174
73                       Trix   General Mills  ...     12  27.753301
74                 Wheat Chex  Ralston Purina  ...      3  49.787445
75                   Wheaties   General Mills  ...      3  51.592193
76        Wheaties Honey Gold   General Mills  ...      8  36.187559

[74 rows x 6 columns]


# Hot
                      name          mfr type  calories  sugars     rating
20  Cream of Wheat (Quick)      Nabisco  Hot       100       0  64.533816
43                   Maypo       Maltex  Hot       100       3  54.850917
57          Quaker Oatmeal  Quaker Oats  Hot       100       1  50.828392

Aggregation

The aggregate() (or agg()) method can be used to compute summary statistics for each group.

cereal.groupby("mfr").agg("mean")

TypeError: dtype 'str' does not support operation 'mean'

cereal.groupby("mfr").agg("mean", numeric_only = True)

Aggregation by column

cereal.groupby("mfr").agg({
  "calories": ['min', 'max'],
  "sugars":   ['median'],
  "rating":   ['sum', 'count']
})

Named aggregation

It is also possible to use special syntax to aggregate specific columns into a named output column,

cereal.groupby("mfr", as_index=False).agg(
  min_cal = ("calories", "min"),
  max_cal = ("calories", max),
  med_sugar = ("sugars", "median"),
  avg_rating = ("rating", np.mean)
)

polars

Polars is a blazingly fast DataFrame library for manipulating structured data. The core is written in Rust, and available for Python, R and NodeJS.

The goal of Polars is to provide a lightning fast DataFrame library that:

• Utilizes all available cores on your machine.
• Optimizes queries to reduce unneeded work/memory allocations.
• Handles datasets much larger than your available RAM.
• Provides a consistent and predictable API.
• Adheres to a strict schema (data-types should be known before running the query).

import polars as pl
pl.__version__

'1.38.1'

Series

pl.Series("ints", [1, 2, 3, 4, 5])

pl.Series("dbls", [1., 2., 3., 4., 5.])

pl.Series("bools", [True, False, True, False, True])

pl.Series("strs", ["A", "B", "C", "D", "E"])

Missing values

pl.Series("ints", 
  [1, 2, 3, None])

pl.Series("bools", 
  [True, False, True, None])

pl.Series("ints", 
  [1, 2, 3, np.nan])

TypeError: unexpected value while building Series of type Int64; found value of type Float64: NaN

Hint: Try setting `strict=False` to allow passing data with mixed types.

pl.Series("dbls", 
  [1., 2., 3., None])

pl.Series("strs", 
  ["A", "B", "C", None])

pl.Series("dbls", 
  [1., 2., 3., np.nan])

Missing value checking

Checking for missing values can be done via the is_null() method

pl.Series("ints", 
  [1, 2, 3, None]).is_null()

pl.Series("dbls", 
  [1., 2., 3., None]).is_null()

pl.Series("dbls", 
  [1., 2., 3., np.nan]).is_null()

pl.Series("bools", 
  [True, False, True, None]).is_null()

null vs NaN

s_null = pl.Series("x", [1., 2., None])

s_nan = pl.Series("x", [1., 2., np.nan])

s_null.is_null()

s_null.is_nan()

s_null.sum()

3.0

s_null.mean()

1.5

s_null.fill_null(0)

s_nan.is_null()

s_nan.is_nan()

s_nan.sum()

nan

s_nan.mean()

nan

s_nan.fill_nan(0)

DataFrames

Data Frames can be constructed in the same way as Pandas,

df = pl.DataFrame(
  {
    "name":   ["anna","bob","carol", "dave", "erin"],
    "id":     np.random.randint(100, 999, 5),
    "weight": np.random.normal(70, 20, 5),
    "height": np.random.normal(170, 15, 5),
    "date":   pd.date_range(start='2/1/2025', periods=5, freq='D')
  },
  schema_overrides = {"id": pl.UInt16, "weight": pl.Float32}
)
df

Expressions

Polars makes use of lazy evaluation to improve its flexibility and computational performance.

bmi_expr = pl.col("weight") / ((pl.col("height")/100) ** 2)
bmi_expr

This represents a potential computation that can be executed later.

Much of the power of Polars comes from the ability to chain together / compose these expressions.

Contexts

Contexts are the environments in which expressions are evaluated - examples of common contexts include: select, with_columns, filter, and group_by.

df.select(bmi = bmi_expr)

df.with_columns(bmi = bmi_expr)

filter()

df.filter(
  pl.col("height") > 160,
  pl.col("id") < 500
)

df.filter(
  (pl.col("height") > 160) | 
  (pl.col("id") < 500)
)

group_by() & agg()

df.group_by(
  id_range = pl.col("id") - pl.col("id") % 100
).agg(
  pl.len(),
  pl.col("name"),
  bmi_expr.alias("bmi"),
  pl.col("weight", "height").mean().name.prefix("avg_"),
  med_height = pl.col("height").median()
)

More expression expansion

num_cols = pl.col(pl.Float64, pl.Float32)

df.with_columns(
  ((num_cols - num_cols.mean())/num_cols.std()).name.suffix("_std")
)

sort() & sort_by()

sort() sorts a DataFrame by one or more columns. sort_by() is the expression-level equivalent for use within contexts.

df.sort("height", descending=True)

df.sort("id", "height")

df.select(
  pl.col("name", "id", "height").sort_by("height", descending=True)
)

Pivoting - long to wide

df_long

df_long.pivot(
  on="type",
  index=["country","year"],
  values="count"
)

Wide to long (unpivot)

df_wide

df_wide.unpivot(
  on=["1999", "2000"],
  index="country",
  variable_name="year",
  value_name="cases"
)

NYC Taxi Data

df = pl.scan_parquet(
  "~/Scratch/nyctaxi/fix/*_fix.parquet"
)
df

df.select(pl.len()).collect()

Query plan

Just like with DuckDB and sqlite we can ask polars to show us the query plan for a given expression.

print( df.select(pl.len()).explain() )

SELECT [len()]
  Parquet SCAN [/Users/rundel/Scratch/nyctaxi/fix/yellow_tripdata_2020-01_fix.parquet, ... 58 other sources]
  PROJECT */19 COLUMNS
  ESTIMATED ROWS: 377895472

Large lazy queries

zone_lookup = pl.read_csv(
  "https://d37ci6vzurychx.cloudfront.net/misc/taxi_zone_lookup.csv"
).rename(
  {"LocationID": "pickup_zone"}
)

query = (
  df
  .filter(pl.col("trip_distance") > 0)
  .rename({"PULocationID": "pickup_zone"})
  .group_by("pickup_zone")
  .agg(
    num_rides = pl.len(),
    avg_fare_per_mile = (pl.col("fare_amount") / pl.col("trip_distance")).mean().round(2)
  ).join(
    zone_lookup.lazy(),
    on = "pickup_zone",
    how = "left"
  )
  .sort("pickup_zone")
)

Result

query.collect()

pandas equivalent

# pyarrow is required for reading parquet files

zone_lookup_pd = pd.read_csv(
  "https://d37ci6vzurychx.cloudfront.net/misc/taxi_zone_lookup.csv"
).rename(columns={"LocationID": "pickup_zone"})

def pd_query():
    df = pd.read_parquet("~/Scratch/nyctaxi/fix/")
    filtered = (
      df[df["trip_distance"] > 0]
      .rename(columns={"PULocationID": "pickup_zone"})
      .assign(fare_per_mile=lambda x: x["fare_amount"] / x["trip_distance"])
    )
    return (
      filtered
      .groupby("pickup_zone")
      .agg(
        num_rides=("fare_per_mile", "size"),
        avg_fare_per_mile=("fare_per_mile", "mean")
      )
      .round({"avg_fare_per_mile": 2})
      .reset_index()
      .merge(zone_lookup_pd, on="pickup_zone", how="left")
      .sort_values("pickup_zone")
    )

Performance

polars

%timeit query.collect()

1.10 s ± 72 ms per loop (mean ± std. dev. of 3 runs, 1 loop each)

pandas

%timeit pd_query()

150.8 s ± 5.9 s per loop (mean ± std. dev. of 3 runs, 1 loop each)