Support array in from_json in R

HyukjinKwon · HyukjinKwon · commit 10feb9d2e9e4 · 2017-03-06T20:55:50.000+09:00
diff --git a/R/pkg/R/functions.R b/R/pkg/R/functions.R
@@ -2430,33 +2430,43 @@ setMethod("date_format", signature(y = "Column", x = "character"),
             column(jc)
           })
 
+setClassUnion("characterOrstructType", c("character", "structType"))
+
 #' from_json
 #'
 #' Parses a column containing a JSON string into a Column of \code{structType} with the specified
 #' \code{schema}. If the string is unparseable, the Column will contains the value NA.
 #'
 #' @param x Column containing the JSON string.
-#' @param schema a structType object to use as the schema to use when parsing the JSON string.
+#' @param schema a structType object or the data type string representing an array or struct type
+#'               used in structField to use as the schema to use when parsing the JSON string.
 #' @param ... additional named properties to control how the json is parsed, accepts the same
 #'            options as the JSON data source.
 #'
 #' @family normal_funcs
 #' @rdname from_json
 #' @name from_json
-#' @aliases from_json,Column,structType-method
+#' @aliases from_json,Column,characterOrstructType-method
 #' @export
 #' @examples
 #' \dontrun{
 #' schema <- structType(structField("name", "string"),
 #' select(df, from_json(df$value, schema, dateFormat = "dd/MM/yyyy"))
 #'}
 #' @note from_json since 2.2.0
-setMethod("from_json", signature(x = "Column", schema = "structType"),
+setMethod("from_json", signature(x = "Column", schema = "characterOrstructType"),
           function(x, schema, ...) {
+            if (is.character(schema)) {
+              jschema <- callJStatic("org.apache.spark.sql.api.r.SQLUtils",
+                                     "getSQLDataType",
+                                     schema)
+            } else {
+              jschema <- schema$jobj
+            }
             options <- varargsToStrEnv(...)
             jc <- callJStatic("org.apache.spark.sql.functions",
                               "from_json",
-                              x@jc, schema$jobj, options)
+                              x@jc, jschema, options)
             column(jc)
           })
 
diff --git a/R/pkg/inst/tests/testthat/test_sparkSQL.R b/R/pkg/inst/tests/testthat/test_sparkSQL.R
@@ -1342,28 +1342,52 @@ test_that("column functions", {
   df <- read.json(mapTypeJsonPath)
   j <- collect(select(df, alias(to_json(df$info), "json")))
   expect_equal(j[order(j$json), ][1], "{\"age\":16,\"height\":176.5}")
-  df <- as.DataFrame(j)
-  schema <- structType(structField("age", "integer"),
-                       structField("height", "double"))
-  s <- collect(select(df, alias(from_json(df$json, schema), "structcol")))
-  expect_equal(ncol(s), 1)
-  expect_equal(nrow(s), 3)
-  expect_is(s[[1]][[1]], "struct")
-  expect_true(any(apply(s, 1, function(x) { x[[1]]$age == 16 } )))
-
-  # passing option
-  df <- as.DataFrame(list(list("col" = "{\"date\":\"21/10/2014\"}")))
-  schema2 <- structType(structField("date", "date"))
-  expect_error(tryCatch(collect(select(df, from_json(df$col, schema2))),
-                        error = function(e) { stop(e) }),
-               paste0(".*(java.lang.NumberFormatException: For input string:).*"))
-  s <- collect(select(df, from_json(df$col, schema2, dateFormat = "dd/MM/yyyy")))
-  expect_is(s[[1]][[1]]$date, "Date")
-  expect_equal(as.character(s[[1]][[1]]$date), "2014-10-21")
-
-  # check for unparseable
-  df <- as.DataFrame(list(list("a" = "")))
-  expect_equal(collect(select(df, from_json(df$a, schema)))[[1]][[1]], NA)
+
+  schemas <- list(structType(structField("age", "integer"), structField("height", "double")),
+                  "struct<age:integer,height:double>")
+  for (schema in schemas) {
+    df <- as.DataFrame(j)
+    s <- collect(select(df, alias(from_json(df$json, schema), "structcol")))
+    expect_equal(ncol(s), 1)
+    expect_equal(nrow(s), 3)
+    expect_is(s[[1]][[1]], "struct")
+    expect_true(any(apply(s, 1, function(x) { x[[1]]$age == 16 } )))
+
+    # passing option
+    df <- as.DataFrame(list(list("col" = "{\"date\":\"21/10/2014\"}")))
+    schema2 <- structType(structField("date", "date"))
+    expect_error(tryCatch(collect(select(df, from_json(df$col, schema2))),
+    error = function(e) { stop(e) }),
+    paste0(".*(java.lang.NumberFormatException: For input string:).*"))
+    s <- collect(select(df, from_json(df$col, schema2, dateFormat = "dd/MM/yyyy")))
+    expect_is(s[[1]][[1]]$date, "Date")
+    expect_equal(as.character(s[[1]][[1]]$date), "2014-10-21")
+
+    # check for unparseable
+    df <- as.DataFrame(list(list("a" = "")))
+    expect_equal(collect(select(df, from_json(df$a, schema)))[[1]][[1]], NA)
+  }
+
+  # check if array type in string is correctly supported.
+  jsonArr <- "[{\"name\":\"Bob\"}, {\"name\":\"Alice\"}]"
+  df <- as.DataFrame(list(list("people" = jsonArr)))
+  arr <- collect(select(df, alias(from_json(df$people, "array<struct<name:string>>"), "arrcol")))
+  expect_equal(ncol(arr), 1)
+  expect_equal(nrow(arr), 1)
+  expect_is(arr[[1]][[1]], "list")
+  expect_equal(length(arr$arrcol[[1]]), 2)
+  expect_equal(arr$arrcol[[1]][[1]]$name, "Bob")
+  expect_equal(arr$arrcol[[1]][[2]]$name, "Alice")
+
+  # check for unparseable data type
+  expect_error(tryCatch(collect(select(df, from_json(df$people, "unknown"))),
+  error = function(e) { stop(e) }),
+  paste0(".*(Invalid type unknown).*"))
+
+  # check for incorrect data type
+  expect_error(tryCatch(collect(select(df, from_json(df$people, "integer"))),
+  error = function(e) { stop(e) }),
+  paste0(".*(data type mismatch: Input schema int must be a struct or an array of structs).*"))
 })
 
 test_that("column binary mathfunctions", {
