From 1095bcbc58e3261f7bdeb88f6ba2c7473d98c293 Mon Sep 17 00:00:00 2001
From: Billy Lanchantin <william.lanchantin@cargosense.com>
Date: Thu, 16 Jan 2025 15:35:54 -0500
Subject: [PATCH] Fix `ungroup` spec (#1056)

* Add more to the ungroup spec

I believe this fixes dialyzer's warning.

* Add a test to clarify how ranges in ungroup work
---
 lib/explorer/data_frame.ex                |  5 ++++-
 test/explorer/data_frame/grouped_test.exs | 11 +++++++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/lib/explorer/data_frame.ex b/lib/explorer/data_frame.ex
index fe1c99ab5..03ae9a586 100644
--- a/lib/explorer/data_frame.ex
+++ b/lib/explorer/data_frame.ex
@@ -5625,7 +5625,10 @@ defmodule Explorer.DataFrame do
 
   """
   @doc type: :single
-  @spec ungroup(df :: DataFrame.t(), groups_or_group :: column_names() | column_name()) ::
+  @spec ungroup(
+          df :: DataFrame.t(),
+          groups_or_group :: column_names() | column_name() | Range.t() | fun()
+        ) ::
           DataFrame.t()
   def ungroup(df, groups \\ ..)
 
diff --git a/test/explorer/data_frame/grouped_test.exs b/test/explorer/data_frame/grouped_test.exs
index d6fb522a9..a14d16ab8 100644
--- a/test/explorer/data_frame/grouped_test.exs
+++ b/test/explorer/data_frame/grouped_test.exs
@@ -66,6 +66,17 @@ defmodule Explorer.DataFrame.GroupedTest do
       assert DF.groups(df2) == []
     end
 
+    test "ungroup by range", %{df: df} do
+      df1 = DF.group_by(df, ["country", "year"])
+      df2 = DF.ungroup(df1, 1..1)
+
+      # Note: the range selected the column at index 1 of `df.names` to ungroup,
+      # not `df.groups`.
+      assert Enum.take(df.names, 2) == ["year", "country"]
+      assert df2.groups == ["year"]
+      assert DF.groups(df2) == ["year"]
+    end
+
     test "raise error for unknown groups", %{df: df} do
       df1 = DF.group_by(df, ["country", "year"])