use polars_core::df;
use polars_core::prelude::*;
use polars_lazy::prelude::*;
use polars_sql::*;

#[test]
fn test_distinct_on() {
    let df = df! {
      "Name" => ["Bob", "Pete", "Pete", "Pete", "Martha", "Martha"],
      "Record Date" => [1, 1, 2, 4, 1, 3],
      "Score" => [8, 2, 9, 3, 2, 6]
    }
    .unwrap()
    .lazy();
    let mut ctx = SQLContext::new();

    ctx.register("df", df.clone());
    let sql = r#"
  SELECT DISTINCT ON ("Name")
      "Name",
      "Record Date",
      "Score"
  FROM
      df
  ORDER BY
      "Name",
      "Record Date" DESC;"#;
    let lf = ctx.execute(sql).unwrap();
    let actual = lf.collect().unwrap();
    let expected = df
        .sort_by_exprs(
            vec![col("Name"), col("Record Date")],
            vec![false, true],
            true,
            false,
        )
        .groupby_stable(vec![col("Name")])
        .agg(vec![col("*").first()]);
    let expected = expected.collect().unwrap();
    assert!(actual.frame_equal(&expected))
}
