Use Distributed.pmap in ParallelTableTransform and ColwiseFeatureTransform (#288)

* Use 'Distributed.pmap' in 'ParallelTableTransform'

* Use 'Distributed.pmap' in 'ColwiseFeatureTransform'

* Apply suggestions

* Update docs

Author: eliascarv

Project.toml

@@ -9,6 +9,7 @@ CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597"
 CoDa = "5900dafe-f573-5c72-b367-76665857777b"
 ColumnSelectors = "9cc86067-7e36-4c61-b350-1ac9833d277f"
 DataScienceTraits = "6cb2f572-2d2b-4ba6-bdb3-e710fa044d6c"
+Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -18,7 +19,6 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
-Transducers = "28d57a85-8fef-5791-bfe6-a80928e7c999"
 TransformsBase = "28dd2a49-a57a-4bfb-84ca-1a49db9b96b8"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
@@ -28,6 +28,7 @@ CategoricalArrays = "0.10"
 CoDa = "1.2"
 ColumnSelectors = "0.1"
 DataScienceTraits = "0.3"
+Distributed = "1.9"
 Distributions = "0.25"
 InverseFunctions = "0.1"
 LinearAlgebra = "1.9"
@@ -37,7 +38,6 @@ Random = "1.9"
 Statistics = "1.9"
 StatsBase = "0.33, 0.34"
 Tables = "1.6"
-Transducers = "0.4"
 TransformsBase = "1.5"
 Unitful = "1.17"
 julia = "1.9"

docs/src/index.md

@@ -25,9 +25,7 @@ missing transform, and contributions are very welcome.
   revertible when the individual transforms are revertible.
 
 - Branches of a pipeline and colwise transforms are run in parallel
-  using multiple threads with the awesome
-  [Transducers.jl](https://github.com/JuliaFolds/Transducers.jl)
-  framework.
+  using multiple processes with the Distributed standard library.
 
 - Pipelines can be reapplied to unseen "test" data using the same cache
   (e.g. constants) fitted with "training" data. For example, a `ZScore`

src/TableTransforms.jl

@@ -23,7 +23,7 @@ using Unitful: AbstractQuantity, AffineQuantity, AffineUnits, Units
 using Distributions: ContinuousUnivariateDistribution, Normal
 using InverseFunctions: NoInverse, inverse as invfun
 using StatsBase: AbstractWeights, Weights, sample
-using Transducers: tcollect
+using Distributed: CachingPool, pmap, workers
 using NelderMead: optimise
 
 import Distributions: quantile, cdf

src/transforms.jl

@@ -184,8 +184,9 @@ function applyfeat(transform::ColwiseFeatureTransform, feat, prep)
   names = Tables.columnnames(cols)
   snames = transform.selector(names)
 
-  # function to transform a single column
-  function colfunc(n)
+  # transform each column in parallel
+  pool = CachingPool(workers())
+  vals = pmap(pool, names) do n
     x = Tables.getcolumn(cols, n)
     if n ∈ snames
       c = colcache(transform, x)
@@ -197,9 +198,6 @@ function applyfeat(transform::ColwiseFeatureTransform, feat, prep)
     (n => y), c
   end
 
-  # parallel map with multiple threads
-  vals = tcollect(colfunc(n) for n in names)
-
   # new table with transformed columns
   𝒯 = (; first.(vals)...)
   newfeat = 𝒯 |> Tables.materializer(feat)
@@ -218,18 +216,14 @@ function revertfeat(transform::ColwiseFeatureTransform, newfeat, fcache)
 
   caches, snames = fcache
 
-  # function to transform a single column
-  function colfunc(i)
-    n = names[i]
-    c = caches[i]
+  # revert each column in parallel
+  pool = CachingPool(workers())
+  vals = pmap(pool, names, caches) do n, c
     y = Tables.getcolumn(cols, n)
     x = n ∈ snames ? colrevert(transform, y, c) : y
     n => x
   end
 
-  # parallel map with multiple threads
-  vals = tcollect(colfunc(i) for i in 1:length(names))
-
   # new table with transformed columns
   (; vals...) |> Tables.materializer(newfeat)
 end
@@ -244,18 +238,14 @@ function reapplyfeat(transform::ColwiseFeatureTransform, feat, fcache)
   # check that cache is valid
   _assert(length(names) == length(caches), "invalid caches for feat")
 
-  # function to transform a single column
-  function colfunc(i)
-    n = names[i]
-    c = caches[i]
+  # transform each column in parallel
+  pool = CachingPool(workers())
+  vals = pmap(pool, names, caches) do n, c
     x = Tables.getcolumn(cols, n)
     y = n ∈ snames ? colapply(transform, x, c) : x
     n => y
   end
 
-  # parallel map with multiple threads
-  vals = tcollect(colfunc(i) for i in 1:length(names))
-
   # new table with transformed columns
   (; vals...) |> Tables.materializer(feat)
 end

src/transforms/parallel.jl

@@ -41,8 +41,8 @@ isrevertible(p::ParallelTableTransform) = any(isrevertible, p.transforms)
 
 function apply(p::ParallelTableTransform, table)
   # apply transforms in parallel
-  f(transform) = apply(transform, table)
-  vals = tcollect(f(t) for t in p.transforms)
+  pool = CachingPool(workers())
+  vals = pmap(t -> apply(t, table), pool, p.transforms)
 
   # retrieve tables and caches
   tables = first.(vals)
@@ -122,9 +122,8 @@ function reapply(p::ParallelTableTransform, table, cache)
   caches = cache[1]
 
   # reapply transforms in parallel
-  f(t, c) = reapply(t, table, c)
-  itr = zip(p.transforms, caches)
-  tables = tcollect(f(t, c) for (t, c) in itr)
+  pool = CachingPool(workers())
+  tables = pmap((t, c) -> reapply(t, table, c), pool, p.transforms, caches)
 
   # features and metadata
   splits = divide.(tables)