Minor clean up prior to a new release

CHANGELOG.md

@@ -9,6 +9,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added `Explorer.DataFrame.partition_by/2`, which splits a dataframe
+  into one dataframe per distinct combination of values in the
+  given columns [#1163](https://github.com/elixir-explorer/explorer/pull/1163)
 - Added `Explorer.Series.rle_id/1`, which assigns a run-length encoding (RLE) ID
   to each row, incrementing every time the value changes from one row to the next.
 - Support for Nx 1.0 [#1161](https://github.com/elixir-explorer/explorer/pull/1161)

lib/explorer/data_frame.ex

@@ -541,10 +541,16 @@ defmodule Explorer.DataFrame do
           opts :: Keyword.t()
         ) ::
           DataFrame.t()
-  def from_query!(conn, query, params, opts \\ []) do
-    case from_query(conn, query, params, opts) do
-      {:ok, df} -> df
-      {:error, error} -> raise error
+  if Code.ensure_loaded?(Adbc) do
+    def from_query!(conn, query, params, opts \\ []) do
+      case from_query(conn, query, params, opts) do
+        {:ok, df} -> df
+        {:error, error} -> raise error
+      end
+    end
+  else
+    def from_query!(_conn, _query, _params, _opts) do
+      raise "you must install :adbc as a dependency in order to use from_query!/3"
     end
   end
 
@@ -3617,35 +3623,61 @@ defmodule Explorer.DataFrame do
   def arrange_with(df, fun, opts \\ []), do: sort_with(df, fun, opts)
 
   @doc """
-  Splits the dataframe into one dataframe per distinct combination of `columns`.
+  Splits the dataframe into a list of dataframes, one for each distinct
+  combination of values in the given columns.
+
+  Every row goes to exactly one dataframe, with `nil` treated like any
+  other value. The dataframes are returned in the order in which each
+  combination first appears, and each of them keeps all the columns,
+  including the ones used for partitioning.
 
-  Every dataframe in the result keeps all the columns, including the ones it
-  was split on, and the list follows the order in which each combination
-  first appears. A lazy dataframe is collected first. Groups on the dataframe
-  are ignored.
+  Groups are ignored and the resulting dataframes are ungrouped.
 
   ## Examples
 
       iex> df = Explorer.DataFrame.new(id: ["b", "a", "b"], x: [1, 2, 3])
       iex> [b, a] = Explorer.DataFrame.partition_by(df, "id")
-      iex> Explorer.DataFrame.to_columns(b, atom_keys: true)
-      %{id: ["b", "b"], x: [1, 3]}
-      iex> Explorer.DataFrame.to_columns(a, atom_keys: true)
-      %{id: ["a"], x: [2]}
+      iex> b
+      #Explorer.DataFrame<
+        Polars[2 x 2]
+        id string ["b", "b"]
+        x s64 [1, 3]
+      >
+      iex> a
+      #Explorer.DataFrame<
+        Polars[1 x 2]
+        id string ["a"]
+        x s64 [2]
+      >
+
+  When partitioning by multiple columns, there is one dataframe for each
+  combination of their values:
+
+      iex> df = Explorer.DataFrame.new(id: ["b", "a", "b", "a"], k: [1, 1, 2, 1], x: [1, 2, 3, 4])
+      iex> dfs = Explorer.DataFrame.partition_by(df, ["id", "k"])
+      iex> Enum.map(dfs, &Explorer.DataFrame.to_columns(&1, atom_keys: true))
+      [
+        %{id: ["b"], k: [1], x: [1]},
+        %{id: ["a", "a"], k: [1, 1], x: [2, 4]},
+        %{id: ["b"], k: [2], x: [3]}
+      ]
 
   """
   @doc type: :single
-  @spec partition_by(df :: DataFrame.t(), columns :: column_names() | column_name()) ::
-          [DataFrame.t()]
+  @spec partition_by(df :: DataFrame.t(), column() | columns()) :: [DataFrame.t()]
+  def partition_by(df, column) when is_column(column) do
+    partition_by(df, [column])
+  end
+
   def partition_by(df, columns) do
-    columns = if is_column(columns), do: [columns], else: columns
-    columns = to_existing_columns(df, columns)
+    case to_existing_columns(df, columns) do
+      [] ->
+        raise ArgumentError,
+              "partition_by/2 expects at least one column, but #{inspect(columns)} selects none"
 
-    if columns == [] do
-      raise ArgumentError, "partition_by/2 needs at least one column to split on"
+      partition_columns ->
+        Shared.apply_dataframe(ungroup(df), :partition_by, [partition_columns])
     end
-
-    Shared.apply_dataframe(ungroup(df), :partition_by, [columns])
   end
 
   @doc """

lib/explorer/polars_backend/data_frame.ex

@@ -22,7 +22,7 @@ defmodule Explorer.PolarsBackend.DataFrame do
   @impl true
   def from_query(conn, query, params) do
     adbc_result =
-      Adbc.Connection.query_pointer(conn, query, params, fn pointer, _num_rows ->
+      Adbc.Connection.query_pointer(conn, query, params, fn %Adbc.StreamResult{pointer: pointer} ->
         Explorer.PolarsBackend.Native.df_from_arrow_stream_pointer(pointer)
       end)
 
@@ -757,9 +757,12 @@ defmodule Explorer.PolarsBackend.DataFrame do
 
   @impl true
   def partition_by(%DataFrame{} = df, columns) do
+    # Partitions keep all columns, so they share the names and dtypes of `df`.
     df
     |> Shared.apply_dataframe(:df_partition_by, [columns])
-    |> Enum.map(&Shared.create_dataframe!/1)
+    |> Enum.map(fn out_df ->
+      Map.put(df, :data, out_df)
+    end)
   end
 
   @impl true

lib/explorer/polars_backend/lazy_frame.ex

@@ -491,11 +491,6 @@ defmodule Explorer.PolarsBackend.LazyFrame do
     end
   end
 
-  @impl true
-  def partition_by(%DF{} = ldf, columns) do
-    ldf |> collect() |> Eager.partition_by(columns)
-  end
-
   @impl true
   def distinct(%DF{} = df, %DF{} = out_df, columns) do
     maybe_columns_to_keep =
@@ -686,6 +681,7 @@ defmodule Explorer.PolarsBackend.LazyFrame do
     mask: 2,
     n_rows: 1,
     estimated_size: 1,
+    partition_by: 2,
     pivot_wider: 5,
     pull: 2,
     put: 4,

lib/explorer/polars_backend/native.ex

@@ -158,8 +158,8 @@ defmodule Explorer.PolarsBackend.Native do
 
   def df_mask(_df, _mask), do: err()
   def df_n_rows(_df), do: err()
-  def df_partition_by(_df, _groups), do: err()
   def df_names(_df), do: err()
+  def df_partition_by(_df, _groups), do: err()
   def df_pivot_wider(_df, _id_columns, _pivot_column, _values_column, _names_prefix), do: err()
   def df_pull(_df, _name), do: err()
   def df_put_column(_df, _series), do: err()

mix.lock

@@ -12,9 +12,9 @@
   "decimal": {:hex, :decimal, "3.1.1", "430d87b04011ce6cbd4fd205be758311a81f87d552d40904abd00f015935b1d0", [:mix], [], "hexpm", "c5f25f2ced74a0587d03e6023f595db8e924c9d3922c8c8ffd9edfc4498cf1f6"},
   "deep_merge": {:hex, :deep_merge, "1.0.2", "476aa7ea61c54de96220051b998d893869069094da65b96101aebf79416f8a1e", [:mix], [], "hexpm", "737a53cdc9758fedbb608bdc213969e65729466c4ef3cd8e8726d0335dff116c"},
   "dll_loader_helper_beam": {:hex, :dll_loader_helper_beam, "1.2.2", "b86f97ec8fc64770c87468e41969eb309d87b29dd5a439b667e5954f85f8f65a", [:rebar3], [], "hexpm", "0e6119edde0ef5e42b4fe22d7dc71b7462e08573cee977c01a26ec5d9cd94a9a"},
-  "earmark_parser": {:hex, :earmark_parser, "1.4.45", "cba8369ab2a1342e419bc2760eec731b17be828941dcf494045d44766227e1d5", [:mix], [], "hexpm", "d3ec045bf122965db20c0bdb420e19ee1415843135327124918473feb4b328e8"},
+  "earmark_parser": {:hex, :earmark_parser, "1.4.46", "67607a0532e810c6f630a515c548d0b24949643f168cc556303bee4cf96105c7", [:mix], [], "hexpm", "9c44636e8a1c68c62f526b2dcd85d941dbbcee7ab82cf64ba06ce28bef8e89f5"},
   "elixir_make": {:hex, :elixir_make, "0.10.0", "16577e2583a79bb79237bbff349619ef5d80afffc07eac6e4faf0d00e2ddaf7d", [:mix], [], "hexpm", "dc1f09fb7fa68866b886abd5f0f3c83553b1a19a52359a899e92af1bb3b31982"},
-  "ex_doc": {:hex, :ex_doc, "0.40.3", "4a972ffe64bc07dc605af487e98fc19b72a4185f55ca031b94c0552d6071c1d9", [:mix], [{:earmark_parser, "~> 1.4.44", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_c, ">= 0.1.0", [hex: :makeup_c, repo: "hexpm", optional: true]}, {:makeup_elixir, "~> 0.14 or ~> 1.0", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1 or ~> 1.0", [hex: :makeup_erlang, repo: "hexpm", optional: false]}, {:makeup_html, ">= 0.1.0", [hex: :makeup_html, repo: "hexpm", optional: true]}], "hexpm", "2756e357742fecd9749b489b85d67c9ce99c465f2e75728d9e6dc8d704b973de"},
+  "ex_doc": {:hex, :ex_doc, "0.40.4", "66f2e42bf588594d5a8aab31cad87f2ddad09d0da1b1a2f379340ec2c2e497cb", [:mix], [{:earmark_parser, "~> 1.4.46", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_c, ">= 0.1.0", [hex: :makeup_c, repo: "hexpm", optional: true]}, {:makeup_elixir, "~> 0.14 or ~> 1.0", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1 or ~> 1.0", [hex: :makeup_erlang, repo: "hexpm", optional: false]}, {:makeup_html, ">= 0.1.0", [hex: :makeup_html, repo: "hexpm", optional: true]}], "hexpm", "6222b9e423d76584ee34df2c82a5ed72c2d53dc153f7f483ad28b378694186cc"},
   "finch": {:hex, :finch, "0.19.0", "c644641491ea854fc5c1bbaef36bfc764e3f08e7185e1f084e35e0672241b76d", [:mix], [{:mime, "~> 1.0 or ~> 2.0", [hex: :mime, repo: "hexpm", optional: false]}, {:mint, "~> 1.6.2 or ~> 1.7", [hex: :mint, repo: "hexpm", optional: false]}, {:nimble_options, "~> 0.4 or ~> 1.0", [hex: :nimble_options, repo: "hexpm", optional: false]}, {:nimble_pool, "~> 1.1", [hex: :nimble_pool, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "fc5324ce209125d1e2fa0fcd2634601c52a787aff1cd33ee833664a5af4ea2b6"},
   "fine": {:hex, :fine, "0.1.6", "4bf7151493443c454aac9f2fa2f34f5fefd0346a83fb5586a016c4a135c63247", [:mix], [], "hexpm", "5638eb4495488e885ebec167fa57973e5c35e1a50c344eb7666c90ec1c4e3b12"},
   "flame": {:hex, :flame, "0.5.2", "d46c4daa19b8921b71e0e57dc69edc01ce1311b1976c160192b05d4253b336e8", [:mix], [{:castore, ">= 0.0.0", [hex: :castore, repo: "hexpm", optional: true]}, {:jason, ">= 0.0.0", [hex: :jason, repo: "hexpm", optional: true]}], "hexpm", "82560ebef6ab3c277875493d0c93494740c930db0b1a3ff1a570eee9206cc6c0"},
@@ -36,12 +36,12 @@
   "ranch": {:hex, :ranch, "1.8.0", "8c7a100a139fd57f17327b6413e4167ac559fbc04ca7448e9be9057311597a1d", [:make, :rebar3], [], "hexpm", "49fbcfd3682fab1f5d109351b61257676da1a2fdbe295904176d5e521a2ddfe5"},
   "req": {:hex, :req, "0.5.8", "50d8d65279d6e343a5e46980ac2a70e97136182950833a1968b371e753f6a662", [:mix], [{:brotli, "~> 0.3.1", [hex: :brotli, repo: "hexpm", optional: true]}, {:ezstd, "~> 1.0", [hex: :ezstd, repo: "hexpm", optional: true]}, {:finch, "~> 0.17", [hex: :finch, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}, {:mime, "~> 2.0.6 or ~> 2.1", [hex: :mime, repo: "hexpm", optional: false]}, {:nimble_csv, "~> 1.0", [hex: :nimble_csv, repo: "hexpm", optional: true]}, {:plug, "~> 1.0", [hex: :plug, repo: "hexpm", optional: true]}], "hexpm", "d7fc5898a566477e174f26887821a3c5082b243885520ee4b45555f5d53f40ef"},
   "rustler": {:hex, :rustler, "0.38.0", "7a8906998ff0d28e3021c0a73264abcda719bda344b2e58307c6805b0f87c9b4", [:mix], [{:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}], "hexpm", "704c03c1bf66be12b031c5a389347b91c81c5cb819a24b068b0de36fe4a5652a"},
-  "rustler_precompiled": {:hex, :rustler_precompiled, "0.8.2", "5f25cbe220a8fac3e7ad62e6f950fcdca5a5a5f8501835d2823e8c74bf4268d5", [:mix], [{:castore, "~> 0.1 or ~> 1.0", [hex: :castore, repo: "hexpm", optional: false]}, {:rustler, "~> 0.23", [hex: :rustler, repo: "hexpm", optional: true]}], "hexpm", "63d1bd5f8e23096d1ff851839923162096364bac8656a4a3c00d1fff8e83ee0a"},
+  "rustler_precompiled": {:hex, :rustler_precompiled, "0.9.0", "3a052eda09f3d2436364645cc1f13279cf95db310eb0c17b0d8f25484b233aa0", [:mix], [{:rustler, "~> 0.23", [hex: :rustler, repo: "hexpm", optional: true]}], "hexpm", "471d97315bd3bf7b64623418b3693eedd8e47de3d1cb79a0ac8f9da7d770d94c"},
   "statistex": {:hex, :statistex, "1.1.1", "73612aa7f79e53c30569be065fd121e380f1cf57bc4c2da5b41be9246da18df9", [:mix], [], "hexpm", "310c4b49b34adf683de3103639006bed233ab54c08a4add65a531448e653857c"},
-  "stream_data": {:hex, :stream_data, "1.1.2", "05499eaec0443349ff877aaabc6e194e82bda6799b9ce6aaa1aadac15a9fdb4d", [:mix], [], "hexpm", "129558d2c77cbc1eb2f4747acbbea79e181a5da51108457000020a906813a1a9"},
+  "stream_data": {:hex, :stream_data, "1.4.0", "026f929db613aabea6208012ae9b8970d3fd5f88b3bdf26831bc536f98c42036", [:mix], [], "hexpm", "2b0ee3a340dcce1c8cf6302a763ee757d1e01c54d6e16d9069062509d68b1dc9"},
   "table": {:hex, :table, "0.1.2", "87ad1125f5b70c5dea0307aa633194083eb5182ec537efc94e96af08937e14a8", [:mix], [], "hexpm", "7e99bc7efef806315c7e65640724bf165c3061cdc5d854060f74468367065029"},
   "table_rex": {:hex, :table_rex, "4.1.0", "fbaa8b1ce154c9772012bf445bfb86b587430fb96f3b12022d3f35ee4a68c918", [:mix], [], "hexpm", "95932701df195d43bc2d1c6531178fc8338aa8f38c80f098504d529c43bc2601"},
   "telemetry": {:hex, :telemetry, "1.4.2", "a0cb522801dffb1c49fe6e30561badffc7b6d0e180db1300df759faa22062855", [:rebar3], [], "hexpm", "928f6495066506077862c0d1646609eed891a4326bee3126ba54b60af61febb1"},
   "toml": {:hex, :toml, "0.7.0", "fbcd773caa937d0c7a02c301a1feea25612720ac3fa1ccb8bfd9d30d822911de", [:mix], [], "hexpm", "0690246a2478c1defd100b0c9b89b4ea280a22be9a7b313a8a058a2408a2fa70"},
-  "tz": {:hex, :tz, "0.28.1", "717f5ffddfd1e475e2a233e221dc0b4b76c35c4b3650b060c8e3ba29dd6632e9", [:mix], [{:castore, "~> 0.1 or ~> 1.0", [hex: :castore, repo: "hexpm", optional: true]}, {:mint, "~> 1.6", [hex: :mint, repo: "hexpm", optional: true]}], "hexpm", "bfdca1aa1902643c6c43b77c1fb0cb3d744fd2f09a8a98405468afdee0848c8a"},
+  "tz": {:hex, :tz, "0.28.4", "f5433914dcf8619d593979ec65bbe878158cad818bc2b53b44d9e4bda509753f", [:mix], [{:castore, "~> 0.1 or ~> 1.0", [hex: :castore, repo: "hexpm", optional: true]}, {:mint, "~> 1.6", [hex: :mint, repo: "hexpm", optional: true]}], "hexpm", "72fe8d526cf30fcc4e07cec4ceb583ba5f2d69c6334c9533419c836a6acd1c56"},
 }

native/explorer/src/dataframe.rs

@@ -81,15 +81,6 @@ pub fn df_n_rows(df: ExDataFrame) -> Result<usize, ExplorerError> {
     Ok(df.height())
 }
 
-#[rustler::nif(schedule = "DirtyCpu")]
-pub fn df_partition_by(
-    df: ExDataFrame,
-    groups: Vec<String>,
-) -> Result<Vec<ExDataFrame>, ExplorerError> {
-    let dfs = df.partition_by_stable(groups, true)?;
-    Ok(dfs.into_iter().map(ExDataFrame::new).collect())
-}
-
 #[rustler::nif]
 pub fn df_width(df: ExDataFrame) -> Result<usize, ExplorerError> {
     Ok(df.width())
@@ -440,6 +431,17 @@ pub fn df_group_indices(
     Ok(series)
 }
 
+#[rustler::nif(schedule = "DirtyCpu")]
+pub fn df_partition_by(
+    df: ExDataFrame,
+    groups: Vec<&str>,
+) -> Result<Vec<ExDataFrame>, ExplorerError> {
+    let include_key = true;
+    let dfs = df.partition_by_stable(groups, include_key)?;
+
+    Ok(dfs.into_iter().map(ExDataFrame::new).collect())
+}
+
 #[rustler::nif(schedule = "DirtyCpu")]
 pub fn df_pivot_wider(
     df: ExDataFrame,

test/explorer/data_frame_test.exs

@@ -3029,29 +3029,89 @@ defmodule Explorer.DataFrameTest do
   end
 
   describe "partition_by/2" do
-    test "one dataframe per distinct combination, in first-seen order" do
+    test "returns one dataframe per value, in the order they first appear" do
       df = DF.new(id: ["b", "a", "b", "a"], k: [1, 1, 2, 1], x: [1.0, 2.0, 3.0, 4.0])
 
       assert [b, a] = DF.partition_by(df, "id")
       assert DF.to_columns(b, atom_keys: true) == %{id: ["b", "b"], k: [1, 2], x: [1.0, 3.0]}
       assert DF.to_columns(a, atom_keys: true) == %{id: ["a", "a"], k: [1, 1], x: [2.0, 4.0]}
+      assert DF.dtypes(b) == DF.dtypes(df)
+      assert DF.dtypes(a) == DF.dtypes(df)
+    end
+
+    test "returns one dataframe per combination of values" do
+      df = DF.new(id: ["b", "a", "b", "a"], k: [1, 1, 2, 1], x: [1.0, 2.0, 3.0, 4.0])
 
       assert [b1, a1, b2] = DF.partition_by(df, [:id, :k])
       assert DF.to_columns(b1, atom_keys: true) == %{id: ["b"], k: [1], x: [1.0]}
       assert DF.to_columns(a1, atom_keys: true) == %{id: ["a", "a"], k: [1, 1], x: [2.0, 4.0]}
       assert DF.to_columns(b2, atom_keys: true) == %{id: ["b"], k: [2], x: [3.0]}
+    end
+
+    test "accepts any column selection" do
+      df = DF.new(id: ["b", "a", "b"], x: [1, 2, 3])
+
+      expected =
+        df
+        |> DF.partition_by("id")
+        |> Enum.map(&DF.to_columns/1)
+
+      for columns <- [:id, 0, [0], 0..0, ~r/^id$/, &(&1 == "id")] do
+        regex_partitioned =
+          df
+          |> DF.partition_by(columns)
+          |> Enum.map(&DF.to_columns/1)
+
+        assert regex_partitioned == expected
+      end
+    end
+
+    test "keeps nil values in their own dataframe" do
+      df = DF.new(id: ["a", nil, "a", nil], x: [1, 2, 3, 4])
+
+      assert [a, nils] = DF.partition_by(df, "id")
+      assert DF.to_columns(a, atom_keys: true) == %{id: ["a", "a"], x: [1, 3]}
+      assert DF.to_columns(nils, atom_keys: true) == %{id: [nil, nil], x: [2, 4]}
+    end
+
+    test "returns an empty list for an empty dataframe" do
+      df = DF.new([id: [], x: []], dtypes: [id: :string, x: {:s, 64}])
 
-      assert_raise ArgumentError, ~r/at least one column/, fn -> DF.partition_by(df, []) end
+      assert DF.partition_by(df, "id") == []
     end
 
-    test "collects a lazy dataframe and ignores groups" do
+    test "ignores groups" do
       df = DF.new(id: ["b", "a", "b"], x: [1, 2, 3])
 
-      assert [b, a] = df |> DF.lazy() |> DF.partition_by("id")
+      assert [b, a] = df |> DF.group_by("x") |> DF.partition_by("id")
       assert DF.to_columns(b, atom_keys: true) == %{id: ["b", "b"], x: [1, 3]}
       assert DF.to_columns(a, atom_keys: true) == %{id: ["a"], x: [2]}
+      assert DF.groups(b) == []
+      assert DF.groups(a) == []
+    end
+
+    test "raises for empty or unknown column selections" do
+      df = DF.new(id: ["b", "a", "b"], x: [1, 2, 3])
+
+      assert_raise ArgumentError,
+                   "partition_by/2 expects at least one column, but [] selects none",
+                   fn -> DF.partition_by(df, []) end
 
-      assert [_, _] = df |> DF.group_by("x") |> DF.partition_by("id")
+      assert_raise ArgumentError,
+                   "partition_by/2 expects at least one column, but ~r/y/ selects none",
+                   fn -> DF.partition_by(df, ~r/y/) end
+
+      assert_raise ArgumentError, ~r/could not find column name "y"/, fn ->
+        DF.partition_by(df, "y")
+      end
+    end
+
+    test "is not available for lazy dataframes" do
+      ldf = DF.new(id: ["b", "a", "b"], x: [1, 2, 3]) |> DF.lazy()
+
+      assert_raise RuntimeError, ~r/`partition_by\/2` is not available/, fn ->
+        DF.partition_by(ldf, "id")
+      end
     end
   end
 

test/support/generator.ex

@@ -11,7 +11,7 @@ defmodule Explorer.Generator do
   """
 
   import Bitwise, only: [<<<: 2]
-  import StreamData
+  import StreamData, except: [date: 0, date: 1]
 
   @type gen(a) :: StreamData.t(a)