fix: dim() should return integer (#577)

pola-rs · Dec 7, 2023 · 56d8e97 · 56d8e97
1 parent 6445153
commit 56d8e97
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 7 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -24,6 +24,8 @@
 - `as_polars_df()` for `data.frame` has a new argument `make_names_unique` (#561).
 - New methods `$str$to_date()`, `$str$to_time()`, `$str$to_datetime()` as 
   alternatives to `$str$strptime()` (#558).
+- The `dim()` function for DataFrame and LazyFrame correctly returns integer instead of
+  double (#577).
 
 # polars 0.11.0
 

diff --git a/R/s3_methods.R b/R/s3_methods.R
@@ -172,7 +172,7 @@ tail.RPolarsLazyFrame = tail.RPolarsDataFrame
 #'
 #' @export
 #' @rdname S3_dim
-dim.RPolarsDataFrame = function(x) x$shape
+dim.RPolarsDataFrame = function(x) as.integer(x$shape)
 
 #' @export
 #' @rdname S3_dim

diff --git a/tests/testthat/test-csv-read.R b/tests/testthat/test-csv-read.R
@@ -29,7 +29,7 @@ test_that("works with single URL", {
   })
   # put messages back in the console
   sink(type = "message")
-  expect_identical(dim(out), c(31, 6))
+  expect_identical(dim(out), c(31L, 6L))
 })
 
 test_that("args separator and eol work", {
@@ -49,11 +49,11 @@ test_that("args skip_rows and skip_rows_after_header work", {
   write.csv(dat, tmpf, row.names = FALSE)
 
   out = pl$read_csv(tmpf, skip_rows = 25)
-  expect_identical(nrow(out), 125)
+  expect_identical(nrow(out), 125L)
   expect_named(out, c("4.8", "3.4", "1.9", "0.2", "setosa"))
 
   out = pl$read_csv(tmpf, skip_rows_after_header = 25)
-  expect_identical(nrow(out), 125)
+  expect_identical(nrow(out), 125L)
   expect_named(out, names(iris))
 })
 
@@ -88,7 +88,7 @@ test_that("arg raise_if_empty works", {
     "no data: empty CSV"
   )
   out = pl$read_csv(tmpf, raise_if_empty = FALSE)
-  expect_identical(dim(out), c(0, 0))
+  expect_identical(dim(out), c(0L, 0L))
 })
 
 # TODO: why does this one fail?

diff --git a/tests/testthat/test-s3_methods.R b/tests/testthat/test-s3_methods.R
@@ -264,3 +264,14 @@ test_that("brackets", {
   expect_equal(pl$Series(letters)[1:5]$to_vector(), letters[1:5])
   expect_equal(pl$Series(letters)[-5]$to_vector(), letters[-5])
 })
+
+
+test_that("dim should integer", {
+  d = dim(as_polars_df(mtcars))
+  expect_identical(d, dim(mtcars))
+  expect_true(is.integer(d))
+
+  d = dim(as_polars_lf(mtcars))
+  expect_identical(d, c(NA_integer_, ncol(mtcars)))
+  expect_true(is.integer(d))
+})
diff --git a/tests/testthat/test-sql.R b/tests/testthat/test-sql.R
@@ -16,7 +16,7 @@ test_that("Intialize SQLContext with LazyFrame like objects", {
       UNION ALL
       SELECT * FROM pl_lf"
     )$collect() |> nrow(),
-    3 * nrow(mtcars)
+    3L * nrow(mtcars)
   )
 })
 
@@ -29,7 +29,7 @@ test_that("SQLContext_register, register_many, unregister", {
   expect_equal(ctx$tables(), "mtcars")
   expect_identical(
     ctx$execute("SELECT * FROM mtcars LIMIT 5")$collect() |> nrow(),
-    5
+    5L
   )
 
   ctx$register("mtcars2", mtcars)