Skip to content

Commit

Permalink
fix first and last for negative row count
Browse files Browse the repository at this point in the history
  • Loading branch information
bkamins committed Nov 9, 2023
1 parent 8c32d53 commit 94c235c
Show file tree
Hide file tree
Showing 3 changed files with 27 additions and 8 deletions.
3 changes: 3 additions & 0 deletions NEWS.md
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,9 @@

## Bug fixes

* Correctly throw an error if negative number of rows is passed
to `first` or `last`
([#3402](https://github.com/JuliaData/DataFrames.jl/pull/3402))
* Always use the default thread pool for multithreaded operations,
instead of using the interactive thread pool when Julia was started
with `-tM,N` with N > 0
Expand Down
18 changes: 14 additions & 4 deletions src/abstractdataframe/abstractdataframe.jl
Original file line number Diff line number Diff line change
Expand Up @@ -558,14 +558,19 @@ Base.first(df::AbstractDataFrame) = df[1, :]
first(df::AbstractDataFrame, n::Integer; view::Bool=false)
Get a data frame with the `n` first rows of `df`.
If `n` is greater than number of rows of `df` get all rows.
Error if `n` is negative.
If `view=false` a freshly allocated `DataFrame` is returned.
If `view=true` then a `SubDataFrame` view into `df` is returned.
$METADATA_FIXED
"""
@inline Base.first(df::AbstractDataFrame, n::Integer; view::Bool=false) =
view ? Base.view(df, 1:min(n ,nrow(df)), :) : df[1:min(n, nrow(df)), :]
@inline function Base.first(df::AbstractDataFrame, n::Integer; view::Bool=false)
n < 0 && throw(ArgumentError("Number of elements must be nonnegative"))
r = min(n, nrow(df))
return view ? Base.view(df, 1:r, :) : df[1:r, :]
end

"""
last(df::AbstractDataFrame)
Expand All @@ -580,14 +585,19 @@ Base.last(df::AbstractDataFrame) = df[nrow(df), :]
last(df::AbstractDataFrame, n::Integer; view::Bool=false)
Get a data frame with the `n` last rows of `df`.
If `n` is greater than number of rows of `df` get all rows.
Error if `n` is negative.
If `view=false` a freshly allocated `DataFrame` is returned.
If `view=true` then a `SubDataFrame` view into `df` is returned.
$METADATA_FIXED
"""
@inline Base.last(df::AbstractDataFrame, n::Integer; view::Bool=false) =
view ? Base.view(df, max(1, nrow(df)-n+1):nrow(df), :) : df[max(1, nrow(df)-n+1):nrow(df), :]
@inline function Base.last(df::AbstractDataFrame, n::Integer; view::Bool=false)
n < 0 && throw(ArgumentError("Number of elements must be nonnegative"))
r = max(1, nrow(df) - n + 1)
return view ? Base.view(df, r:nrow(df), :) : df[r:nrow(df), :]
end

"""
describe(df::AbstractDataFrame; cols=:)
Expand Down
14 changes: 10 additions & 4 deletions test/dataframe.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1180,10 +1180,16 @@ end
@test_throws BoundsError first(DataFrame(x=[]))
@test_throws BoundsError last(DataFrame(x=[]))

@test first(df, 6) == DataFrame(A=1:6)
@test first(df, 1) == DataFrame(A=1)
@test last(df, 6) == DataFrame(A=5:10)
@test last(df, 1) == DataFrame(A=10)
for v in (true, false)
@test first(df, 6, view=v) == DataFrame(A=1:6)
@test first(df, 1, view=v) == DataFrame(A=1)
@test first(df, 0, view=v) == DataFrame()
@test_throws ArgumentError first(df, -1, view=v)
@test last(df, 6, view=v) == DataFrame(A=5:10)
@test last(df, 1, view=v) == DataFrame(A=10)
@test last(df, 0, view=v) == DataFrame()
@test_throws ArgumentError last(df, -1, view=v)
end

@inferred first(df, 6)
@inferred last(df, 6)
Expand Down

0 comments on commit 94c235c

Please sign in to comment.