Skip to content

Commit

Permalink
Add getindex/setindex! methods for Colon
Browse files Browse the repository at this point in the history
  • Loading branch information
simonster committed Jun 24, 2015
1 parent fd88acf commit c793c5b
Showing 1 changed file with 44 additions and 0 deletions.
44 changes: 44 additions & 0 deletions src/dataframe/dataframe.jl
Original file line number Diff line number Diff line change
Expand Up @@ -240,6 +240,9 @@ function Base.getindex{T <: ColumnIndex}(df::DataFrame, col_inds::AbstractVector
return DataFrame(new_columns, Index(_names(df)[selected_columns]))
end

# df[:] => (Sub)?DataFrame
Base.getindex(df::DataFrame, col_inds::Colon) = copy(df)

# df[SingleRowIndex, SingleColumnIndex] => Scalar
function Base.getindex(df::DataFrame, row_ind::Real, col_ind::ColumnIndex)
selected_column = index(df)[col_ind]
Expand All @@ -266,6 +269,22 @@ function Base.getindex{R <: Real, T <: ColumnIndex}(df::DataFrame, row_inds::Abs
return DataFrame(new_columns, Index(_names(df)[selected_columns]))
end

# df[:, SingleColumnIndex] => (Sub)?AbstractVector
# df[:, MultiColumnIndex] => (Sub)?DataFrame
Base.getindex{T<:ColumnIndex}(df::DataFrame, row_inds::Colon, col_inds::Union(T, AbstractVector{T})) = df[col_inds]

# df[SingleRowIndex, :] => (Sub)?DataFrame
Base.getindex(df::DataFrame, row_ind::Real, col_inds::Colon) = df[[row_ind], col_inds]

# df[MultiRowIndex, :] => (Sub)?DataFrame
function Base.getindex{R<:Real}(df::DataFrame, row_inds::AbstractVector{R}, col_inds::Colon)
new_columns = Any[dv[row_inds] for dv in df.columns]
return DataFrame(new_columns, copy(index(df)))
end

# df[:, :] => (Sub)?DataFrame
Base.getindex(df::DataFrame, ::Colon, ::Colon) = copy(df)

##############################################################################
##
## setindex!()
Expand Down Expand Up @@ -403,6 +422,9 @@ function Base.setindex!{T <: ColumnIndex}(df::DataFrame,
return df
end

# df[:] = AbstractVector or Single Item
Base.setindex!(df::DataFrame, v, ::Colon) = (df[1:size(df, 2)] = v; df)

# df[SingleRowIndex, SingleColumnIndex] = Single Item
function Base.setindex!(df::DataFrame,
v::Any,
Expand Down Expand Up @@ -562,6 +584,28 @@ function Base.setindex!{R <: Real, T <: ColumnIndex}(df::DataFrame,
return df
end

# df[:] = DataFrame, df[:, :] = DataFrame
function Base.setindex!(df::DataFrame,
new_df::DataFrame,
row_inds::Colon,
col_inds::Colon=Colon())
df.columns = copy(new_df.columns)
df.colindex = copy(new_df.colindex)
df
end

# df[:, :] = ...
Base.setindex!(df::DataFrame, v, ::Colon, ::Colon) =
(df[1:size(df, 1), 1:size(df, 2)] = v; df)

# df[Any, :] = ...
Base.setindex!(df::DataFrame, v, row_inds, ::Colon) =
(df[row_inds, 1:size(df, 2)] = v; df)

# df[:, Any] = ...
Base.setindex!(df::DataFrame, v, ::Colon, col_inds) =
(df[col_inds] = v; df)

# Special deletion assignment
Base.setindex!(df::DataFrame, x::Nothing, col_ind::Int) = delete!(df, col_ind)

Expand Down

0 comments on commit c793c5b

Please sign in to comment.