Add recursive cleanClosure for function access.

pkg/R/utils.R

@@ -350,7 +350,11 @@ processClosure <- function(node, oldEnv, argNames, newEnv) {
         # Set parameter 'inherits' to FALSE since we do not need to search in
         # attached package environments.
         if (exists(nodeChar, envir=func.env, inherits = FALSE)) {
-          assign(nodeChar, get(nodeChar, envir=func.env), envir = newEnv)
+          obj <- get(nodeChar, envir=func.env)
+          if (is.function(obj)) {
+            obj <- cleanClosure(obj)
+          }
+          assign(nodeChar, obj, envir = newEnv)
           break
         } else {
           # Continue to search in enclosure.
@@ -366,9 +370,11 @@ processClosure <- function(node, oldEnv, argNames, newEnv) {
 # outside a UDF, and stores them in a new environment.
 # param
 #   func A function whose closure needs to be captured.
-#   newEnv A new function environment to store necessary function dependencies.
-cleanClosure <- function(func, newEnv) {
-  if (is.function(func) && is.environment(newEnv)) {
+# return value
+#   a new function that has an correct environment (closure).
+cleanClosure <- function(func) {
+  if (is.function(func) {
+    newEnv <- new.env(parent = .GlobalEnv)
     # .defVars is a character vector of variables names defined in the function.
     assign(".defVars", c(), envir = .sparkREnv)
     func.body <- body(func)
@@ -377,5 +383,7 @@ cleanClosure <- function(func, newEnv) {
     argsNames <- argNames[-length(argNames)]  # Remove the ending NULL in pairlist.
     # Recursively examine variables in the function body.
     processClosure(func.body, oldEnv, argNames, newEnv)
+    environment(func) <- newEnv
   }
+  func
 }

pkg/inst/tests/test_utils.R

@@ -36,12 +36,12 @@ test_that("reserialize on RDD", {
   unlink(fileName)
 })
 
-test_that("clean.closure on R functions", {
+test_that("cleanClosure on R functions", {
   y <- c(1, 2, 3)
   g <- function(x) { x + 1 }
   f <- function(x) { g(x) + y }
-  env <- new.env()
-  cleanClosure(f, env)
+  newF <- cleanClosure(f)
+  env <- environment(newF)
   expect_equal(length(ls(env)), 2)  # y, g
   actual <- get("y", envir = env)
   expect_equal(actual, y)
@@ -53,20 +53,32 @@ test_that("clean.closure on R functions", {
   funcEnv <- new.env(parent = env2)
   f <- function(x) { min(g(x) + y) }
   environment(f) <- funcEnv  # enclosing relationship: f -> funcEnv -> env2 -> .GlobalEnv
-  env <- new.env()
-  SparkR:::cleanClosure(f, env)
+  newF <- SparkR:::cleanClosure(f)
+  env <- environment(newF)
   expect_equal(length(ls(env)), 2)  # "min" should not be included
   actual <- get("y", envir = env)
   expect_equal(actual, y)
   actual <- get("g", envir = env)
   expect_equal(actual, g)
+
+  g <- function(x) { x + y }
+  f <- function(x) { lapply(x, g) + 1 }
+  newF <- SparkR:::cleanClosure(f)
+  env <- environment(newF)
+  expect_equal(length(ls(env)), 1)  # Only "g", "y" should be in the environemnt of g.
+  expect_equal(ls(env), "g")
+  newG <- get("g", envir = env)
+  env <- environment(newG)
+  expect_equal(length(ls(env)), 1)
+  actual <- get("y", envir = env)
+  expect_equal(actual, y)
 
   # Test for function (and variable) definitions.
   f <- function(x) {
     g <- function(y) { y * 2 }
     g(x)
   }
-  env <- new.env()
-  SparkR:::cleanClosure(f, env)
+  newF <- SparkR:::cleanClosure(f)
+  env <- environment(newF)
   expect_equal(length(ls(env)), 0)  # "y" and "g" should not be included.
 })