code: def test_agg_after_head() -> None:
      a = [1, 1, 1, 2, 2, 3, 3, 3, 3]
  
      df = pl.DataFrame(
          {"a": a, "b": pl.arange(1, len(a) + 1, eager=True)}
      )
  
      expected = pl.DataFrame({"a": [1, 2, 3], "b": [6, 9, 21]})
  
      for maintain_order in [True, False]:
          out = df.group_by("a", maintain_order=maintain_order).agg(
              [pl.col("b").head(3).sum()]
          )
  
          if not maintain_order:
              out = out.sort("a")
  
          assert_frame_equal(out, expected)
  
imports: from __future__ import annotations
  from datetime import datetime, timedelta
  from typing import Any
  import numpy
  import pandas
  import polars
  from polars.testing import assert_frame_equal
label_desc: Use image node_id 3 for calling a global function with example usage: test_agg_after_head() without return types