chore: Implement zoom v2 with internal dm_zoom2_to(), dm_update_zoom2ed(), dm_insert_zoom2ed(), dm_discard_zoom2ed() (#2399)

* Initial plan

* feat: implement `dm_zoom2_to()`, `dm_update_zoom2ed()`, and `dm_insert_zoom2ed()`

Co-authored-by: krlmlr <1741643+krlmlr@users.noreply.github.com>

* Remove col_tracker attribute, add comprehensive dplyr verb snapshot tests

Co-authored-by: krlmlr <1741643+krlmlr@users.noreply.github.com>

* Rewrite dplyr verb tests: simple explicit tests, expand with right_join, full_join, relocate, group_by, count, 3-table join

Co-authored-by: krlmlr <1741643+krlmlr@users.noreply.github.com>

* Fix test names for consistency

Co-authored-by: krlmlr <1741643+krlmlr@users.noreply.github.com>

* Fix: uuid-based registry for attr-stripping verbs, 3-table joins, add missing verb tests

Co-authored-by: krlmlr <1741643+krlmlr@users.noreply.github.com>

* Use exists() instead of intersect() for uuid cleanup

Co-authored-by: krlmlr <1741643+krlmlr@users.noreply.github.com>

* Drop registry, store zoom2 info in dm_key_info, use cli::cli_abort, add dm_discard_zoom2ed

Co-authored-by: krlmlr <1741643+krlmlr@users.noreply.github.com>

* Don't export for now

* Silent

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: krlmlr <1741643+krlmlr@users.noreply.github.com>
Co-authored-by: Kirill Müller <krlmlr@users.noreply.github.com>
Co-authored-by: Kirill Müller <kirill@cynkra.com>

Author: Copilot

R/deconstruct.R

@@ -54,7 +54,8 @@ new_keyed_tbl <- function(
   uks = NULL,
   fks_in = NULL,
   fks_out = NULL,
-  uuid = NULL
+  uuid = NULL,
+  zoom2 = NULL
 ) {
   check_dots_empty()
 
@@ -68,14 +69,20 @@ new_keyed_tbl <- function(
   }
 
   class(x) <- unique(c("dm_keyed_tbl", class(x)))
-  attr(x, "dm_key_info") <- list(
+  dm_key_info <- list(
     pk = pk,
     uks = uks,
     fks_in = fks_in,
     fks_out = fks_out,
     uuid = uuid
   )
 
+  if (!is.null(zoom2)) {
+    dm_key_info$zoom2 <- zoom2
+  }
+
+  attr(x, "dm_key_info") <- dm_key_info
+
   x
 }
 
@@ -85,7 +92,8 @@ new_keyed_tbl_from_keys_info <- function(tbl, keys_info) {
     pk = keys_info$pk,
     fks_in = keys_info$fks_in,
     fks_out = keys_info$fks_out,
-    uuid = keys_info$uuid
+    uuid = keys_info$uuid,
+    zoom2 = keys_info$zoom2
   )
 }
 

R/dplyr.R

@@ -458,7 +458,8 @@ summarise.dm_keyed_tbl <- function(.data, ..., .by = NULL, .groups = NULL) {
   new_keyed_tbl(
     summarised_tbl,
     pk = new_pk,
-    uuid = keys_info$uuid
+    uuid = keys_info$uuid,
+    zoom2 = keys_info$zoom2
   )
 }
 
@@ -492,7 +493,8 @@ reframe.dm_keyed_tbl <- function(.data, ..., .by = NULL) {
   # so the primary key is not preserved
   new_keyed_tbl(
     reframed_tbl,
-    uuid = keys_info$uuid
+    uuid = keys_info$uuid,
+    zoom2 = keys_info$zoom2
   )
 }
 
@@ -555,7 +557,8 @@ count.dm_keyed_tbl <- function(
   counted_tbl <- count(tbl, ..., wt = {{ wt }}, sort = sort, name = name, .drop = .drop)
   new_keyed_tbl(
     counted_tbl,
-    uuid = keys_info$uuid
+    uuid = keys_info$uuid,
+    zoom2 = keys_info$zoom2
   )
 }
 
@@ -589,7 +592,8 @@ tally.dm_keyed_tbl <- function(x, wt = NULL, sort = FALSE, name = NULL) {
   tallied_tbl <- tally(tbl, wt = {{ wt }}, sort = sort, name = name)
   new_keyed_tbl(
     tallied_tbl,
-    uuid = keys_info$uuid
+    uuid = keys_info$uuid,
+    zoom2 = keys_info$zoom2
   )
 }
 
@@ -729,7 +733,8 @@ left_join.dm_keyed_tbl <- function(
     pk = join_spec$new_pk,
     fks_in = join_spec$new_fks_in,
     fks_out = join_spec$new_fks_out,
-    uuid = join_spec$new_uuid
+    uuid = join_spec$new_uuid,
+    zoom2 = keyed_get_info(x)$zoom2
   )
 }
 
@@ -822,7 +827,8 @@ inner_join.dm_keyed_tbl <- function(
     pk = join_spec$new_pk,
     fks_in = join_spec$new_fks_in,
     fks_out = join_spec$new_fks_out,
-    uuid = join_spec$new_uuid
+    uuid = join_spec$new_uuid,
+    zoom2 = keyed_get_info(x)$zoom2
   )
 }
 
@@ -910,7 +916,8 @@ full_join.dm_keyed_tbl <- function(
     pk = join_spec$new_pk,
     fks_in = join_spec$new_fks_in,
     fks_out = join_spec$new_fks_out,
-    uuid = join_spec$new_uuid
+    uuid = join_spec$new_uuid,
+    zoom2 = keyed_get_info(x)$zoom2
   )
 }
 
@@ -1003,7 +1010,8 @@ right_join.dm_keyed_tbl <- function(
     pk = join_spec$new_pk,
     fks_in = join_spec$new_fks_in,
     fks_out = join_spec$new_fks_out,
-    uuid = join_spec$new_uuid
+    uuid = join_spec$new_uuid,
+    zoom2 = keyed_get_info(x)$zoom2
   )
 }
 
@@ -1198,7 +1206,8 @@ cross_join.dm_keyed_tbl <- function(x, y, ..., copy = NULL, suffix = c(".x", ".y
 
   new_keyed_tbl(
     joined_tbl,
-    uuid = keyed_get_info(x)$uuid
+    uuid = keyed_get_info(x)$uuid,
+    zoom2 = keyed_get_info(x)$zoom2
   )
 }
 

R/zoom2.R

@@ -0,0 +1,155 @@
+#' Mark table for manipulation (v2)
+#'
+#' @description
+#' `r lifecycle::badge("experimental")`
+#'
+#' `dm_zoom2_to()` zooms to the given table, returning a keyed table
+#' that has the dm object as an attribute.
+#' Key column tracking (primary and foreign) is the responsibility
+#' of the `dm_keyed_tbl` object.
+#'
+#' `dm_update_zoom2ed()` overwrites the originally zoomed table
+#' with the manipulated table.
+#'
+#' `dm_insert_zoom2ed()` adds the manipulated table as a new table to the dm.
+#'
+#' `dm_discard_zoom2ed()` discards the zoomed table and returns the
+#' original `dm` as it was before zooming.
+#'
+#' @inheritParams dm_zoom_to
+#'
+#' @return For `dm_zoom2_to()`: A `dm_keyed_tbl` object with zoom2 info
+#'   stored in `dm_key_info`.
+#'
+#' @noRd
+#' @examplesIf rlang::is_installed(c("nycflights13", "DiagrammeR"))
+#' flights_keyed <- dm_zoom2_to(dm_nycflights13(), flights)
+#'
+#' flights_keyed
+#'
+#' flights_keyed_transformed <-
+#'   flights_keyed %>%
+#'   mutate(am_pm_dep = ifelse(dep_time < 1200, "am", "pm"))
+#'
+#' # replace table `flights` with the zoomed table
+#' flights_keyed_transformed %>%
+#'   dm_update_zoom2ed()
+#'
+#' # insert the zoomed table as a new table
+#' flights_keyed_transformed %>%
+#'   dm_insert_zoom2ed("extended_flights")
+#'
+#' # discard changes and return original dm
+#' flights_keyed_transformed %>%
+#'   dm_discard_zoom2ed()
+dm_zoom2_to <- function(dm, table) {
+  check_not_zoomed(dm)
+  table_name <- dm_tbl_name(dm, {{ table }})
+
+  keyed_tables <- dm_get_keyed_tables_impl(dm)
+  keyed_tbl <- keyed_tables[[table_name]]
+
+  # Store zoom2 info inside dm_key_info so it survives all dplyr/tidyr verbs
+  keys_info <- keyed_get_info(keyed_tbl)
+  keys_info$zoom2 <- list(dm = dm, table_name = table_name)
+  attr(keyed_tbl, "dm_key_info") <- keys_info
+
+  keyed_tbl
+}
+
+#' @rdname dm_zoom2_to
+#' @param zoomed_tbl A `dm_keyed_tbl` object returned by `dm_zoom2_to()`
+#'   or modified via dplyr operations.
+#'
+#' @return For `dm_update_zoom2ed()`, `dm_insert_zoom2ed()` and
+#'   `dm_discard_zoom2ed()`: A `dm` object.
+#'
+#' @noRd
+dm_update_zoom2ed <- function(zoomed_tbl) {
+  zoom2_info <- zoom2_get_info(zoomed_tbl)
+  dm <- zoom2_info$dm
+  table_name <- zoom2_info$table_name
+
+  keyed_tables <- dm_get_keyed_tables_impl(dm)
+  keyed_tables[[table_name]] <- zoom2_clean_keys_info(zoomed_tbl)
+
+  # Preserve table order by using names from original keyed_tables
+  new_dm(keyed_tables)
+}
+
+#' @rdname dm_zoom2_to
+#' @param new_tbl_name Name of the new table.
+#' @inheritParams vctrs::vec_as_names
+#'
+#' @noRd
+dm_insert_zoom2ed <- function(zoomed_tbl, new_tbl_name = NULL, repair = "unique", quiet = FALSE) {
+  zoom2_info <- zoom2_get_info(zoomed_tbl)
+  dm <- zoom2_info$dm
+  table_name <- zoom2_info$table_name
+
+  if (is.null(new_tbl_name)) {
+    new_tbl_name <- table_name
+  }
+
+  keyed_tables <- dm_get_keyed_tables_impl(dm)
+
+  # Repair names if needed
+  names_list <- repair_table_names(
+    old_names = names(keyed_tables),
+    new_names = new_tbl_name,
+    repair,
+    quiet
+  )
+
+  # Rename existing tables if name repair occurred
+  names(keyed_tables) <- names_list$old_new_names[names(keyed_tables)]
+  new_tbl_name <- names_list$new_names
+
+  # Add the new table
+  keyed_tables[[new_tbl_name]] <- zoom2_clean_keys_info(zoomed_tbl)
+
+  out <- new_dm(keyed_tables)
+
+  # Transfer color from original table to new table
+  orig_def <- dm_get_def(dm)
+  orig_display <- orig_def$display[orig_def$table == table_name]
+  if (!is.na(orig_display)) {
+    out <- dm_set_colors(out, !!!set_names(new_tbl_name, orig_display))
+  }
+
+  out
+}
+
+#' @rdname dm_zoom2_to
+#' @noRd
+dm_discard_zoom2ed <- function(zoomed_tbl) {
+  zoom2_info <- zoom2_get_info(zoomed_tbl)
+  zoom2_info$dm
+}
+
+
+# Internal helpers --------------------------------------------------------
+
+zoom2_get_info <- function(zoomed_tbl) {
+  if (!is_dm_keyed_tbl(zoomed_tbl)) {
+    cli::cli_abort("This object was not created by {.fn dm_zoom2_to}.")
+  }
+
+  keys_info <- keyed_get_info(zoomed_tbl)
+  zoom2 <- keys_info$zoom2
+
+  if (is.null(zoom2)) {
+    cli::cli_abort("This object was not created by {.fn dm_zoom2_to}.")
+  }
+
+  zoom2
+}
+
+zoom2_clean_keys_info <- function(x) {
+  if (is_dm_keyed_tbl(x)) {
+    keys_info <- keyed_get_info(x)
+    keys_info$zoom2 <- NULL
+    attr(x, "dm_key_info") <- keys_info
+  }
+  x
+}