Allow accessing the dataframe inside query

elixir-explorer · Sep 4, 2024 · 49bb939 · 49bb939
1 parent 77e0201
commit 49bb939
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/lib/explorer/query.ex b/lib/explorer/query.ex
@@ -739,6 +739,11 @@ defmodule Explorer.Query do
     end
   end
 
+  @doc """
+  Returns the dataframe scoped by this query.
+  """
+  defmacro df(), do: df_var()
+
   @doc false
   def __across__(df, selector) do
     df

diff --git a/test/explorer/query_test.exs b/test/explorer/query_test.exs
@@ -5,6 +5,12 @@ defmodule Explorer.QueryTest do
   require Explorer.DataFrame, as: DF
   doctest Explorer.Query
 
+  test "allows accessing the dataframe" do
+    assert DF.new(a: [1, 2, 3])
+           |> DF.filter(df()["a"] < ^if(true, do: 3, else: 1))
+           |> DF.to_columns(atom_keys: true) == %{a: [1, 2]}
+  end
+
   test "allows Kernel operations within pin" do
     assert DF.new(a: [1, 2, 3])
            |> DF.filter(a < ^if(true, do: 3, else: 1))