pandas 3 support: convert Arrow extension columns; lift pandas<3 pin

NEWS.md

@@ -1,9 +1,12 @@
 # nat.python 0.2.0.9000 (development version)
 
-* `simple_python()` now pins the baseline install to `pandas < 3`. pandas 3.0
-  makes Arrow-backed strings the default dtype, which `pandas2df()` does not yet
-  convert back to R; the pin keeps the provisioned environment functional until
-  that support lands (#6).
+* `pandas2df()` now converts pandas extension-array columns that reticulate
+  leaves unconverted, in particular pandas 3.0's default Arrow-backed string
+  dtype (PDEP-14): string columns become R character vectors and other Arrow
+  columns (e.g. `int64[pyarrow]` ids) map to the same R types as their
+  native-dtype equivalents (#6).
+* `simple_python()` installs current pandas (3.x supported), no longer pinned to
+  `pandas < 3`, now that `pandas2df()` handles the Arrow-backed string dtype.
 * CI now provisions Python through `simple_python()` itself (the end-user path),
   rather than a bespoke `reticulate::py_install()` call.
 

R/convert.R

@@ -24,6 +24,11 @@
 #'     read as strings so arbitrary-precision Python ints round-trip. Genuine
 #'     list-valued columns (e.g. multi-select values) are left intact.
 #'   * datetime columns are normalised to `POSIXct` in UTC.
+#'   * pandas extension-array columns that reticulate leaves unconverted -- most
+#'     importantly pandas 3.0's default Arrow-backed string dtype -- are
+#'     recovered: string columns become character vectors, other Arrow columns
+#'     are classified like object columns (so Arrow-backed ids map like native
+#'     ones).
 #'
 #'   The optional `use_arrow` path round-trips through a Feather file and needs
 #'   the Suggested `arrow` package; `bigint` does not apply to it.
@@ -138,6 +143,24 @@ pandas2df_inmem <- function(df, tibble = FALSE, bigint = "auto") {
     res <- tibble::as_tibble(res)
   }
 
+  # pandas extension arrays reticulate cannot convert come back as raw python
+  # objects. The one that now arrives by default is pandas 3.0's Arrow-backed
+  # string dtype (PDEP-14), but any explicit ArrowDtype column lands here too.
+  # Convert each from its string values: a declared string dtype becomes an R
+  # character vector; any other extension column is classified like an object
+  # column, so Arrow-backed ids and numbers map to the same R types as their
+  # native-dtype equivalents.
+  ext_cols <- names(res)[vapply(res, function(v)
+    inherits(v, "python.builtin.object"), logical(1))]
+  for (col in ext_cols) {
+    vals <- pandas_series_character_values(reticulate::py_get_item(df, col))
+    if (is.null(vals)) next
+    res[[col]] <- if (col %in% names(dtypes) && is_string_dtype(dtypes[[col]]))
+      vals
+    else
+      classify_object_values(vals, bigint = bigint, col = col)
+  }
+
   # splice the fast-converted int columns back into their original positions
   splice_fast_int <- function(res) {
     for (col in names(fast_int)) res[[col]] <- fast_int[[col]]
@@ -205,6 +228,13 @@ pandas_py_to_r_frame <- function(df) {
   )
 }
 
+# Does a pandas dtype string denote a string dtype? Covers the StringDtype
+# family ("string", "string[python]", "string[pyarrow]"), Arrow large strings,
+# and pandas 3.0's new default string dtype, which reports simply as "str".
+is_string_dtype <- function(dt) {
+  grepl("^(str|string|large_string)(\\[.*\\])?$", tolower(dt))
+}
+
 pandas_dataframe_dtypes <- function(df) {
   dtype_series <- reticulate::py_get_attr(df, "dtypes")
   dtype_strings <- reticulate::py_call(

R/env.R

@@ -77,12 +77,11 @@ simple_python <- function(pyinstall = c("basic", "full", "extra", "minimal",
     # nat.python's own baseline: pandas2df() needs pandas, and numpy rides in
     # with it. Every richer bundle builds on top of this.
     #
-    # Pinned to pandas < 3 for now: pandas 3.0 makes Arrow-backed strings the
-    # default dtype (PDEP-14), which pandas2df() does not yet convert back to R
-    # (string columns return as raw ArrowStringArray objects). Lift this once
-    # pandas2df() handles the Arrow string dtype.
-    cli::cli_inform("Installing pandas (<3 for now; brings numpy)")
-    ourpip("pandas<3")
+    # No version pin: pandas 3.0 makes Arrow-backed strings the default dtype
+    # (PDEP-14), and pandas2df() converts those (and other Arrow extension
+    # columns) back to native R vectors, so current pandas is fine.
+    cli::cli_inform("Installing pandas (brings numpy)")
+    ourpip("pandas")
   }
   if (pyinstall %in% c("basic", "full", "extra")) {
     cli::cli_inform("Installing cloudvolume")

tests/testthat/test-convert.R

@@ -198,3 +198,50 @@ test_that("pandas2df fast-path preserves column order and nullable ids", {
   expect_identical(out$n, c(1L, 2L, 3L))
   expect_identical(out$label, c("a", "b", "c"))
 })
+
+test_that("is_string_dtype recognises the string dtype family", {
+  expect_true(is_string_dtype("str"))                 # pandas 3.0 default
+  expect_true(is_string_dtype("string"))
+  expect_true(is_string_dtype("string[python]"))
+  expect_true(is_string_dtype("string[pyarrow]"))
+  expect_true(is_string_dtype("large_string[pyarrow]"))
+  expect_false(is_string_dtype("object"))
+  expect_false(is_string_dtype("int64[pyarrow]"))
+  expect_false(is_string_dtype("struct"))             # not a string dtype
+})
+
+test_that("pandas2df converts Arrow-backed string columns to character", {
+  skip_if_no_module("pandas")
+  skip_if_no_module("pyarrow")
+  # An explicit Arrow-backed string column reproduces pandas 3.0's default
+  # dtype (an ArrowStringArray) even under pandas 2: reticulate leaves it as a
+  # raw python object, and pandas2df must recover an R character vector (with a
+  # missing cell -> NA), while a neighbouring numeric column is unaffected.
+  df <- reticulate::py_eval(
+    paste0("__import__('pandas').DataFrame({",
+           "'label': __import__('pandas').array(",
+           "['a', 'b', None], dtype='string[pyarrow]'), ",
+           "'n': [1, 2, 3]})"),
+    convert = FALSE)
+  out <- pandas2df(df)
+  expect_type(out$label, "character")
+  expect_identical(out$label, c("a", "b", NA))
+  expect_identical(out$n, c(1L, 2L, 3L))
+})
+
+test_that("pandas2df converts Arrow-backed integer id columns", {
+  skip_if_no_module("pandas")
+  skip_if_no_module("pyarrow")
+  # A general (non-string) Arrow extension column: 64-bit ids in an
+  # int64[pyarrow] column must map to the same integer64 the native int64 path
+  # produces, not survive as a raw python object.
+  ids <- c("720575940621039145", "720575940626877799")
+  df <- reticulate::py_eval(
+    paste0("__import__('pandas').DataFrame({",
+           "'id': __import__('pandas').array([",
+           paste(ids, collapse = ", "), "], dtype='int64[pyarrow]')})"),
+    convert = FALSE)
+  out <- pandas2df(df)
+  expect_s3_class(out$id, "integer64")
+  expect_identical(as.character(out$id), ids)
+})