from __future__ import annotations

import pytest

from daft.expressions import col
from tests.conftest import assert_df_equals


@pytest.mark.parametrize(
    "keys",
    [
        pytest.param(["Borough"], id="NumGroupByKeys:1"),
        pytest.param(["Borough", "Complaint Type"], id="NumGroupByKeys:2"),
    ],
)
def test_distinct_all_columns(daft_df, service_requests_csv_pd_df, repartition_nparts, keys, with_morsel_size):
    """Sums across groups."""
    daft_df = daft_df.repartition(repartition_nparts).select(*[col(k) for k in keys]).distinct()

    service_requests_csv_pd_df = (
        service_requests_csv_pd_df.groupby(keys).sum("Unique Key").reset_index().drop("Unique Key", axis=1)
    )
    daft_pd_df = daft_df.to_pandas()
    assert_df_equals(daft_pd_df, service_requests_csv_pd_df, sort_key=keys)
