R/pkg/tests/fulltests/test_sparkSQL_arrow.R - master - githubmirror/spark

R/pkg/tests/fulltests/test_sparkSQL_arrow.R
267 строк9 KB
Dongjoon Hyun
[SPARK-44127][R][TESTS] Reenable test_sparkSQL_arrow.R in Java 21
24 авг 2023, 10:22
24 авг 2023, 10:2217f27ab
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

library(testthat)

context("SparkSQL Arrow optimization")

sparkSession <- sparkR.session(
  master = sparkRTestMaster,
  enableHiveSupport = FALSE,
  sparkConfig = list(spark.sql.execution.arrow.sparkr.enabled = "true"))

test_that("createDataFrame/collect Arrow optimization", {
  skip_if_not_installed("arrow")

  conf <- callJMethod(sparkSession, "conf")
  arrowEnabled <- sparkR.conf("spark.sql.execution.arrow.sparkr.enabled")[[1]]

  callJMethod(conf, "set", "spark.sql.execution.arrow.sparkr.enabled", "false")
  tryCatch({
    expected <- collect(createDataFrame(mtcars))
  },
  finally = {
    callJMethod(conf, "set", "spark.sql.execution.arrow.sparkr.enabled", arrowEnabled)
  })

  expect_equal(collect(createDataFrame(mtcars)), expected)
})

test_that("createDataFrame/collect Arrow optimization - many partitions (partition order test)", {
  skip_if_not_installed("arrow")
  expect_equal(collect(createDataFrame(mtcars, numPartitions = 32)),
               collect(createDataFrame(mtcars, numPartitions = 1)))
})

test_that("createDataFrame/collect Arrow optimization - type specification", {
  skip_if_not_installed("arrow")
  rdf <- data.frame(list(list(a = 1,
                              b = "a",
                              c = TRUE,
                              d = 1.1,
                              e = 1L,
                              f = as.Date("1990-02-24"),
                              g = as.POSIXct("1990-02-24 12:34:56"))))

  arrowEnabled <- sparkR.conf("spark.sql.execution.arrow.sparkr.enabled")[[1]]
  conf <- callJMethod(sparkSession, "conf")

  callJMethod(conf, "set", "spark.sql.execution.arrow.sparkr.enabled", "false")
  tryCatch({
    expected <- collect(createDataFrame(rdf))
  },
  finally = {
    callJMethod(conf, "set", "spark.sql.execution.arrow.sparkr.enabled", arrowEnabled)
  })

  expect_true(all(collect(createDataFrame(rdf)) == expected))
})

test_that("dapply() Arrow optimization", {
  skip_if_not_installed("arrow")
  df <- createDataFrame(mtcars)

  conf <- callJMethod(sparkSession, "conf")
  arrowEnabled <- sparkR.conf("spark.sql.execution.arrow.sparkr.enabled")[[1]]

  callJMethod(conf, "set", "spark.sql.execution.arrow.sparkr.enabled", "false")
  tryCatch({
    ret <- dapply(df,
                  function(rdf) {
                    stopifnot(is.data.frame(rdf))
                    rdf
                  },
                  schema(df))
    expected <- collect(ret)
  },
  finally = {
    callJMethod(conf, "set", "spark.sql.execution.arrow.sparkr.enabled", arrowEnabled)
  })

  ret <- dapply(df,
                function(rdf) {
                  stopifnot(is.data.frame(rdf))
                  # mtcars' hp is more then 50.
                  stopifnot(all(rdf$hp > 50))
                  rdf
                },
                schema(df))
  actual <- collect(ret)
  expect_equal(actual, expected)
  expect_equal(count(ret), nrow(mtcars))
})

test_that("dapply() Arrow optimization - type specification", {
  skip_if_not_installed("arrow")
  # Note that regular dapply() seems not supporting date and timestamps
  # whereas Arrow-optimized dapply() does.
  rdf <- data.frame(list(list(a = 1,
                              b = "a",
                              c = TRUE,
                              d = 1.1,
                              e = 1L)))
  # numPartitions are set to 8 intentionally to test empty partitions as well.
  df <- createDataFrame(rdf, numPartitions = 8)

  conf <- callJMethod(sparkSession, "conf")
  arrowEnabled <- sparkR.conf("spark.sql.execution.arrow.sparkr.enabled")[[1]]

  callJMethod(conf, "set", "spark.sql.execution.arrow.sparkr.enabled", "false")
  tryCatch({
    ret <- dapply(df, function(rdf) { rdf }, schema(df))
    expected <- collect(ret)
  },
  finally = {
    callJMethod(conf, "set", "spark.sql.execution.arrow.sparkr.enabled", arrowEnabled)
  })

  ret <- dapply(df, function(rdf) { rdf }, schema(df))
  actual <- collect(ret)
  expect_equal(actual, expected)
})

test_that("dapply() Arrow optimization - type specification (date and timestamp)", {
  skip_if_not_installed("arrow")
  rdf <- data.frame(list(list(a = as.Date("1990-02-24"),
                              b = as.POSIXct("1990-02-24 12:34:56"))))
  df <- createDataFrame(rdf)
  ret <- dapply(df, function(rdf) { rdf }, schema(df))
  expect_true(all(collect(ret) == rdf))
})

test_that("gapply() Arrow optimization", {
  skip_if_not_installed("arrow")
  df <- createDataFrame(mtcars)

  conf <- callJMethod(sparkSession, "conf")
  arrowEnabled <- sparkR.conf("spark.sql.execution.arrow.sparkr.enabled")[[1]]

  callJMethod(conf, "set", "spark.sql.execution.arrow.sparkr.enabled", "false")
  tryCatch({
    ret <- gapply(df,
                 "gear",
                 function(key, grouped) {
                   if (length(key) > 0) {
                     stopifnot(is.numeric(key[[1]]))
                   }
                   stopifnot(is.data.frame(grouped))
                   grouped
                 },
                 schema(df))
    expected <- collect(ret)
  },
  finally = {
    callJMethod(conf, "set", "spark.sql.execution.arrow.sparkr.enabled", arrowEnabled)
  })

  ret <- gapply(df,
               "gear",
               function(key, grouped) {
                 if (length(key) > 0) {
                   stopifnot(is.numeric(key[[1]]))
                 }
                 stopifnot(is.data.frame(grouped))
                 stopifnot(length(colnames(grouped)) == 11)
                 # mtcars' hp is more then 50.
                 stopifnot(all(grouped$hp > 50))
                 grouped
               },
               schema(df))
  actual <- collect(ret)
  expect_equal(actual, expected)
  expect_equal(count(ret), nrow(mtcars))
})

test_that("gapply() Arrow optimization - type specification", {
  skip_if_not_installed("arrow")
  # Note that regular gapply() seems not supporting date and timestamps
  # whereas Arrow-optimized gapply() does.
  rdf <- data.frame(list(list(a = 1,
                              b = "a",
                              c = TRUE,
                              d = 1.1,
                              e = 1L)))
  df <- createDataFrame(rdf)

  conf <- callJMethod(sparkSession, "conf")
  arrowEnabled <- sparkR.conf("spark.sql.execution.arrow.sparkr.enabled")[[1]]

  callJMethod(conf, "set", "spark.sql.execution.arrow.sparkr.enabled", "false")
  tryCatch({
    ret <- gapply(df,
                  "a",
                  function(key, grouped) { grouped }, schema(df))
    expected <- collect(ret)
  },
  finally = {
    callJMethod(conf, "set", "spark.sql.execution.arrow.sparkr.enabled", arrowEnabled)
  })

  ret <- gapply(df,
                "a",
                function(key, grouped) { grouped }, schema(df))
  actual <- collect(ret)
  expect_equal(actual, expected)
})

test_that("gapply() Arrow optimization - type specification (date and timestamp)", {
  skip_if_not_installed("arrow")
  rdf <- data.frame(list(list(a = as.Date("1990-02-24"),
                              b = as.POSIXct("1990-02-24 12:34:56"))))
  df <- createDataFrame(rdf)
  ret <- gapply(df,
                "a",
                function(key, grouped) { grouped }, schema(df))
  expect_true(all(collect(ret) == rdf))
})

test_that("Arrow optimization - unsupported types", {
  skip_if_not_installed("arrow")

  expect_error(checkSchemaInArrow(structType("a FLOAT")), "not support float type")
  expect_error(checkSchemaInArrow(structType("a BINARY")), "not support binary type")
  expect_error(checkSchemaInArrow(structType("a ARRAY<INT>")), "not support array type")
  expect_error(checkSchemaInArrow(structType("a MAP<INT, INT>")), "not support map type")
  expect_error(checkSchemaInArrow(structType("a STRUCT<a: INT>")),
               "not support nested struct type")
})

test_that("SPARK-32478: gapply() Arrow optimization - error message for schema mismatch", {
  skip_if_not_installed("arrow")
  df <- createDataFrame(list(list(a = 1L, b = "a")))

  expect_error(
    count(gapply(df, "a", function(key, group) { group }, structType("a int, b int"))),
    "expected IntegerType, IntegerType, got IntegerType, StringType")
})

test_that("SPARK-43789: Automatically pick the number of partitions based on Arrow batch size", {
  skip_if_not_installed("arrow")

  conf <- callJMethod(sparkSession, "conf")
  maxRecordsPerBatch <- sparkR.conf("spark.sql.execution.arrow.maxRecordsPerBatch")[[1]]

  callJMethod(conf, "set", "spark.sql.execution.arrow.maxRecordsPerBatch", "10")
  tryCatch({
    expect_equal(getNumPartitionsRDD(toRDD(createDataFrame(mtcars))), 4)
  },
  finally = {
    callJMethod(conf, "set", "spark.sql.execution.arrow.maxRecordsPerBatch", maxRecordsPerBatch)
  })
})

sparkR.session.stop()