Add load_ipc_stream to the LazyFrame impl

elixir-explorer · Dec 17, 2022 · b326afd · b326afd
1 parent 9425b0b
commit b326afd
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git a/lib/explorer/polars_backend/lazy_frame.ex b/lib/explorer/polars_backend/lazy_frame.ex
@@ -192,6 +192,14 @@ defmodule Explorer.PolarsBackend.LazyFrame do
     end
   end
 
+  @impl true
+  def load_ipc_stream(contents, columns) do
+    case Eager.load_ipc_stream(contents, columns) do
+      {:ok, df} -> {:ok, Eager.to_lazy(df)}
+      {:error, error} -> {:error, error}
+    end
+  end
+
   # Groups
 
   # TODO: Make the functions of non-implemented functions

diff --git a/test/explorer/data_frame/lazy_test.exs b/test/explorer/data_frame/lazy_test.exs
@@ -263,4 +263,18 @@ defmodule Explorer.DataFrame.LazyTest do
 
     assert DF.to_columns(df1) == DF.to_columns(df)
   end
+
+  test "load_ipc_stream/2 - with defaults", %{df: df} do
+    df = DF.slice(df, 0, 10)
+    contents = DF.dump_ipc_stream!(df)
+
+    ldf = DF.load_ipc_stream!(contents, lazy: true)
+
+    # no-op 
+    assert DF.to_lazy(ldf) == ldf
+
+    df1 = DF.collect(ldf)
+
+    assert DF.to_columns(df1) == DF.to_columns(df)
+  end
 end