Merge pull request #224 from JuliaML/fix

Fix #223 + Some refactors in `tablerows` and `Parallel`

Author: juliohm

src/tablerows.jl

@@ -52,7 +52,7 @@ struct CTableRows{T}
 
   function CTableRows(table)
     cols = Tables.columns(table)
-    nrows = _nrows(cols)
+    nrows = _nrows(table)
     new{typeof(cols)}(cols, nrows)
   end
 end
@@ -81,15 +81,17 @@ Tables.getcolumn(row::CTableRow, nm::Symbol) = Tables.getcolumn(getcols(row), nm
 
 struct RTableRows{T}
   rows::T
+  nrows::Int
 
   function RTableRows(table)
     rows = Tables.rows(table)
-    new{typeof(rows)}(rows)
+    nrows = _nrows(table)
+    new{typeof(rows)}(rows, nrows)
   end
 end
 
 # iterator interface
-Base.length(rows::RTableRows) = length(rows.rows)
+Base.length(rows::RTableRows) = rows.nrows
 function Base.iterate(rows::RTableRows, args...)
   next = iterate(rows.rows, args...)
   if isnothing(next)
@@ -116,9 +118,15 @@ Tables.getcolumn(row::RTableRow, nm::Symbol) = Tables.getcolumn(getrow(row), nm)
 # UTILS
 #-------
 
-function _nrows(cols)
-  names = Tables.columnnames(cols)
-  isempty(names) && return 0
-  column = Tables.getcolumn(cols, first(names))
-  length(column)
+function _nrows(table)
+  if Tables.rowaccess(table)
+    rows = Tables.rows(table)
+    length(rows)
+  else
+    cols = Tables.columns(table)
+    names = Tables.columnnames(cols)
+    isempty(names) && return 0
+    column = Tables.getcolumn(cols, first(names))
+    length(column)
+  end
 end

src/transforms/parallel.jl

@@ -53,6 +53,9 @@ function apply(p::ParallelTableTransform, table)
   feats = first.(splits)
   metas = last.(splits)
 
+  # check the number of rows of generated tables
+  @assert allequal(_nrows(f) for f in feats) "parallel branches must produce the same number of rows"
+
   # table with concatenated features
   newfeat = tablehcat(feats)
 
@@ -128,6 +131,9 @@ function reapply(p::ParallelTableTransform, table, cache)
   feats = first.(splits)
   metas = last.(splits)
 
+  # check the number of rows of generated tables
+  @assert allequal(_nrows(f) for f in feats) "parallel branches must produce the same number of rows"
+
   # table with concatenated features
   newfeat = tablehcat(feats)
 
@@ -140,24 +146,23 @@ end
 
 function tablehcat(tables)
   # concatenate columns
-  allvars, allvals = [], []
-  varsdict = Set{Symbol}()
-  for 𝒯 in tables
-    cols = Tables.columns(𝒯)
-    vars = Tables.columnnames(cols)
-    vals = [Tables.getcolumn(cols, var) for var in vars]
-    for (var, val) in zip(vars, vals)
-      while var ∈ varsdict
-        var = Symbol(var, :_)
+  allnames = Symbol[]
+  allcolumns = []
+  for table in tables
+    cols = Tables.columns(table)
+    names = Tables.columnnames(cols)
+    for name in names
+      column = Tables.getcolumn(cols, name)
+      while name ∈ allnames
+        name = Symbol(name, :_)
       end
-      push!(varsdict, var)
-      push!(allvars, var)
-      push!(allvals, val)
+      push!(allnames, name)
+      push!(allcolumns, column)
     end
   end
 
   # table with concatenated columns
-  𝒯 = (; zip(allvars, allvals)...)
+  𝒯 = (; zip(allnames, allcolumns)...)
   𝒯 |> Tables.materializer(first(tables))
 end
 

test/tablerows.jl

@@ -33,6 +33,7 @@
   @test row[:a] == 1
   @test row["a"] == 1
   # iterator interface
+  @test length(row) == 2
   item, state = iterate(row)
   @test item == 1
   item, state = iterate(row, state)
@@ -73,6 +74,7 @@
   @test row[:b] == 4
   @test row["b"] == 4
   # iterator interface
+  @test length(row) == 2
   item, state = iterate(row)
   @test item == 1
   item, state = iterate(row, state)

test/transforms/parallel.jl

@@ -46,4 +46,15 @@
   @test t == n
   tₒ = revert(T, n, c)
   @test tₒ == t
+
+  # https://github.com/JuliaML/TableTransforms.jl/issues/223
+  t = (a=1:4, b=5:8)
+  left = Select(:a) → Filter(row -> row.a < 4)
+  right = Select(:b) → Filter(row -> row.b < 7)
+  T = left ⊔ right
+  @test_throws AssertionError apply(T, t)
+  t1 = (a=1:4, b=4:7)
+  t2 = (a=1:4, b=5:8)
+  n, c = apply(T, t1)
+  @test_throws AssertionError reapply(T, t2, c)
 end