diff --git a/lib/explorer/shared.ex b/lib/explorer/shared.ex index caedd8b4b..a9b88dabc 100644 --- a/lib/explorer/shared.ex +++ b/lib/explorer/shared.ex @@ -207,16 +207,18 @@ defmodule Explorer.Shared do names end - def to_existing_columns(%{names: names} = df, %Range{} = columns, raise?) do + def to_existing_columns(%{names: names} = df, first..last//step = columns, raise?) do if raise? do n_cols = Explorer.DataFrame.n_columns(df) # With `Enum.slice/2`, negative indices are counted from the end. - [slice_min, slice_max] = - [columns.first, columns.last] + [slice_min, slice_pseudo_max] = + [first, last] |> Enum.map(&if(&1 < 0, do: n_cols + &1, else: &1)) |> Enum.sort() + slice_max = slice_min + step * (Range.size(slice_min..slice_pseudo_max//step) - 1) + if slice_min < 0 or slice_max >= n_cols do raise ArgumentError, "range #{inspect(columns)} is out of bounds for a dataframe with #{n_cols} column(s)" diff --git a/test/explorer/data_frame_test.exs b/test/explorer/data_frame_test.exs index a7ee10d69..5685bfabf 100644 --- a/test/explorer/data_frame_test.exs +++ b/test/explorer/data_frame_test.exs @@ -2690,6 +2690,7 @@ defmodule Explorer.DataFrameTest do assert DF.to_columns(df[[:a, :c]]) == %{"a" => [1, 2, 3], "c" => [4.0, 5.1, 6.2]} assert DF.to_columns(df[0..-2//1]) == %{"a" => [1, 2, 3], "b" => ["a", "b", "c"]} assert DF.to_columns(df[-3..-1]) == DF.to_columns(df) + assert DF.to_columns(df[1..3//3]) == %{"b" => ["a", "b", "c"]} assert DF.to_columns(df[..]) == DF.to_columns(df) assert %Series{} = s1 = df[0]