SPARK-19342 bug fixed in collect method for collecting timestamp column, if column values are all NAs, the type is logical

titicaca · titicaca · commit 6a0eb3f8789b · 2017-01-25T11:09:04.000+08:00
diff --git a/R/pkg/R/DataFrame.R b/R/pkg/R/DataFrame.R
@@ -1136,17 +1136,11 @@ setMethod("collect",
 
                   # Note that "binary" columns behave like complex types.
                   if (!is.null(PRIMITIVE_TYPES[[colType]]) && colType != "binary") {
-                    valueIndex <- which(!is.na(col))
-                    if (length(valueIndex) > 0 && valueIndex[1] > 1) {
-                      colTail <- col[-(1 : (valueIndex[1] - 1))]
-                      vec <- do.call(c, colTail)
-                      classVal <- class(vec)
-                      vec <- c(rep(NA, valueIndex[1] - 1), vec)
-                      class(vec) <- classVal
-                    } else {
-                      vec <- do.call(c, col)
-                    }
+                    vec <- do.call(c, col)
                     stopifnot(class(vec) != "list")
+                    # If vec is an vector with only NAs, the type is logical
+                    if (length(vec[!is.na(vec)]) > 0)
+                      class(vec) <- PRIMITIVE_TYPES[[colType]]
                     df[[colIndex]] <- vec
                   } else {
                     df[[colIndex]] <- col
diff --git a/R/pkg/inst/tests/testthat/test_sparkSQL.R b/R/pkg/inst/tests/testthat/test_sparkSQL.R
@@ -2748,6 +2748,24 @@ test_that("Call DataFrameWriter.load() API in Java without path and check argume
                  "Unnamed arguments ignored: 2, 3, a.")
 })
 
+test_that("Collect on DataFrame when NAs exists at the top of a timestamp column") {
+  ldf <- data.frame(col1 = c(0, 1, 2),
+                   col2 = c(as.POSIXct("2017-01-01 00:00:01"),
+                            NA,
+                            as.POSIXct("2017-01-01 12:00:01")))
+  sdf1 <- createDataFrame(ldf)
+  ldf1 <- collect(sdf1)
+  expect_equal(dtypes(sdf1), list(c("col1", "double"), c("col2", "timestamp")))
+  expect_equal(class(ldf1$col1), "numeric")
+  expect_equal(class(ldf1$col2), "POSIXct")
+  
+  sdf2 <- filter(sdf1, "col1 > 0")
+  ldf2 <- collect(sdf2)
+  expect_equal(dtypes(sdf2), list(c("col1", "double"), c("col2", "timestamp")))
+  expect_equal(class(ldf2$col1), "numeric")
+  expect_equal(class(ldf2$col2), "POSIXct")
+}
+
 unlink(parquetPath)
 unlink(orcPath)
 unlink(jsonPath)
