DTableColumn improvements (#12)

Author: krynju

.github/workflows/CI.yml

@@ -21,6 +21,7 @@ jobs:
       matrix:
         version:
           - '1.7'
+          - '1.8'
           - 'nightly'
         os:
           - ubuntu-latest

src/table/dataframes_interface.jl

@@ -1,7 +1,8 @@
+import DataAPI: All, Between, BroadcastedSelector, Cols
+import DataFrames: AsTable, ByRow, ColumnIndex, MultiColumnIndex, normalize_selection
 import InvertedIndices: BroadcastedInvertedIndex
-import DataAPI: Between, All, Cols, BroadcastedSelector
-import DataFrames: ColumnIndex, MultiColumnIndex,
-    ByRow, AsTable, normalize_selection
+
+
 
 make_pair_concrete(@nospecialize(x::Pair)) =
     make_pair_concrete(x.first) => make_pair_concrete(x.second)

src/table/dtable.jl

@@ -21,11 +21,12 @@ mutable struct DTable
     chunks::VTYPE
     tabletype
     schema::Union{Nothing,Tables.Schema}
-    DTable(chunks::VTYPE, tabletype) = new(chunks, tabletype, nothing)
 end
 
-DTable(chunks::Vector{Dagger.EagerThunk}, args...) = DTable(VTYPE(chunks), args...)
-DTable(chunks::Vector{Dagger.Chunk}, args...) = DTable(VTYPE(chunks), args...)
+DTable(chunks::Vector, tabletype) = DTable(VTYPE(chunks), tabletype, nothing)
+DTable(chunks::Vector, tabletype, schema) = DTable(VTYPE(chunks), tabletype, schema)
+
+
 
 """
     DTable(table; tabletype=nothing) -> DTable

src/table/dtable_column.jl

@@ -1,101 +1,87 @@
-
 mutable struct DTableColumn{T,TT}
     dtable::DTable
-    current_chunk::Int
     col::Int
     colname::Symbol
     chunk_lengths::Vector{Int}
-    current_iterator::Union{Nothing,TT}
-    chunkstore::Union{Nothing,Vector{T}}
+    _chunk::Int
+    _iter::Union{Nothing,TT}
+    _chunkstore::Union{Nothing,Vector{T}}
 end
 
-__ff = (ch, col) -> Tables.getcolumn(Tables.columns(ch), col)
+function getcolumn_chunk(chunk_contents, col::Int)
+    return Tables.getcolumn(Tables.columns(chunk_contents), col)
+end
 
-function DTableColumn(dtable::DTable, col::Int)
-    column_eltype = Tables.schema(Tables.columns(dtable)).types[col]
-    iterator_type = fetch(Dagger.spawn((ch, _col) -> typeof(iterate(__ff(ch, _col))), dtable.chunks[1], col))
+function DTableColumn(d::DTable, col::Int)
+    column_eltype = Tables.schema(Tables.columns(d)).types[col]
+
+    iterator_type = Nothing
+    c_idx = 1
+    while iterator_type === Nothing && c_idx <= nchunks(d)
+        iterator_type = fetch(Dagger.spawn(
+            (ch, _col) -> typeof(iterate(getcolumn_chunk(ch, _col))),
+            d.chunks[c_idx],
+            col
+        ))
+        c_idx += 1
+    end
 
     DTableColumn{column_eltype,iterator_type}(
-        dtable,
-        0,
+        d,
         col,
-        _columnnames_svector(dtable)[col],
-        chunk_lengths(dtable),
+        _columnnames_svector(d)[col],
+        chunk_lengths(d),
+        0,
         nothing,
         nothing,
     )
 end
 
 
-function getindex(dtablecolumn::DTableColumn, idx::Int)
-    chunk_idx = 0
-    s = 1
-    for (i, e) in enumerate(dtablecolumn.chunk_lengths)
-        if s <= idx < s + e
-            chunk_idx = i
-            break
-        end
-        s = s + e
-    end
-    chunk_idx == 0 && throw(BoundsError())
-    offset = idx - s + 1
-    chunk = fetch(Dagger.spawn(__ff, dtablecolumn.dtable.chunks[chunk_idx], dtablecolumn.col))
-
-    row, iter = iterate(Tables.rows(chunk))
-    for _ in 1:(offset-1)
-        row, iter = iterate(Tables.rows(chunk), iter)
-    end
-    Tables.getcolumn(row, dtablecolumn.col)
-end
+DTableColumn(d::DTable, col::String) =
+    DTableColumn(d, only(indexin([col], string.(_columnnames_svector(d)))))
+DTableColumn(d::DTable, col::Symbol) = DTableColumn(d, string(col))
 
-length(dtablecolumn::DTableColumn) = sum(dtablecolumn.chunk_lengths)
+length(dtc::DTableColumn) = sum(dtc.chunk_lengths)
 
 
-function pull_next_chunk(dtablecolumn::DTableColumn, chunkidx::Int)
-    while dtablecolumn.current_iterator === nothing
-        chunkidx += 1
-        if chunkidx <= length(dtablecolumn.dtable.chunks)
-            dtablecolumn.chunkstore =
-                fetch(Dagger.spawn(__ff, dtablecolumn.dtable.chunks[chunkidx], dtablecolumn.col))
+function pull_next_chunk!(dtc::DTableColumn)
+    # find first non-empty chunk
+    while dtc._iter === nothing
+        dtc._chunk += 1
+        if dtc._chunk <= nchunks(dtc.dtable)
+            dtc._chunkstore = fetch(Dagger.spawn(
+                getcolumn_chunk,
+                dtc.dtable.chunks[dtc._chunk],
+                dtc.col
+            ))
         else
-            return chunkidx
+            return nothing
         end
-        dtablecolumn.current_iterator = iterate(dtablecolumn.chunkstore)
+        # iterate in case this chunk is empty
+        dtc._iter = iterate(dtc._chunkstore)
     end
-    return chunkidx
+    return nothing
 end
 
 
-function iterate(dtablecolumn::DTableColumn)
-    if length(dtablecolumn) == 0
-        return nothing
-    end
-    dtablecolumn.chunkstore = nothing
-    dtablecolumn.current_iterator = nothing
-    chunkidx = pull_next_chunk(dtablecolumn, 0)
-    ci = dtablecolumn.current_iterator
-    if ci === nothing
-        return nothing
-    else
-        return (ci[1], (chunkidx, ci[2]))
-    end
+function iterate(dtc::DTableColumn)
+    length(dtc) == 0 && return nothing
+
+    # on every iteration start reset the cache
+    dtc._chunkstore = nothing
+    dtc._iter = nothing
+    dtc._chunk = 0
+
+    # pull the first chunk
+    pull_next_chunk!(dtc)
+
+    return dtc._iter
 end
 
-function iterate(dtablecolumn::DTableColumn, iter)
-    (chunkidx, i) = iter
-    cs = dtablecolumn.chunkstore
-    ci = nothing
-    if cs !== nothing
-        ci = iterate(cs, i)
-    else
-        return nothing
-    end
-    dtablecolumn.current_iterator = ci
-    chunkidx = pull_next_chunk(dtablecolumn, chunkidx)
-    ci = dtablecolumn.current_iterator
-    if ci === nothing
-        return nothing
-    else
-        return (ci[1], (chunkidx, ci[2]))
-    end
+function iterate(dtc::DTableColumn, iter)
+    dtc._chunkstore === nothing && return nothing
+    dtc._iter = iterate(dtc._chunkstore, iter)
+    pull_next_chunk!(dtc)
+    return dtc._iter
 end

src/table/operations.jl

@@ -216,7 +216,7 @@ function filter(f, d::DTable)
         m = TableOperations.filter(_f, _chunk)
         Tables.materializer(_chunk)(m)
     end
-    DTable(map(c -> Dagger.spawn(chunk_wrap, c, f), d.chunks), d.tabletype)
+    DTable(map(c -> Dagger.spawn(chunk_wrap, c, f), d.chunks), d.tabletype, d.schema)
 end
 
 

test/column.jl

@@ -0,0 +1,26 @@
+@testset "DTableColumn" begin
+    S = 10_000
+    col_a = collect(1:S)
+    col_b = rand(S)
+    nt = (a=col_a, b=col_b)
+    d = DTable(nt, S ÷ 10)
+
+    @test collect(DTableColumn(d, 1)) == col_a
+    @test collect(DTableColumn(d, "a")) == col_a
+    @test collect(DTableColumn(d, :a)) == col_a
+    @test collect(DTableColumn(d, 2)) == col_b
+    @test collect(DTableColumn(d, "b")) == col_b
+    @test collect(DTableColumn(d, :b)) == col_b
+
+    d2 = filter(x -> x.a <= S / 2, d)
+    @test collect(DTableColumn(d2, 1)) == col_a[1:Int(S / 2)]
+    @test collect(DTableColumn(d2, 2)) == col_b[1:Int(S / 2)]
+
+    d2 = filter(x -> x.a >= S / 2, d)
+    @test collect(DTableColumn(d2, 1)) == col_a[Int(S / 2):end]
+    @test collect(DTableColumn(d2, 2)) == col_b[Int(S / 2):end]
+
+    d2 = filter(x -> x.a < 0, d)
+    @test collect(DTableColumn(d2, 1)) == col_a[1:-1]
+    @test collect(DTableColumn(d2, 2)) == col_b[1:-1]
+end

test/runtests.jl

@@ -14,4 +14,5 @@ using Distributed
 @testset "DTables.jl" begin
     include("table.jl")
     include("table_dataframes.jl")
+    include("column.jl")
 end