timearray: add simple 2D getindex support (#423)

* timearray: add simple 2D getindex support

close #411

NEWS.md

@@ -28,14 +28,23 @@
 
 * Fix issues of `TimeArray` column names copying. (#418)
 
-* `ta[]` throws `BoundsError` now. (#420)
+* `timearray[]` throws `BoundsError` now. (#420)
 
   ```julia
   julia> cl[]
   ERROR: BoundsError: attempt to access TimeArray{Float64,1,Date,Array{Float64,1}}
     at index []
   ```
 
+* 2D `getindex` supports. (#423)
+
+  ```julia
+  ohlc[1:42,   [:High, :Low]]
+  ohlc[42:end, [:High, :Low]]
+  ohlc[:,      [:High, :Low]]
+  ohlc[42,     [:High, :Low]]
+  ```
+
 
 ### 0.15.0
 

docs/src/indexing.md

@@ -52,10 +52,11 @@ ohlc[Date(2000, 1, 3):Day(1):Date(2000, 2, 4)]
 
 ### Symbol
 
-| Example           | Description                            | Indexing value   |
-|-------------------|----------------------------------------|------------------|
-| `[:Open]`         | The column named `:Open`               | single symbol    |
-| `[:Open, :Close]` | The columns named `:Open` and `:Close` | multiple symbols |
+| Example             | Description                            | Indexing value   |
+|---------------------|----------------------------------------|------------------|
+| `[:Open]`           | The column named `:Open`               | single symbol    |
+| `[:Open, :Close]`   | The columns named `:Open` and `:Close` | multiple symbols |
+| `[[:Open, :Close]]` | The columns named `:Open` and `:Close` | multiple symbols |
 
 Examples in REPL:
 
@@ -67,13 +68,15 @@ using Dates
 ```@repl symbol-indexing
 ohlc[:Open]
 ohlc[:Open, :Close]
+cols = [:Open, :Close]
+ohlc[cols]
 ```
 
 ## Mixed approach
 
 | Example                     | Description                    | Indexing value                |
 |-----------------------------|--------------------------------|-------------------------------|
-| `[:Open][1:3]`              | `:Open` column & first 3 rows  | single symbol & integer range |
+| `[1:3, :Open]`              | `:Open` column & first 3 rows  | single symbol & integer range |
 | `[:Open][Date(2000, 1, 3)]` | `:Open` column and Jan 3, 2000 | single symbol & Date          |
 
 Examples in REPL:
@@ -84,6 +87,6 @@ using Dates
 ```
 
 ```@repl mixed-indexing
-ohlc[:Open][1:3]
+ohlc[1:3, :Open]
 ohlc[:Open][Date(2000, 1, 3)]
 ```

src/timearray.jl

@@ -349,10 +349,32 @@ function getindex(ta::TimeArray, s::Symbol)
 end
 
 # array of columns by name
-function getindex(ta::TimeArray, ss::Symbol...)
-    ns = [findcol(ta, s) for s in ss]
-    TimeArray(timestamp(ta), values(ta)[:, ns], collect(ss), meta(ta))
-end
+getindex(ta::TimeArray, ss::Symbol...) = getindex(ta, collect(ss))
+getindex(ta::TimeArray, ss::Vector{Symbol}) =
+    TimeArray(ta; values = values(ta)[:, map(s -> findcol(ta, s), ss)], colnames = ss)
+
+# ta[rows, cols]
+getindex(ta::TimeArray,
+         rows::Union{AbstractVector{<:Integer},Colon},
+         cols::AbstractVector{Symbol}) =
+    TimeArray(
+        ta;
+        timestamp = timestamp(ta)[rows],
+        values = values(ta)[rows, map(s -> findcol(ta, s), cols)],
+        colnames = cols,
+        unchecked = true)
+
+# ta[n, cols]
+getindex(ta::TimeArray, n::Integer, cols) =
+    getindex(ta, [n], cols)
+
+# ta[rows, col]
+getindex(ta::TimeArray, rows, col::Symbol) =
+    getindex(ta, rows, [col])
+
+# ta[n, col]
+getindex(ta::TimeArray, n::Integer, col::Symbol) =
+    getindex(ta, [n], [col])
 
 # single date
 function getindex(ta::TimeArray{T,N,D}, d::D) where {T,N,D}
@@ -376,7 +398,10 @@ getindex(ta::TimeArray, k::TimeArray{Bool,1}) = ta[findwhen(k)]
 # getindex{T,N}(ta::TimeArray{T,N}, d::DAYOFWEEK) = ta[dayofweek(timestamp(ta)) .== d]
 
 # Define end keyword
-lastindex(ta::TimeArray) = length(timestamp(ta))
+lastindex(ta::TimeArray, d::Integer = 1) =
+    (d == 1) ? length(timestamp(ta)) :
+    (d == 2) ? length(colnames(ta)) :
+    1
 
 eachindex(ta::TimeArray) = Base.OneTo(length(timestamp(ta)))
 

test/timearray.jl

@@ -329,6 +329,68 @@ end
         @test_throws MethodError ohlc[merge(op .> cl, op .< cl)]
     end
 
+    @testset "getindex on Vector{Symbol}" begin
+        hl = [:High, :Low]
+        ta = ohlc[hl]
+
+        @test size(values(ta)) == (length(timestamp(ohlc)), 2)
+        @test colnames(ta) == hl
+    end
+
+    @testset "2D getindex on [Vector{Int}, Vector{Symbol}]" begin
+        hl = [:High, :Low]
+        ta = ohlc[1:10, hl]
+
+        @test size(values(ta)) == (10, 2)
+        @test colnames(ta) == hl
+
+        ta = ohlc[10:end, hl]
+        @test size(values(ta)) == (length(timestamp(ohlc)) - 9, 2)
+        @test colnames(ta) == hl
+
+        ta = ohlc[:, hl]
+        @test size(values(ta)) == (length(timestamp(ohlc)), 2)
+        @test colnames(ta) == hl
+
+        # test KeyError
+        @test_throws KeyError ohlc[1:2, [:Unknown]]
+        @test_throws KeyError ohlc[1:end, [:Unknown]]
+        @test_throws KeyError ohlc[:, [:Unknown]]
+    end
+
+    @testset "2D getindex on [Integer, Vector{Symbol}]" begin
+        hl = [:High, :Low]
+        ta = ohlc[42, hl]
+        @test size(values(ta)) == (1, 2)
+        @test colnames(ta) == hl
+
+        ta = ohlc[end, hl]
+        @test size(values(ta)) == (1, 2)
+        @test colnames(ta) == hl
+
+        @test_throws KeyError ohlc[42, [:Unknown]]
+    end
+
+    @testset "2D getindex on [Vector{Int}, Symbol]" begin
+        ta = ohlc[1:42, :Open]
+        @test size(values(ta)) == (42, 1)
+        @test colnames(ta) == [:Open]
+
+        ta = ohlc[42:43, :Close]
+        @test size(values(ta)) == (2, 1)
+        @test colnames(ta) == [:Close]
+
+        @test_throws KeyError ohlc[1:42, :Unknown]
+    end
+
+    @testset "2D getindex on [Integer, Symbol]" begin
+        ta = ohlc[42, :Open]
+        @test size(values(ta)) == (1, 1)
+        @test colnames(ta) == [:Open]
+
+        @test_throws KeyError ohlc[42, :Unknown]
+    end
+
     @testset "Base.eachindex" begin
         @test eachindex(cl)   == 1:length(cl)
         @test eachindex(ohlc) == 1:length(ohlc)