Skip to content

Commit 45c040e

Browse files
committed
df.col .= x is in-place and df.col = scalar is allowed
Fixes #3200
1 parent fdfa2f7 commit 45c040e

File tree

8 files changed

+242
-162
lines changed

8 files changed

+242
-162
lines changed

src/dataframe/dataframe.jl

Lines changed: 32 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -632,11 +632,11 @@ Base.getindex(df::DataFrame, row_ind::typeof(!), col_inds::MultiColumnIndex) =
632632
##############################################################################
633633

634634
# Will automatically add a new column if needed
635-
function insert_single_column!(df::DataFrame, v::AbstractVector, col_ind::ColumnIndex)
636-
if ncol(df) != 0 && nrow(df) != length(v)
635+
function insert_single_column!(df::DataFrame, v::Any, col_ind::ColumnIndex; copycols = false)
636+
dv = _preprocess_column(v, nrow(df), copycols)
637+
if ncol(df) != 0 && nrow(df) != length(dv)
637638
throw(ArgumentError("New columns must have the same length as old columns"))
638639
end
639-
dv = isa(v, AbstractRange) ? collect(v) : v
640640
firstindex(dv) != 1 && _onebased_check_error()
641641

642642
if haskey(index(df), col_ind)
@@ -664,24 +664,22 @@ function insert_single_entry!(df::DataFrame, v::Any, row_ind::Integer, col_ind::
664664
end
665665
end
666666

667-
# df[!, SingleColumnIndex] = AbstractVector
668-
function Base.setindex!(df::DataFrame, v::AbstractVector, ::typeof(!), col_ind::ColumnIndex)
667+
# df[!, SingleColumnIndex] = value
668+
function Base.setindex!(df::DataFrame, v::Any, ::typeof(!), col_ind::ColumnIndex)
669669
insert_single_column!(df, v, col_ind)
670670
return df
671671
end
672672

673-
# df.col = AbstractVector
673+
# df.col = value
674674
# separate methods are needed due to dispatch ambiguity
675675
Base.setproperty!(df::DataFrame, col_ind::Symbol, v::AbstractVector) =
676676
(df[!, col_ind] = v)
677677
Base.setproperty!(df::DataFrame, col_ind::AbstractString, v::AbstractVector) =
678678
(df[!, col_ind] = v)
679-
Base.setproperty!(::DataFrame, col_ind::Symbol, v::Any) =
680-
throw(ArgumentError("It is only allowed to pass a vector as a column of a DataFrame. " *
681-
"Instead use `df[!, col_ind] .= v` if you want to use broadcasting."))
682-
Base.setproperty!(::DataFrame, col_ind::AbstractString, v::Any) =
683-
throw(ArgumentError("It is only allowed to pass a vector as a column of a DataFrame. " *
684-
"Instead use `df[!, col_ind] .= v` if you want to use broadcasting."))
679+
Base.setproperty!(df::DataFrame, col_ind::Symbol, v::Any) =
680+
(df[!, col_ind] = v)
681+
Base.setproperty!(df::DataFrame, col_ind::AbstractString, v::Any) =
682+
(df[!, col_ind] = v)
685683

686684
# df[SingleRowIndex, SingleColumnIndex] = Single Item
687685
function Base.setindex!(df::DataFrame, v::Any, row_ind::Integer, col_ind::ColumnIndex)
@@ -786,6 +784,28 @@ for T1 in (:AbstractVector, :Not, :Colon, :(typeof(!))),
786784
end
787785
end
788786

787+
for T1 in (:(typeof(!)),),
788+
T2 in MULTICOLUMNINDEX_TUPLE
789+
@eval function Base.setindex!(df::DataFrame,
790+
v::AbstractVector,
791+
row_inds::$T1,
792+
col_inds::$T2)
793+
throw(ArgumentError("a vector can not be assigned to multiple rows and columns, consider reshaping to a matrix first"))
794+
end
795+
796+
@eval function Base.setindex!(df::DataFrame,
797+
v::Any,
798+
row_inds::$T1,
799+
col_inds::$T2)
800+
idxs = index(df)[col_inds]
801+
for col in idxs
802+
# this will drop metadata appropriately
803+
df[row_inds, col] = v
804+
end
805+
return df
806+
end
807+
end
808+
789809
"""
790810
copy(df::DataFrame; copycols::Bool=true)
791811

src/other/broadcasting.jl

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,7 @@ function Base.maybeview(df::AbstractDataFrame, rows, cols)
153153
return view(df, rows, cols)
154154
end
155155

156+
# df[:, cols] .= ...
156157
function Base.dotview(df::AbstractDataFrame, ::Colon, cols::ColumnIndex)
157158
if haskey(index(df), cols)
158159
_drop_all_nonnote_metadata!(parent(df))
@@ -168,10 +169,15 @@ function Base.dotview(df::AbstractDataFrame, ::Colon, cols::ColumnIndex)
168169
return LazyNewColDataFrame(df, Symbol(cols))
169170
end
170171

172+
# df[!, cols] .= ...
171173
function Base.dotview(df::AbstractDataFrame, ::typeof(!), cols)
172174
if !(cols isa ColumnIndex)
173175
return ColReplaceDataFrame(df, convert(Vector{Int}, index(df)[cols]))
174176
end
177+
if haskey(index(df), cols)
178+
_drop_all_nonnote_metadata!(parent(df))
179+
return view(df, :, cols)
180+
end
175181
if cols isa SymbolOrString
176182
if columnindex(df, cols) == 0 && !is_column_insertion_allowed(df)
177183
throw(ArgumentError("creating new columns in a SubDataFrame that subsets " *
@@ -184,7 +190,13 @@ function Base.dotview(df::AbstractDataFrame, ::typeof(!), cols)
184190
end
185191

186192
if isdefined(Base, :dotgetproperty) # Introduced in Julia 1.7
193+
# df.col .= ...
187194
function Base.dotgetproperty(df::AbstractDataFrame, col::SymbolOrString)
195+
if haskey(index(df), col)
196+
_drop_all_nonnote_metadata!(parent(df))
197+
return df[!, col]
198+
end
199+
188200
if columnindex(df, col) == 0 && !is_column_insertion_allowed(df)
189201
throw(ArgumentError("creating new columns in a SubDataFrame that subsets " *
190202
"columns of its parent data frame is disallowed"))

src/subdataframe/subdataframe.jl

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -209,7 +209,7 @@ end
209209
# and then define methods for them)
210210
# consider merging SubDataFrame and DataFrame setindex! methods
211211

212-
function Base.setindex!(sdf::SubDataFrame, v::AbstractVector,
212+
function Base.setindex!(sdf::SubDataFrame, v::Any,
213213
::typeof(!), col_ind::ColumnIndex)
214214
if col_ind isa Union{Signed, Unsigned} && !(1 <= col_ind <= ncol(sdf))
215215
throw(ArgumentError("Cannot assign to non-existent column: $col_ind"))
@@ -219,15 +219,17 @@ function Base.setindex!(sdf::SubDataFrame, v::AbstractVector,
219219
throw(ArgumentError("creating new columns in a SubDataFrame that subsets " *
220220
"columns of its parent data frame is disallowed"))
221221
end
222+
v = _preprocess_column(v, nrow(sdf), false)
222223
sdf[:, col_ind] = v
223224
else
224225
pdf = parent(sdf)
225226
p_col_ind = parentcols(index(sdf), col_ind)
226227
old_col = pdf[!, p_col_ind]
227228
T = eltype(old_col)
228-
S = eltype(v)
229+
S = v isa AbstractVector ? eltype(v) : typeof(v)
229230
newcol = similar(old_col, promote_type(T, S), length(old_col))
230231
newcol .= old_col
232+
v = _preprocess_column(v, nrow(sdf), false)
231233
newcol[rows(sdf)] = v
232234
pdf[!, p_col_ind] = newcol
233235
end
@@ -278,12 +280,10 @@ Base.setproperty!(df::SubDataFrame, col_ind::Symbol, v::AbstractVector) =
278280
(df[!, col_ind] = v)
279281
Base.setproperty!(df::SubDataFrame, col_ind::AbstractString, v::AbstractVector) =
280282
(df[!, col_ind] = v)
281-
Base.setproperty!(::SubDataFrame, col_ind::Symbol, v::Any) =
282-
throw(ArgumentError("It is only allowed to pass a vector as a column of a SubDataFrame. " *
283-
"Instead use `df[!, col_ind] .= v` if you want to use broadcasting."))
284-
Base.setproperty!(::SubDataFrame, col_ind::AbstractString, v::Any) =
285-
throw(ArgumentError("It is only allowed to pass a vector as a column of a SubDataFrame. " *
286-
"Instead use `df[!, col_ind] .= v` if you want to use broadcasting."))
283+
Base.setproperty!(df::SubDataFrame, col_ind::Symbol, v::Any) =
284+
(df[!, col_ind] = v)
285+
Base.setproperty!(df::SubDataFrame, col_ind::AbstractString, v::Any) =
286+
(df[!, col_ind] = v)
287287

288288
##############################################################################
289289
##

0 commit comments

Comments
 (0)