Skip to content

Commit c793c5b

Browse files
committed
Add getindex/setindex! methods for Colon
Needed since JuliaLang/julia#10331
1 parent fd88acf commit c793c5b

File tree

1 file changed

+44
-0
lines changed

1 file changed

+44
-0
lines changed

src/dataframe/dataframe.jl

+44
Original file line numberDiff line numberDiff line change
@@ -240,6 +240,9 @@ function Base.getindex{T <: ColumnIndex}(df::DataFrame, col_inds::AbstractVector
240240
return DataFrame(new_columns, Index(_names(df)[selected_columns]))
241241
end
242242

243+
# df[:] => (Sub)?DataFrame
244+
Base.getindex(df::DataFrame, col_inds::Colon) = copy(df)
245+
243246
# df[SingleRowIndex, SingleColumnIndex] => Scalar
244247
function Base.getindex(df::DataFrame, row_ind::Real, col_ind::ColumnIndex)
245248
selected_column = index(df)[col_ind]
@@ -266,6 +269,22 @@ function Base.getindex{R <: Real, T <: ColumnIndex}(df::DataFrame, row_inds::Abs
266269
return DataFrame(new_columns, Index(_names(df)[selected_columns]))
267270
end
268271

272+
# df[:, SingleColumnIndex] => (Sub)?AbstractVector
273+
# df[:, MultiColumnIndex] => (Sub)?DataFrame
274+
Base.getindex{T<:ColumnIndex}(df::DataFrame, row_inds::Colon, col_inds::Union(T, AbstractVector{T})) = df[col_inds]
275+
276+
# df[SingleRowIndex, :] => (Sub)?DataFrame
277+
Base.getindex(df::DataFrame, row_ind::Real, col_inds::Colon) = df[[row_ind], col_inds]
278+
279+
# df[MultiRowIndex, :] => (Sub)?DataFrame
280+
function Base.getindex{R<:Real}(df::DataFrame, row_inds::AbstractVector{R}, col_inds::Colon)
281+
new_columns = Any[dv[row_inds] for dv in df.columns]
282+
return DataFrame(new_columns, copy(index(df)))
283+
end
284+
285+
# df[:, :] => (Sub)?DataFrame
286+
Base.getindex(df::DataFrame, ::Colon, ::Colon) = copy(df)
287+
269288
##############################################################################
270289
##
271290
## setindex!()
@@ -403,6 +422,9 @@ function Base.setindex!{T <: ColumnIndex}(df::DataFrame,
403422
return df
404423
end
405424

425+
# df[:] = AbstractVector or Single Item
426+
Base.setindex!(df::DataFrame, v, ::Colon) = (df[1:size(df, 2)] = v; df)
427+
406428
# df[SingleRowIndex, SingleColumnIndex] = Single Item
407429
function Base.setindex!(df::DataFrame,
408430
v::Any,
@@ -562,6 +584,28 @@ function Base.setindex!{R <: Real, T <: ColumnIndex}(df::DataFrame,
562584
return df
563585
end
564586

587+
# df[:] = DataFrame, df[:, :] = DataFrame
588+
function Base.setindex!(df::DataFrame,
589+
new_df::DataFrame,
590+
row_inds::Colon,
591+
col_inds::Colon=Colon())
592+
df.columns = copy(new_df.columns)
593+
df.colindex = copy(new_df.colindex)
594+
df
595+
end
596+
597+
# df[:, :] = ...
598+
Base.setindex!(df::DataFrame, v, ::Colon, ::Colon) =
599+
(df[1:size(df, 1), 1:size(df, 2)] = v; df)
600+
601+
# df[Any, :] = ...
602+
Base.setindex!(df::DataFrame, v, row_inds, ::Colon) =
603+
(df[row_inds, 1:size(df, 2)] = v; df)
604+
605+
# df[:, Any] = ...
606+
Base.setindex!(df::DataFrame, v, ::Colon, col_inds) =
607+
(df[col_inds] = v; df)
608+
565609
# Special deletion assignment
566610
Base.setindex!(df::DataFrame, x::Nothing, col_ind::Int) = delete!(df, col_ind)
567611

0 commit comments

Comments
 (0)