diff --git a/ext/BangBangDataFramesExt.jl b/ext/BangBangDataFramesExt.jl index 4b24e7a..ce834ec 100644 --- a/ext/BangBangDataFramesExt.jl +++ b/ext/BangBangDataFramesExt.jl @@ -14,6 +14,22 @@ end BangBang.append!!(df::DataFrames.DataFrame, source) = df_append!!(df, source) BangBang.copyappendable(df::DataFrames.AbstractDataFrame) = copy(df) + function BangBang.setindex!!( + df::DataFrames.DataFrame, + v, + row::Integer, + col::Union{Integer,Symbol,AbstractString}, + ) + col_idx = col isa Integer ? Int(col) : DataFrames.columnindex(df, col) + columns = getfield(df, :columns) + old_col = columns[col_idx] + new_col = BangBang.setindex!!(old_col, v, row) + if new_col !== old_col + columns[col_idx] = new_col + end + return df + end + _getvalue(x, pos, name) = getproperty(x, name) _getvalue(x::AbstractVector, pos, name) = x[pos] _getvalue(x::Tuple, pos, name) = x[pos] diff --git a/test/test_dataframes.jl b/test/test_dataframes.jl index 5f18aed..372365b 100644 --- a/test/test_dataframes.jl +++ b/test/test_dataframes.jl @@ -1,6 +1,6 @@ module TestDataFrames -using BangBang: append!!, push!! +using BangBang: append!!, push!!, setindex!! using CategoricalArrays: CategoricalArray using DataFrames: DataFrame using InitialValues: InitialValue @@ -84,4 +84,32 @@ end end end +@testset "setindex!!" begin + @testset "basic" begin + df = DataFrame("a" => [0.0]) + @test setindex!!(df, 1.0, 1, "a") === df + @test df[1, "a"] === 1.0 + end + + @testset "Symbol column" begin + df = DataFrame(a = [0.0]) + @test setindex!!(df, 2.0, 1, :a) === df + @test df[1, :a] === 2.0 + end + + @testset "Integer column" begin + df = DataFrame(a = [0.0]) + @test setindex!!(df, 3.0, 1, 1) === df + @test df[1, 1] === 3.0 + end + + @testset "type widening" begin + df = DataFrame(a = [1, 2, 3]) + result = setindex!!(df, 1.5, 2, :a) + @test result === df + @test df[2, :a] === 1.5 + @test eltype(df.a) === Float64 + end +end + end # module