diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index af863e8bb..699f56aa1 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -39,7 +39,9 @@ jobs: - name: Verify sess source identity and bootstrap run: pnpm run test:sess-source - name: Install R test dependencies - run: install.packages(c("remotes", "pkgbuild", "languageserver", "lintr", "renv")) + run: | + install.packages(c("remotes", "pkgbuild", "languageserver", "lintr", "renv", + "tinytest", "arrow", "DBI", "RSQLite", "dbplyr", "dplyr", "bit64")) shell: Rscript {0} - run: pnpm run build - name: Verify installed sess identity @@ -47,8 +49,12 @@ jobs: source("R/sess_source.R") stopifnot(!sess_install_required("dist/resources/sess")) shell: Rscript {0} - - name: Run sess tests - run: tinytest::test_package("sess") + - name: Run sess tests including lazy Data Viewer regressions + run: | + stopifnot(all(vapply(c("arrow", "DBI", "RSQLite", "dbplyr", "dplyr", "bit64"), + requireNamespace, logical(1), quietly = TRUE))) + results <- tinytest::test_package("sess") + stopifnot(all(as.logical(results))) shell: Rscript {0} - name: Compile test files run: pnpm run clean && pnpm exec tsc diff --git a/sess/DESCRIPTION b/sess/DESCRIPTION index 77a5b5f43..f31354e5b 100644 --- a/sess/DESCRIPTION +++ b/sess/DESCRIPTION @@ -1,7 +1,7 @@ Package: sess Type: Package Title: High-Performance IPC Bridge for R Sessions -Version: 3.2.0 +Version: 3.2.0.9000 Authors@R: c( person(given = "Randy", family = "Lai", @@ -39,7 +39,12 @@ Suggests: httpgd, jgd, svglite, - tinytest + tinytest, + arrow, + DBI, + RSQLite, + dbplyr, + dplyr NeedsCompilation: no Config/vscode-R/Interactive: 1 Config/roxygen2/markdown: TRUE diff --git a/sess/R/dataview-arrow.R b/sess/R/dataview-arrow.R new file mode 100644 index 000000000..0c10377de --- /dev/null +++ b/sess/R/dataview-arrow.R @@ -0,0 +1,965 @@ +# Lazy Arrow support for the data viewer + +dataview_is_arrow_lazy <- function(data) { + inherits(data, "Dataset") || + inherits(data, "arrow_dplyr_query") +} + +dataview_arrow_reader_batch_size <- 5000L +dataview_arrow_cache_block_size <- 1000L +dataview_arrow_sort_block_size <- 5000L +dataview_arrow_cache_rows <- 20000L + +dataview_arrow_require <- function() { + if (!requireNamespace("arrow", quietly = TRUE)) { + stop("Viewing Arrow datasets requires the optional 'arrow' package") + } +} + +dataview_arrow_nested_columns <- function(data) { + schema <- if (inherits(data, "Schema")) data else getExportedValue("arrow", "infer_schema")(data) + vapply(schema$fields, function(field) { + inherits( + field$type, + c("ListType", "LargeListType", "FixedSizeListType", "MapType", "StructType") + ) + }, logical(1)) +} + +dataview_arrow_conversion <- function(data) { + if (is.data.frame(data)) return(integer()) + schema <- getExportedValue("arrow", "infer_schema")(data) + which(vapply(schema$fields, function(field) { + inherits( + field$type, c("Int64Type", "Date32Type", "Date64Type", "TimestampType", "DurationType") + ) + }, logical(1))) +} + +dataview_arrow_data_frame <- function(data, conversion = NULL) { + old_options <- options(arrow.int64_downcast = FALSE) + on.exit(options(old_options), add = TRUE) + + page <- as.data.frame(data, optional = TRUE) + if (inherits(data, "ArrowTabular")) { + if (is.null(conversion)) conversion <- dataview_arrow_conversion(data) + for (position in conversion) { + values <- tryCatch(as.vector(data[[position]]), error = function(e) NULL) + if (!is.null(values) && length(values) == nrow(page)) { + page[[position]] <- values + } + } + } + for (position in seq_len(ncol(page))) { + if (inherits(page[[position]], "vctrs_list_of")) { + page[[position]] <- as.list(page[[position]]) + } + } + page +} + +dataview_arrow_bind_column <- function(values) { + template <- values[[1L]] + if (is.data.frame(template)) { + # Struct columns contain rows, not a list of fields to concatenate. + return(dataview_arrow_bind_pages(values)) + } + if (inherits(template, "integer64")) { + value <- unlist(lapply(values, unclass), use.names = FALSE) + class(value) <- "integer64" + return(value) + } + if (inherits(template, "Date")) { + return(structure( + unlist(lapply(values, unclass), use.names = FALSE), + class = "Date" + )) + } + if (inherits(template, "POSIXct")) { + value <- structure( + unlist(lapply(values, unclass), use.names = FALSE), + class = class(template) + ) + attr(value, "tzone") <- attr(template, "tzone", exact = TRUE) + return(value) + } + if (inherits(template, "difftime")) { + return(as.difftime( + unlist(lapply(values, unclass), use.names = FALSE), + units = attr(template, "units", exact = TRUE) + )) + } + do.call(c, values) +} + +dataview_arrow_bind_pages <- function(pages) { + if (length(pages) == 1L) { + return(pages[[1L]]) + } + + columns <- lapply(seq_len(ncol(pages[[1L]])), function(position) { + dataview_arrow_bind_column(lapply(pages, function(page) page[[position]])) + }) + names(columns) <- names(pages[[1L]]) + n <- sum(vapply(pages, nrow, integer(1))) + structure( + columns, + class = "data.frame", + row.names = c(NA_integer_, -n) + ) +} + +dataview_arrow_reader_state <- function(data = NULL) { + dataview_arrow_require() + state <- new.env(parent = emptyenv()) + state$reader <- NULL + # Current RecordBatch buffer from the forward Arrow reader. + state$batch <- NULL + state$batch_row <- 0L + state$next_row <- 1L + # Source-row cache used by normal forward/backward scrolling. + state$row_cache <- list() + # Display-row cache keyed by filter/sort query and display block. + state$query_cache <- list() + state$fragment_index <- NULL + state$row_group_index <- NULL + state$source_conversion <- if (!is.null(data)) dataview_arrow_conversion(data) else NULL + state$conversion <- state$source_conversion + state$projection <- NULL + state$page_data <- NULL + state$result_cache_dir <- NULL + state$result_cache_batches <- list() + state$result_cache_rows <- 0L + state$result_complete <- FALSE + state +} + +dataview_arrow_result_cache_init <- function(state, view_id) { + if (!inherits(state$data, "arrow_dplyr_query")) return(invisible(NULL)) + + root <- file.path(tempdir(), "vscode-r-dataview") + dir.create(root, recursive = TRUE, showWarnings = FALSE) + result_cache_dir <- tempfile(pattern = paste0(view_id, "-"), tmpdir = root) + if (!dir.create(result_cache_dir, showWarnings = FALSE)) { + stop("unable to create temporary Arrow data viewer cache") + } + + state$arrow_reader$result_cache_dir <- result_cache_dir + state$arrow_reader$result_cache_batches <- list() + state$arrow_reader$result_cache_rows <- 0L + state$arrow_reader$result_complete <- FALSE + state$arrow_reader$result_query_key <- NULL + state$arrow_reader$result_query_indices <- NULL + invisible(NULL) +} + +dataview_arrow_result_cache_cleanup <- function(state) { + if (is.null(state$arrow_reader)) return(invisible(NULL)) + + dataview_arrow_reader_reset(state) + result_cache_dir <- state$arrow_reader$result_cache_dir + if (!is.null(result_cache_dir) && dir.exists(result_cache_dir)) { + unlink(result_cache_dir, recursive = TRUE, force = TRUE) + } + state$arrow_reader$result_cache_dir <- NULL + state$arrow_reader$result_cache_batches <- list() + state$arrow_reader$result_cache_rows <- 0L + state$arrow_reader$result_complete <- FALSE + state$arrow_reader$result_query_key <- NULL + state$arrow_reader$result_query_indices <- NULL + invisible(NULL) +} + +dataview_arrow_cache_get <- function(reader_state, row_idx) { + # Reuse cached source rows for normal scrolling without rescanning Arrow. + if (!length(row_idx)) { + return(NULL) + } + + for (i in seq_along(reader_state$row_cache)) { + cached <- reader_state$row_cache[[i]] + if (min(row_idx) < cached$first_row || max(row_idx) > cached$last_row) { + next + } + + cached_idx <- match(row_idx, cached$row_idx) + if (all(!is.na(cached_idx))) { + reader_state$row_cache <- c( + reader_state$row_cache[-i], + list(cached) + ) + return(cached$data[cached_idx, , drop = FALSE]) + } + } + NULL +} + +dataview_arrow_cache_add <- function(reader_state, row_idx, data) { + reader_state$row_cache[[length(reader_state$row_cache) + 1L]] <- list( + first_row = min(row_idx), + last_row = max(row_idx), + row_idx = row_idx, + data = data + ) + + while (length(reader_state$row_cache) > 1L && + sum(vapply( + reader_state$row_cache, + function(cached) length(cached$row_idx), + integer(1) + )) > dataview_arrow_cache_rows) { + reader_state$row_cache <- reader_state$row_cache[-1L] + } +} + +dataview_arrow_query_cache_get <- function(reader_state, query_key, block_start) { + # Reuse cached display blocks for the current filter/sort result. + for (i in rev(seq_along(reader_state$query_cache))) { + cached <- reader_state$query_cache[[i]] + if (!identical(cached$query_key, query_key) || + cached$block_start != block_start) { + next + } + reader_state$query_cache <- c( + reader_state$query_cache[-i], + list(cached) + ) + return(cached$data) + } + NULL +} + +dataview_arrow_query_cache_add <- function( + reader_state, + query_key, + block_start, + data +) { + reader_state$query_cache[[length(reader_state$query_cache) + 1L]] <- list( + query_key = query_key, + block_start = block_start, + data = data + ) + while (length(reader_state$query_cache) > 1L && + sum(vapply( + reader_state$query_cache, + function(cached) nrow(cached$data), + integer(1) + )) > dataview_arrow_cache_rows) { + reader_state$query_cache <- reader_state$query_cache[-1L] + } +} + +dataview_arrow_page_state <- function(state, fields = NULL) { + positions <- dataview_page_positions(state, fields) + reader_state <- state$arrow_reader + + if (inherits(state$data, "arrow_dplyr_query")) { + if (!identical(positions, reader_state$projection)) { + reader_state$row_cache <- list() + reader_state$query_cache <- list() + reader_state$projection <- positions + } + state$arrow_source <- state$data + state$columns <- state$columns[c(1L, positions + 1L)] + state$column_fields <- as.character(positions) + return(state) + } + + if (!identical(positions, reader_state$projection)) { + dataview_arrow_reader_reset(state) + reader_state$row_cache <- list() + reader_state$query_cache <- list() + reader_state$page_data <- if (length(positions) && length(positions) < ncol(state$data)) { + state$data$WithSchema(state$data$schema[names(state$data)[positions]]) + } else { + state$data + } + reader_state$projection <- positions + conversion <- match(reader_state$source_conversion, positions, nomatch = 0L) + reader_state$conversion <- conversion[conversion > 0L] + } + state$arrow_source <- state$data + state$data <- reader_state$page_data + state$columns <- state$columns[c(1L, positions + 1L)] + state$column_fields <- as.character(positions) + state +} + +dataview_arrow_query_columns <- function(state, sort_model, filter_model) { + positions <- dataview_query_positions(state, sort_model, filter_model) + if (!length(positions)) return(list()) + names <- names(state$data)[positions - 1L] + reader <- if (inherits(state$data, "arrow_dplyr_query")) { + dataview_arrow_reader_open( + dplyr::select( + dplyr::ungroup(dplyr::collapse(state$data)), + dplyr::all_of(names) + ) + ) + } else { + getExportedValue("arrow", "Scanner")$create( + state$data, + projection = names, + batch_size = dataview_arrow_reader_batch_size + )$ToRecordBatchReader() + } + on.exit(try(reader$Close(), silent = TRUE), add = TRUE) + conversion <- match(state$arrow_reader$source_conversion, positions - 1L, nomatch = 0L) + columns <- dataview_arrow_data_frame(reader$read_table(), conversion[conversion > 0L]) + setNames(as.list(columns), as.character(positions)) +} + +dataview_arrow_reader_reset <- function(state, clear_cache = FALSE) { + reader_state <- state$arrow_reader + if (!is.null(reader_state$reader)) { + try(reader_state$reader$Close(), silent = TRUE) + } + reader_state$reader <- NULL + reader_state$batch <- NULL + reader_state$batch_row <- 0L + reader_state$next_row <- 1L + if (clear_cache) { + for (batch in reader_state$result_cache_batches) { + if (!is.null(batch$path)) unlink(batch$path, force = TRUE) + } + reader_state$result_cache_batches <- list() + reader_state$result_cache_rows <- 0L + reader_state$result_complete <- FALSE + reader_state$result_query_key <- NULL + reader_state$result_query_indices <- NULL + reader_state$row_cache <- list() + reader_state$query_cache <- list() + } +} + +dataview_arrow_reader_open <- function(data) { + if (inherits(data, "arrow_dplyr_query")) { + # Scanner$create(query) does not execute ordering or aggregation nodes. + return(getExportedValue("arrow", "as_record_batch_reader")(data)) + } + Scanner <- getExportedValue("arrow", "Scanner") + Scanner$create(data, batch_size = dataview_arrow_reader_batch_size)$ToRecordBatchReader() +} + +dataview_arrow_column <- function(data, position) { + old_options <- options(arrow.int64_downcast = FALSE) + on.exit(options(old_options), add = TRUE) + + name <- names(data)[[position]] + Scanner <- getExportedValue("arrow", "Scanner") + reader <- Scanner$create( + data, + projection = name, + batch_size = dataview_arrow_reader_batch_size + )$ToRecordBatchReader() + on.exit(try(reader$Close(), silent = TRUE), add = TRUE) + dataview_arrow_data_frame(reader$read_table())[[name]] +} + +dataview_arrow_result_cache_add <- function(reader_state, batch) { + if (is.null(reader_state$result_cache_dir) || !nrow(batch)) { + return(invisible(NULL)) + } + + batch_number <- length(reader_state$result_cache_batches) + 1L + first_row <- reader_state$next_row + last_row <- first_row + nrow(batch) - 1L + reader_state$result_cache_batches[[batch_number]] <- list( + data = batch, + path = NULL, + first_row = first_row, + last_row = last_row + ) + reader_state$result_cache_rows <- last_row + + in_memory_rows <- function() { + sum(vapply( + reader_state$result_cache_batches, + function(cached) { + if (is.null(cached$data)) 0L else cached$last_row - cached$first_row + 1L + }, + integer(1) + )) + } + + while (length(reader_state$result_cache_batches) > 1L && + in_memory_rows() > dataview_arrow_cache_rows) { + position <- which(vapply( + reader_state$result_cache_batches, + function(cached) !is.null(cached$data), + logical(1) + ))[[1L]] + cached <- reader_state$result_cache_batches[[position]] + path <- file.path( + reader_state$result_cache_dir, + sprintf("batch-%08d.arrow", position) + ) + temporary_path <- paste0(path, ".tmp") + on.exit(unlink(temporary_path, force = TRUE), add = TRUE) + getExportedValue("arrow", "write_ipc_file")( + cached$data, temporary_path + ) + if (!file.rename(temporary_path, path)) { + stop("unable to finalize temporary Arrow data viewer cache") + } + cached$data <- NULL + cached$path <- path + reader_state$result_cache_batches[[position]] <- cached + } + invisible(NULL) +} + +dataview_arrow_reader_ensure_batch <- function(reader_state) { + failed <- function(e) { + # A failed reader may return EOF on its next call. Never publish its partial + # result as complete, or append a new execution to the old cached rows. + dataview_arrow_reader_reset(list(arrow_reader = reader_state), clear_cache = TRUE) + stop(e) + } + tryCatch({ + while (is.null(reader_state$batch) || + reader_state$batch_row >= nrow(reader_state$batch)) { + reader_state$batch <- reader_state$reader$read_next_batch() + reader_state$batch_row <- 0L + if (is.null(reader_state$batch)) { + reader_state$result_complete <- TRUE + return(FALSE) + } + dataview_arrow_result_cache_add(reader_state, reader_state$batch) + } + TRUE + }, error = failed, interrupt = failed) +} +dataview_arrow_reader_take <- function(reader_state, n, collect = TRUE) { + pages <- list() + while (n > 0L) { + if (!dataview_arrow_reader_ensure_batch(reader_state)) { + break + } + + take <- min(n, nrow(reader_state$batch) - reader_state$batch_row) + if (collect) { + pages[[length(pages) + 1L]] <- + dataview_arrow_data_frame( + reader_state$batch$Slice(reader_state$batch_row, take), reader_state$conversion + ) + } + reader_state$batch_row <- reader_state$batch_row + take + reader_state$next_row <- reader_state$next_row + take + n <- n - take + } + + if (!collect) { + return(invisible(NULL)) + } + if (length(pages) == 1L) { + return(pages[[1L]]) + } + dataview_arrow_bind_pages(pages) +} + +dataview_arrow_reader_select <- function(reader_state, row_idx) { + pages <- list() + while (length(row_idx)) { + if (row_idx[[1L]] > reader_state$next_row) { + dataview_arrow_reader_take( + reader_state, + row_idx[[1L]] - reader_state$next_row, + collect = FALSE + ) + } + if (!dataview_arrow_reader_ensure_batch(reader_state)) { + break + } + + batch_last <- reader_state$next_row + + nrow(reader_state$batch) - reader_state$batch_row - 1L + selected <- row_idx[row_idx <= batch_last] + positions <- reader_state$batch_row + + selected - reader_state$next_row + 1L + pages[[length(pages) + 1L]] <- + dataview_arrow_data_frame( + reader_state$batch[positions, , drop = FALSE], reader_state$conversion + ) + + dataview_arrow_reader_take( + reader_state, + selected[[length(selected)]] - reader_state$next_row + 1L, + collect = FALSE + ) + row_idx <- row_idx[row_idx > batch_last] + } + + # At a consumed batch boundary, check EOF without scanning the query. If a + # next batch exists, ensure_batch retains it in the cache for later pages. + if (!is.null(reader_state$result_cache_dir) && + !is.null(reader_state$batch) && + reader_state$batch_row >= nrow(reader_state$batch)) { + dataview_arrow_reader_ensure_batch(reader_state) + } + + if (!length(pages)) { + return(NULL) + } + if (length(pages) == 1L) { + return(pages[[1L]]) + } + dataview_arrow_bind_pages(pages) +} + +dataview_arrow_result_cache_slice <- function(state, row_idx, projection = NULL) { + reader_state <- state$arrow_reader + if (is.null(reader_state$result_cache_dir) || !length(row_idx)) return(NULL) + + target <- max(row_idx) + while (target > reader_state$result_cache_rows && + !isTRUE(reader_state$result_complete)) { + if (is.null(reader_state$reader)) { + reader_state$reader <- dataview_arrow_reader_open(state$data) + } + if (!dataview_arrow_reader_ensure_batch(reader_state)) break + if (target <= reader_state$result_cache_rows) break + + remaining <- nrow(reader_state$batch) - reader_state$batch_row + if (remaining > 0L) { + dataview_arrow_reader_take(reader_state, remaining, collect = FALSE) + } + } + + available <- row_idx[row_idx <= reader_state$result_cache_rows] + if (is.null(projection)) projection <- reader_state$projection + if (is.null(projection)) projection <- seq_len(ncol(state$data)) + fields <- names(state$data)[projection] + if (!length(available)) { + schema <- dataview_schema(state$data) + return(schema[integer(), projection, drop = FALSE]) + } + + pages <- list() + page_rows <- integer() + + for (batch in reader_state$result_cache_batches) { + selected <- available[ + available >= batch$first_row & available <= batch$last_row + ] + if (!length(selected)) next + + positions <- selected - batch$first_row + 1L + page <- if (!length(fields)) { + data.frame(row.names = seq_along(selected)) + } else if (!is.null(batch$data)) { + dataview_arrow_data_frame( + batch$data[positions, projection, drop = FALSE] + ) + } else { + data <- getExportedValue("arrow", "read_ipc_file")( + batch$path, + col_select = fields, + as_data_frame = FALSE + ) + dataview_arrow_data_frame(data$Take(positions - 1L)) + } + pages[[length(pages) + 1L]] <- page + page_rows <- c(page_rows, selected) + } + + if (!length(pages)) return(NULL) + page <- if (length(pages) == 1L) { + pages[[1L]] + } else { + dataview_arrow_bind_pages(pages) + } + page[match(available, page_rows), , drop = FALSE] +} +dataview_arrow_fragment_index <- function(state) { + reader_state <- state$arrow_reader + if (isFALSE(reader_state$fragment_index)) { + return(NULL) + } + if (!is.null(reader_state$fragment_index)) { + return(reader_state$fragment_index) + } + + data <- state$arrow_source %||% state$data + files <- if (inherits(data, "FileSystemDataset")) data$files else character() + if (length(files) < 2L) { + reader_state$fragment_index <- FALSE + return(NULL) + } + + FileSystemDatasetFactory <- getExportedValue( + "arrow", "FileSystemDatasetFactory" + ) + filesystem <- data$filesystem + format <- data$format + schema <- data$schema + fields <- names(schema) + missing_columns <- vector("list", length(files)) + fragments <- lapply(seq_along(files), function(i) { + tryCatch( + { + factory <- FileSystemDatasetFactory$create( + filesystem, + paths = files[[i]], + format = format + ) + missing_columns[[i]] <<- setdiff(fields, names(factory$Inspect())) + factory$Finish(schema = schema) + }, + error = function(e) NULL + ) + }) + if (any(vapply(fragments, is.null, logical(1)))) { + reader_state$fragment_index <- FALSE + return(NULL) + } + + counts <- vapply( + fragments, + function(fragment) as.numeric(fragment$num_rows), + numeric(1) + ) + if (any(!is.finite(counts)) || sum(counts) != state$total_rows) { + reader_state$fragment_index <- FALSE + return(NULL) + } + + keep <- counts > 0 + fragments <- fragments[keep] + missing_columns <- missing_columns[keep] + counts <- counts[keep] + ends <- cumsum(counts) + starts <- if (length(counts)) c(1, head(ends, -1L) + 1) else numeric() + + # R does not expose a Dataset's fragment partition expressions. A field + # absent from a physical file is either a partition constant or a null field + # introduced by schema unification. Sample those values from the ORIGINAL + # dataset, retaining directory/Hive partitioning and the declared field types. + # Only one row per nonempty file and the missing fields are materialized. + missing_names <- unique(unlist(missing_columns, use.names = FALSE)) + partition_values <- if (length(missing_names)) { + dataview_slice(data$WithSchema(schema[missing_names]), starts) + } else { + NULL + } + + reader_state$fragment_index <- list( + datasets = fragments, + starts = starts, + ends = ends, + missing_columns = missing_columns, + partition_values = partition_values + ) + reader_state$fragment_index +} + +dataview_arrow_fragment_slice <- function(state, row_idx) { + # Read requested rows from only the dataset files that contain them. + index <- dataview_arrow_fragment_index(state) + if (is.null(index) || !length(row_idx)) { + return(NULL) + } + + fragment_pos <- findInterval(row_idx - 1, index$ends) + 1L + if (any(fragment_pos < 1L | fragment_pos > length(index$datasets))) { + return(NULL) + } + + # Use one Arrow scan over ONLY the selected files, in their source order. + # This avoids one scanner/conversion per file for widely scattered sorts. + fragments <- sort(unique(fragment_pos)) + counts <- index$ends[fragments] - index$starts[fragments] + 1 + starts <- c(1, head(cumsum(counts), -1L) + 1) + selected_idx <- row_idx - index$starts[fragment_pos] + + starts[match(fragment_pos, fragments)] + data <- if (length(fragments) == 1L) { + index$datasets[[fragments]] + } else { + getExportedValue("arrow", "open_dataset")(index$datasets[fragments]) + } + if (!is.null(state$column_fields)) data <- data$WithSchema(state$data$schema) + page <- dataview_slice(data, selected_idx, state$arrow_reader$conversion) + + for (name in intersect(names(index$partition_values), names(page))) { + missing <- vapply(index$missing_columns, function(fields) name %in% fields, logical(1)) + positions <- which(missing[fragment_pos]) + if (is.data.frame(page[[name]])) { + page[[name]][positions, ] <- + index$partition_values[[name]][fragment_pos[positions], , drop = FALSE] + } else { + page[[name]][positions] <- index$partition_values[[name]][fragment_pos[positions]] + } + } + page +} + +dataview_arrow_row_group_index <- function(state) { + reader_state <- state$arrow_reader + if (isFALSE(reader_state$row_group_index)) return(NULL) + if (!is.null(reader_state$row_group_index)) return(reader_state$row_group_index) + + data <- state$arrow_source %||% state$data + files <- if (inherits(data, "FileSystemDataset")) data$files else character() + if (length(files) != 1L || !identical(data$format$type, "parquet")) { + reader_state$row_group_index <- FALSE + return(NULL) + } + + reader <- tryCatch( + getExportedValue("arrow", "ParquetFileReader")$create( + data$filesystem$OpenInputFile(files[[1L]]) + ), + error = function(e) NULL + ) + if (is.null(reader) || reader$num_row_groups < 2L || + any(dataview_arrow_nested_columns(reader$GetSchema()))) { + # ReadRowGroups indexes Parquet leaf columns, not top-level Arrow fields. + # Inspect the physical schema even if the Dataset hides nested columns. + reader_state$row_group_index <- FALSE + return(NULL) + } + + groups <- seq_len(reader$num_row_groups) - 1L + counts <- vapply( + groups, + function(group) as.numeric(reader$ReadRowGroup(group, integer())$num_rows), + numeric(1) + ) + if (any(!is.finite(counts)) || sum(counts) != state$total_rows) { + reader_state$row_group_index <- FALSE + return(NULL) + } + + ends <- cumsum(counts) + starts <- c(1, head(ends, -1L) + 1) + fields <- names(reader$GetSchema()) + missing_names <- setdiff(names(data), fields) + partition_values <- if (length(missing_names)) { + dataview_slice(data$WithSchema(data$schema[missing_names]), 1L) + } else { + NULL + } + + reader_state$row_group_index <- list( + reader = reader, + starts = starts, + ends = ends, + fields = fields, + partition_values = partition_values + ) + reader_state$row_group_index +} + +dataview_arrow_row_group_slice <- function(state, row_idx) { + # Read requested rows directly from the Parquet row groups that contain them. + index <- dataview_arrow_row_group_index(state) + if (is.null(index) || !length(row_idx)) return(NULL) + + group_pos <- findInterval(row_idx - 1, index$ends) + 1L + if (any(group_pos < 1L | group_pos > length(index$ends))) return(NULL) + + groups <- sort(unique(group_pos)) + counts <- index$ends[groups] - index$starts[groups] + 1 + starts <- c(1, head(cumsum(counts), -1L) + 1) + selected_idx <- row_idx - index$starts[group_pos] + + starts[match(group_pos, groups)] + + fields <- names(state$data) + columns <- match(fields, index$fields) + present <- !is.na(columns) + page <- if (any(present)) { + table <- index$reader$ReadRowGroups( + groups - 1L, as.integer(columns[present] - 1L) + ) + # Select in Arrow so only requested rows are materialized as R columns. + dataview_arrow_data_frame(table$Take(selected_idx - 1L)) + } else { + data.frame(row.names = seq_along(row_idx)) + } + + for (name in fields[!present]) { + value <- index$partition_values[[name]] + if (is.null(value)) return(NULL) + page[[name]] <- if (is.data.frame(value)) { + value[rep(1L, nrow(page)), , drop = FALSE] + } else { + value[rep(1L, nrow(page))] + } + } + page[, fields, drop = FALSE] +} + +dataview_arrow_query_forward_slice <- function(state, row_idx) { + # Reuse the forward reader when requested source rows can be reached efficiently. + if (!length(row_idx) || + (length(row_idx) > 1L && any(diff(row_idx) <= 0L))) { + return(NULL) + } + + reader_state <- state$arrow_reader + first_row <- row_idx[[1L]] + if (isTRUE(reader_state$result_complete)) { + return(NULL) + } + if (first_row < reader_state$next_row) { + return(NULL) + } + if (inherits(state$data, "FileSystemDataset") && + (!isFALSE(reader_state$fragment_index) || + !isFALSE(reader_state$row_group_index)) && + tail(row_idx, 1L) - reader_state$next_row > + length(row_idx) + dataview_arrow_reader_batch_size) { + # A distant or sparse request should not consume all preceding data. + return(NULL) + } + if (is.null(reader_state$reader)) { + reader_state$reader <- dataview_arrow_reader_open(state$data) + } + if (first_row > reader_state$next_row) { + dataview_arrow_reader_take( + reader_state, + first_row - reader_state$next_row, + collect = FALSE + ) + } + dataview_arrow_reader_select(reader_state, row_idx) +} + +dataview_arrow_query_fetch <- function(state, row_idx) { + # Choose the cheapest Arrow path for fetching the requested source rows. + if (!length(row_idx)) { + return(dataview_schema(state$data)) + } + + if (inherits(state$data, "arrow_dplyr_query")) { + page <- dataview_arrow_query_forward_slice(state, row_idx) + if (!is.null(page)) { + projection <- state$arrow_reader$projection + if (is.null(projection)) projection <- seq_len(ncol(page)) + return(page[, projection, drop = FALSE]) + } + + return(dataview_arrow_result_cache_slice(state, row_idx)) + } + + if (!is.data.frame(state$data)) { + page <- dataview_arrow_query_forward_slice(state, row_idx) + if (!is.null(page)) { + return(page) + } + page <- dataview_arrow_fragment_slice(state, row_idx) + if (!is.null(page)) { + return(page) + } + page <- dataview_arrow_row_group_slice(state, row_idx) + if (!is.null(page)) { + return(page) + } + } + dataview_slice(state$data, row_idx, state$arrow_reader$conversion) +} + +dataview_arrow_block_slice <- function(state, row_idx) { + # Fetch contiguous source rows in cache-sized blocks for normal scrolling. + reader_state <- state$arrow_reader + pages <- list() + block_starts <- dataview_block_starts( + row_idx, dataview_arrow_cache_block_size + ) + + for (block_start in block_starts) { + block_end <- min( + state$total_rows, + block_start + dataview_arrow_cache_block_size - 1L + ) + block_row_idx <- seq.int(block_start, block_end) + cached <- dataview_arrow_cache_get(reader_state, block_row_idx) + + if (is.null(cached)) { + cached <- dataview_arrow_query_fetch(state, block_row_idx) + dataview_arrow_cache_add(reader_state, block_row_idx, cached) + } + + selected <- row_idx[row_idx >= block_start & row_idx <= block_end] + pages[[length(pages) + 1L]] <- cached[ + selected - block_start + 1L, + , + drop = FALSE + ] + } + + if (length(pages) == 1L) { + return(pages[[1L]]) + } + dataview_arrow_bind_pages(pages) +} + +dataview_arrow_slice <- function(state, row_idx) { + # Fetch source rows directly, using block caching for contiguous requests. + if (!length(row_idx)) { + schema <- dataview_schema(state$data) + if (inherits(state$data, "arrow_dplyr_query")) { + positions <- state$arrow_reader$projection %||% seq_len(ncol(schema)) + schema <- schema[, positions, drop = FALSE] + } + return(schema) + } + reader_state <- state$arrow_reader + + if (inherits(state$data, "arrow_dplyr_query")) { + return(dataview_arrow_query_fetch(state, row_idx)) + } + + if (length(row_idx) == 1L || all(diff(row_idx) == 1L)) { + return(dataview_arrow_block_slice(state, row_idx)) + } + + cached <- dataview_arrow_cache_get(reader_state, row_idx) + if (!is.null(cached)) { + return(cached) + } + + page <- dataview_arrow_query_fetch(state, row_idx) + dataview_arrow_cache_add(reader_state, row_idx, page) + page +} + +dataview_arrow_query_slice <- function(state, display_idx) { + # Fetch filtered/sorted display rows using query-specific cached blocks. + if (!length(display_idx)) { + return(dataview_schema(state$data)) + } + + # Prefetch larger display blocks; random Dataset pages use file fragments or + # Parquet row groups when available. + reader_state <- state$arrow_reader + block_size <- if (isTRUE(state$query_has_sort)) { + dataview_arrow_sort_block_size + } else { + dataview_arrow_cache_block_size + } + block_starts <- dataview_block_starts(display_idx, block_size) + pages <- lapply(block_starts, function(block_start) { + block_end <- min( + length(state$query_indices), + block_start + block_size - 1L + ) + cached <- dataview_arrow_query_cache_get( + reader_state, state$query_key, block_start + ) + if (is.null(cached)) { + block_display_idx <- seq.int(block_start, block_end) + cached <- dataview_arrow_query_fetch( + state, state$query_indices[block_display_idx] + ) + dataview_arrow_query_cache_add( + reader_state, state$query_key, block_start, cached + ) + } + selected <- display_idx[display_idx >= block_start & display_idx <= block_end] + cached[selected - block_start + 1L, , drop = FALSE] + }) + if (length(pages) == 1L) return(pages[[1L]]) + dataview_arrow_bind_pages(pages) +} diff --git a/sess/R/dataview-dbi.R b/sess/R/dataview-dbi.R new file mode 100644 index 000000000..0994c8fdb --- /dev/null +++ b/sess/R/dataview-dbi.R @@ -0,0 +1,274 @@ +dataview_is_dbi_lazy <- function(data) { + inherits(data, "tbl_sql") +} + +dataview_dbi_cache_block_size <- 5000L +dataview_dbi_cache_rows <- 20000L + +dataview_dbi_source <- function(data, view_id = NULL) { + if (!requireNamespace("DBI", quietly = TRUE) || + !requireNamespace("dbplyr", quietly = TRUE)) { + stop("Viewing lazy database tables requires the optional 'DBI' and 'dbplyr' packages") + } + con <- dbplyr::remote_con(data) + if (!inherits(con, "DBIConnection") || !DBI::dbIsValid(con)) { + stop("the lazy table does not have a valid DBI connection") + } + # Reserve the connection before even querying metadata: some DBI drivers + # invalidate an existing result whenever another statement is executed. + for (id in names(.sess_env$dataviews)) { + previous <- .sess_env$dataviews[[id]] + if (!identical(previous$dbi$con, con)) next + if (!identical(id, view_id)) { + stop("Connection already has a Data Viewer. Use a separate connection for another viewer.") + } + # Reopening the same viewer replaces its execution without leaving a + # partially fetched cursor for the metadata query to invalidate. + dataview_dbi_result_reset(previous$dbi_cache, clear_cache = TRUE) + } + # Let dbplyr translate metadata SQL for the connection's dialect. Execute the + # original top-level query when fetching rows, preserving its own operations. + source_sql <- suppressWarnings(dbplyr::sql_render(data, subquery = TRUE)) + schema <- DBI::dbGetQuery(con, as.character(dbplyr::sql_query_fields(con, source_sql))) + list(con = con, query_sql = as.character(dbplyr::sql_render(data)), schema = schema) +} + +dataview_dbi_to_state <- function(data, view_id = NULL) { + source <- dataview_dbi_source(data, view_id) + colnames <- names(source$schema) + headers <- c(" ", trimws(colnames)) + fields <- as.character(seq_along(headers) - 1L) + cols <- c(list(integer()), as.list(source$schema)) + columns <- .mapply(get_column_def, list(headers, fields, cols), NULL) + columns <- lapply(columns, function(column) { + column$filter <- jsonlite::unbox(FALSE) + column$sortable <- jsonlite::unbox(FALSE) + column + }) + list( + dbi = source, + columns = columns, + column_names = colnames, + total_rows = NULL, + uses_result_cache = TRUE, + dbi_cache = dataview_dbi_cache_state() + ) +} + +dataview_dbi_cache_state <- function() { + state <- new.env(parent = emptyenv()) + state$result <- NULL + state$result_cache_dir <- NULL + state$result_cache_blocks <- list() + state$result_cache_rows <- 0L + state$result_complete <- FALSE + state +} + +dataview_dbi_result_reset <- function(cache_state, clear_cache = FALSE) { + if (!is.null(cache_state$result) && DBI::dbIsValid(cache_state$result)) { + try(DBI::dbClearResult(cache_state$result), silent = TRUE) + } + cache_state$result <- NULL + if (clear_cache) { + # A new execution must never append to rows from the failed cursor. + if (!is.null(cache_state$result_cache_dir)) { + # Include files finalized just before an interrupted cache publication. + unlink(list.files(cache_state$result_cache_dir, full.names = TRUE), force = TRUE) + } + cache_state$result_cache_blocks <- list() + cache_state$result_cache_rows <- 0L + cache_state$result_complete <- FALSE + } +} + +dataview_dbi_result_cache_init <- function(state, view_id) { + if (!isTRUE(state$uses_result_cache)) return(invisible(NULL)) + if (!requireNamespace("arrow", quietly = TRUE)) { + stop("Viewing lazy database tables requires the optional 'arrow' package") + } + + root <- file.path(tempdir(), "vscode-r-dataview") + dir.create(root, recursive = TRUE, showWarnings = FALSE) + result_cache_dir <- tempfile( + pattern = paste0(view_id, "-dbi-"), tmpdir = root + ) + if (!dir.create(result_cache_dir, showWarnings = FALSE)) { + stop("unable to create temporary database data viewer cache") + } + + state$dbi_cache$result_cache_dir <- result_cache_dir + state$dbi_cache$result_cache_blocks <- list() + state$dbi_cache$result_cache_rows <- 0L + state$dbi_cache$result_complete <- FALSE + invisible(NULL) +} + +dataview_dbi_result_cache_cleanup <- function(state) { + if (is.null(state$dbi_cache)) return(invisible(NULL)) + + dataview_dbi_result_reset(state$dbi_cache) + result_cache_dir <- state$dbi_cache$result_cache_dir + if (!is.null(result_cache_dir) && dir.exists(result_cache_dir)) { + unlink(result_cache_dir, recursive = TRUE, force = TRUE) + } + state$dbi_cache$result_cache_dir <- NULL + state$dbi_cache$result_cache_blocks <- list() + state$dbi_cache$result_cache_rows <- 0L + state$dbi_cache$result_complete <- FALSE + invisible(NULL) +} + +dataview_dbi_result_cache_add <- function(cache_state, block) { + if (!nrow(block)) return(invisible(NULL)) + + block_number <- length(cache_state$result_cache_blocks) + 1L + first_row <- cache_state$result_cache_rows + 1L + last_row <- first_row + nrow(block) - 1L + cache_state$result_cache_blocks[[block_number]] <- list( + data = block, + path = NULL, + first_row = first_row, + last_row = last_row + ) + cache_state$result_cache_rows <- last_row + + in_memory_rows <- function() { + sum(vapply( + cache_state$result_cache_blocks, + function(cached) { + if (is.null(cached$data)) 0L else cached$last_row - cached$first_row + 1L + }, + integer(1) + )) + } + + while (length(cache_state$result_cache_blocks) > 1L && + in_memory_rows() > dataview_dbi_cache_rows) { + position <- which(vapply( + cache_state$result_cache_blocks, + function(cached) !is.null(cached$data), + logical(1) + ))[[1L]] + cached <- cache_state$result_cache_blocks[[position]] + path <- file.path( + cache_state$result_cache_dir, + sprintf("block-%08d.arrow", position) + ) + temporary_path <- paste0(path, ".tmp") + on.exit(unlink(temporary_path, force = TRUE), add = TRUE) + getExportedValue("arrow", "write_ipc_file")( + getExportedValue("arrow", "Table")$create(cached$data), + temporary_path + ) + if (!file.rename(temporary_path, path)) { + stop("unable to finalize temporary database data viewer cache") + } + cached$data <- NULL + cached$path <- path + cache_state$result_cache_blocks[[position]] <- cached + } + invisible(NULL) +} + +dataview_dbi_result_cache_slice <- function(state, row_idx, positions) { + cache_state <- state$dbi_cache + target <- if (length(row_idx)) max(row_idx) else 0L + + if (target > cache_state$result_cache_rows && + !isTRUE(cache_state$result_complete)) { + if (!is.null(cache_state$result) && !DBI::dbIsValid(cache_state$result)) { + stop(paste( + "The Data Viewer result is no longer valid; its connection was reused or disconnected.", + "Reopen the viewer with a dedicated connection." + )) + } + tryCatch({ + if (is.null(cache_state$result)) { + cache_state$result <- DBI::dbSendQuery( + state$dbi$con, state$dbi$query_sql + ) + } + while (target > cache_state$result_cache_rows && + !isTRUE(cache_state$result_complete)) { + block <- DBI::dbFetch(cache_state$result, n = dataview_dbi_cache_block_size) + dataview_dbi_result_cache_add(cache_state, block) + if (!nrow(block) || DBI::dbHasCompleted(cache_state$result)) { + cache_state$result_complete <- TRUE + dataview_dbi_result_reset(cache_state) + } + } + }, error = function(e) { + suspendInterrupts(dataview_dbi_result_reset(cache_state, clear_cache = TRUE)) + stop(e) + }, interrupt = function(e) { + suspendInterrupts(dataview_dbi_result_reset(cache_state, clear_cache = TRUE)) + stop(e) + }) + } + + available <- row_idx[row_idx <= cache_state$result_cache_rows] + if (!length(available)) { + return(state$dbi$schema[integer(), positions, drop = FALSE]) + } + + pages <- list() + page_rows <- integer() + for (cached in cache_state$result_cache_blocks) { + selected <- available[ + available >= cached$first_row & available <= cached$last_row + ] + if (!length(selected)) next + + rows <- selected - cached$first_row + 1L + page <- if (!length(positions)) { + data.frame(row.names = seq_along(selected)) + } else if (!is.null(cached$data)) { + cached$data[rows, positions, drop = FALSE] + } else { + data <- getExportedValue("arrow", "read_ipc_file")( + cached$path, + col_select = state$column_names[positions], + as_data_frame = FALSE + ) + dataview_arrow_data_frame(data$Take(rows - 1L)) + } + pages[[length(pages) + 1L]] <- page + page_rows <- c(page_rows, selected) + } + + page <- if (length(pages) == 1L) pages[[1L]] else do.call(rbind, pages) + page[match(available, page_rows), , drop = FALSE] +} + +dataview_dbi_page <- function( + state, start_row, end_row, sort_model, filter_model, fields = NULL +) { + positions <- dataview_page_positions(state, fields) + # Ignore stale viewer controls: database results are fetched in source order. + total <- if (isTRUE(state$dbi_cache$result_complete)) state$dbi_cache$result_cache_rows else NULL + row_idx <- if (is.null(total)) { + if (start_row < end_row) seq.int(start_row + 1L, end_row) else integer() + } else { + dataview_page_indices(start_row, end_row, total) + } + page <- dataview_dbi_result_cache_slice(state, row_idx, positions) + + for (position in seq_len(ncol(page))) { + if (inherits(page[[position]], "POSIXt")) { + page[[position]] <- sub( + "\\.?0+$", "", format(page[[position]], "%Y-%m-%dT%H:%M:%OS6") + ) + } else if (inherits(page[[position]], "Date")) { + page[[position]] <- format(page[[position]], "%Y-%m-%d") + } else if (inherits(page[[position]], "integer64")) { + page[[position]] <- as.character(page[[position]]) + } else if (state$columns[[positions[[position]] + 1L]]$type == "textColumn") { + page[[position]] <- dataview_format_column(page[[position]]) + } + } + + page_row_idx <- row_idx[seq_len(nrow(page))] + rows <- dataview_bind_rows(page, page_row_idx, as.character(positions)) + c(list(rows = rows), dataview_lazy_totals(state)) +} diff --git a/sess/R/dataview-lazy.R b/sess/R/dataview-lazy.R new file mode 100644 index 000000000..fbeebd843 --- /dev/null +++ b/sess/R/dataview-lazy.R @@ -0,0 +1,89 @@ +# Viewer operations over one cached execution of a transformed lazy query. + +dataview_query_positions <- function(state, sort_model, filter_model) { + positions <- integer() + for (field in names(filter_model)) { + pos <- dataview_field_position(field, length(state$columns)) + if (!is.na(pos) && pos > 1L && !isFALSE(state$columns[[pos]]$filter)) { + positions <- c(positions, pos) + } + } + for (item in sort_model) { + pos <- dataview_field_position(as.character(item$colId %||% ""), length(state$columns)) + if (!is.na(pos) && pos > 1L && !isFALSE(state$columns[[pos]]$sortable)) { + positions <- c(positions, pos) + } + } + unique(positions) +} + +dataview_lazy_cache <- function(state) { + if (!is.null(state$dbi)) state$dbi_cache else state$arrow_reader +} + +dataview_lazy_cache_slice <- function(state, rows, positions) { + if (!is.null(state$dbi)) { + dataview_dbi_result_cache_slice(state, rows, positions) + } else { + dataview_arrow_result_cache_slice(state, rows, positions) + } +} + +dataview_lazy_page_indices <- function(state, start_row, end_row, sort_model, filter_model) { + cache <- dataview_lazy_cache(state) + key <- dataview_query_key(sort_model, filter_model) + if (!identical(key, cache$result_query_key)) { + # Keep key columns in source order across in-memory and IPC projections. + positions <- sort(dataview_query_positions(state, sort_model, filter_model)) + indices <- NULL + if (length(positions)) { + # Finish the existing execution, spilling batches as usual. An infinite + # target drains the reader without assembling the full result in R. + dataview_lazy_cache_slice(state, Inf, integer()) + if (cache$result_cache_rows > .Machine$integer.max) { + stop("query result is too large for the current data viewer row index") + } + indices <- integer() + if (cache$result_cache_rows > 0L) { + query_state <- state + query_state$total_rows <- cache$result_cache_rows + columns <- dataview_lazy_cache_slice( + state, seq_len(cache$result_cache_rows), positions - 1L + ) + query_state$query_columns <- setNames(as.list(columns), as.character(positions)) + indices <- dataview_apply_filter_model(query_state, filter_model, NULL) + indices <- dataview_apply_sort_model(query_state, sort_model, indices) + } + } + # Publish only after success. A cancelled/failed operation can be retried, + # and changing controls always starts from the original cached result. + cache$result_query_indices <- indices + cache$result_query_key <- key + } + + total <- dataview_lazy_totals(state)$totalRows + display_idx <- if (is.null(total)) { + if (start_row < end_row) seq.int(start_row + 1L, end_row) else integer() + } else { + dataview_page_indices(start_row, end_row, total) + } + if (is.null(cache$result_query_indices)) display_idx else cache$result_query_indices[display_idx] +} + +dataview_lazy_totals <- function(state) { + cache <- dataview_lazy_cache(state) + original_total <- if (isTRUE(cache$result_complete)) cache$result_cache_rows else state$total_rows + total <- if (is.null(cache$result_query_indices)) { + original_total + } else { + length(cache$result_query_indices) + } + c( + list(lastRow = total %||% -1L), + if (!is.null(total)) { + list(totalRows = total, totalUnfiltered = original_total) + } else { + list(fetchedRows = cache$result_cache_rows) + } + ) +} diff --git a/sess/R/handlers.R b/sess/R/handlers.R index ad0f10f50..99c7815c6 100644 --- a/sess/R/handlers.R +++ b/sess/R/handlers.R @@ -39,6 +39,7 @@ workspace_child_count <- function(object) { get_workspace_data <- function() { env <- .GlobalEnv all_names <- workspace_env_names(env) + workspace_dimensions <- new.env(parent = emptyenv()) objs <- lapply(all_names, function(name) { if (bindingIsActive(name, env)) { @@ -56,6 +57,13 @@ get_workspace_data <- function() { obj_type <- typeof(obj) obj_length <- length(obj) obj_dim <- dim(obj) + if (inherits(obj, "arrow_dplyr_query") && !is.null(obj_dim)) { + assign( + name, + list(object = obj, dim = obj_dim), + envir = workspace_dimensions + ) + } first_class <- if (length(obj_class)) obj_class[[1]] else obj_type info <- list( class = class(obj), @@ -94,6 +102,7 @@ get_workspace_data <- function() { info }) names(objs) <- all_names + .sess_env$workspace_dimensions <- workspace_dimensions list( globalenv = objs, @@ -696,12 +705,25 @@ dataview_is_table <- function(data) { is.data.frame(data) || is.matrix(data) || inherits(data, "ArrowTabular") || + dataview_is_arrow_lazy(data) || + dataview_is_dbi_lazy(data) || inherits(data, "polars_data_frame") } dataview_schema <- function(data) { + if (inherits(data, "ArrowTabular") || dataview_is_arrow_lazy(data)) { + old_options <- options(arrow.int64_downcast = FALSE) + on.exit(options(old_options), add = TRUE) + } + if (inherits(data, "ArrowTabular")) { - return(data$Slice(0L, 0L)$to_data_frame()) + return(dataview_arrow_data_frame(data$Slice(0L, 0L))) + } + if (inherits(data, "arrow_dplyr_query")) { + return(dataview_arrow_data_frame(utils::head(data, 0L))) + } + if (dataview_is_arrow_lazy(data)) { + return(dataview_arrow_data_frame(data[integer(0), ])) } if (inherits(data, "polars_data_frame")) { return(as.data.frame(data$slice(0L, 0L))) @@ -709,12 +731,27 @@ dataview_schema <- function(data) { data[0, , drop = FALSE] } -dataview_slice <- function(data, row_idx) { +dataview_slice <- function(data, row_idx, conversion = NULL) { if (inherits(data, "ArrowTabular")) { if (!length(row_idx)) { - return(data$Slice(0L, 0L)$to_data_frame()) + return(dataview_arrow_data_frame(data$Slice(0L, 0L))) + } + return(dataview_arrow_data_frame(data[row_idx, ], conversion)) + } + if (dataview_is_arrow_lazy(data)) { + if (!length(row_idx)) { + return(dataview_schema(data)) + } + if (inherits(data, "arrow_dplyr_query")) { + reader_state <- dataview_arrow_reader_state() + reader_state$conversion <- conversion + reader_state$reader <- dataview_arrow_reader_open(data) + on.exit(try(reader_state$reader$Close(), silent = TRUE), add = TRUE) + selected <- sort(unique(row_idx)) + page <- dataview_arrow_reader_select(reader_state, selected) + return(page[match(row_idx, selected), , drop = FALSE]) } - return(data[row_idx, ]$to_data_frame()) + return(dataview_arrow_data_frame(data[row_idx, , drop = FALSE], conversion)) } if (inherits(data, "polars_data_frame")) { if (!length(row_idx)) { @@ -739,6 +776,13 @@ dataview_column <- function(data, position) { if (inherits(data, "ArrowTabular")) { return(as.vector(data[[position]])) } + if (inherits(data, "Dataset") && !is.data.frame(data)) { + return(dataview_arrow_column(data, position)) + } + if (dataview_is_arrow_lazy(data)) { + name <- names(data)[[position]] + return(as.data.frame(data[, name, drop = FALSE])[[name]]) + } if (inherits(data, "polars_data_frame")) { return(as.data.frame(data[, position])[[1L]]) } @@ -749,12 +793,17 @@ dataview_column <- function(data, position) { } dataview_to_state <- function(data) { + if (dataview_is_dbi_lazy(data)) return(dataview_dbi_to_state(data)) if (!dataview_is_table(data)) { - stop("data must be a data frame, matrix, Arrow table, or Polars data frame") + stop("data must be a data frame, matrix, Arrow table/dataset/query, or Polars data frame") } - n <- nrow(data) + n <- if (inherits(data, "arrow_dplyr_query")) NULL else nrow(data) + if (!is.null(n) && (length(n) != 1L || is.na(n))) { + stop("unable to determine the number of rows") + } colnames <- if (inherits(data, "ArrowTabular") || + dataview_is_arrow_lazy(data) || inherits(data, "polars_data_frame")) { names(data) } else { @@ -783,18 +832,29 @@ dataview_to_state <- function(data) { dataview_column(schema, position) }) ) + columns <- .mapply( + get_column_def, + list(headers, fields, cols), + NULL + ) + if (inherits(data, "ArrowTabular") || dataview_is_arrow_lazy(data)) { + nested <- dataview_arrow_nested_columns(data) + for (position in which(nested)) { + columns[[position + 1L]]$filter <- jsonlite::unbox(FALSE) + columns[[position + 1L]]$sortable <- jsonlite::unbox(FALSE) + } + } list( data = data, row_index = row_index, - columns = .mapply( - get_column_def, - list(headers, fields, cols), - NULL - ), + columns = columns, total_rows = n, query_key = NULL, - query_indices = NULL + query_indices = NULL, + query_has_sort = FALSE, + uses_result_cache = inherits(data, "arrow_dplyr_query"), + arrow_reader = if (dataview_is_arrow_lazy(data)) dataview_arrow_reader_state(data) else NULL ) } @@ -822,20 +882,41 @@ dataview_set_state <- function(view_id, state) { state } -dataview_register <- function(data, view_id = NULL, live = FALSE) { +dataview_register <- function(data, view_id = NULL, live = FALSE, total_rows = NULL) { if (is.null(view_id)) { view_id <- dataview_new_id() } - state <- dataview_to_state(data) + state <- if (dataview_is_dbi_lazy(data)) { + dataview_dbi_to_state(data, view_id) + } else { + dataview_to_state(data) + } + if (inherits(data, "arrow_dplyr_query")) { + if (is.numeric(total_rows) && length(total_rows) == 1L && + is.finite(total_rows) && total_rows >= 0 && total_rows == floor(total_rows)) { + state$total_rows <- total_rows + } + dataview_arrow_result_cache_init(state, view_id) + } + if (!is.null(state$dbi) && isTRUE(state$uses_result_cache)) { + dataview_dbi_result_cache_init(state, view_id) + } + + previous_state <- .sess_env$dataviews[[view_id]] + if (!is.null(previous_state$arrow_reader)) { + dataview_arrow_result_cache_cleanup(previous_state) + } + if (!is.null(previous_state$dbi_cache)) { + dataview_dbi_result_cache_cleanup(previous_state) + } state$live <- live state$revision <- .sess_env$dataview_revision %||% 0 state <- dataview_set_state(view_id, state) - list( - view_id = view_id, - state_generation = state$state_generation, - total_rows = state$total_rows, - columns = dataview_columns(state) + c( + list(view_id = view_id, state_generation = state$state_generation), + if (!is.null(state$total_rows)) list(total_rows = state$total_rows), + list(columns = dataview_columns(state)) ) } @@ -874,7 +955,8 @@ dataview_column_values <- function(state, position, row_idx = NULL) { if (!is.null(row_idx) && is.matrix(state$data)) { return(state$data[row_idx, position - 1L]) } - values <- dataview_column(state$data, position - 1L) + values <- state$query_columns[[as.character(position)]] + if (is.null(values)) values <- dataview_column(state$data, position - 1L) if (is.null(row_idx)) values else values[row_idx] } @@ -931,6 +1013,14 @@ dataview_format_column <- function(values, truncate = TRUE) { formatted[is.na(values)] <- NA_character_ return(formatted) } + if (is.data.frame(values)) { + return(lapply(seq_len(nrow(values)), function(row) { + lapply(as.list(values[row, , drop = FALSE]), function(value) { + value <- dataview_format_column(value) + if (length(value) == 1L) value[[1L]] else value + }) + })) + } if (is.list(values) && is.null(dim(values)) && !inherits(values, "POSIXlt")) { return(lapply(values, dataview_cell_preview)) } @@ -980,7 +1070,11 @@ dataview_match_condition <- function(values, cond) { if (inherits(values, "Date")) { ds <- as.Date(values) d1 <- as.Date(if (is.null(cond$dateFrom)) cond$filter else cond$dateFrom) - d2 <- as.Date(if (is.null(cond$dateTo)) cond$filterTo else cond$dateTo) + d2 <- if (cond_type == "inRange") { + as.Date(if (is.null(cond$dateTo)) cond$filterTo else cond$dateTo) + } else { + NULL + } } else { ds <- as.POSIXct(values) tz <- attr(ds, "tzone", exact = TRUE) %||% "" @@ -988,10 +1082,14 @@ dataview_match_condition <- function(values, cond) { if (is.null(cond$dateFrom)) cond$filter else cond$dateFrom, tz = tz ) - d2 <- as.POSIXct( - if (is.null(cond$dateTo)) cond$filterTo else cond$dateTo, - tz = tz - ) + d2 <- if (cond_type == "inRange") { + as.POSIXct( + if (is.null(cond$dateTo)) cond$filterTo else cond$dateTo, + tz = tz + ) + } else { + NULL + } } result <- switch(cond_type, equals = ds == d1, @@ -1072,6 +1170,37 @@ dataview_field_position <- function(field, column_count) { position } +dataview_page_positions <- function(state, fields = NULL) { + positions <- seq_len(length(state$columns) - 1L) + if (!is.null(fields)) { + fields <- unlist(fields, use.names = FALSE) + positions <- positions[as.character(positions) %in% fields] + } + positions +} + +dataview_page_indices <- function(start_row, end_row, total) { + end_row <- min(total, end_row) + if (start_row >= total || start_row >= end_row) { + return(integer()) + } + seq.int(start_row + 1L, end_row) +} + +dataview_block_starts <- function(row_idx, block_size) { + if (!length(row_idx)) { + return(integer()) + } + unique((row_idx - 1L) %/% block_size * block_size + 1L) +} + +dataview_bind_rows <- function(page, row_idx, fields) { + rows <- cbind(data.frame(row_idx, check.names = FALSE), page) + names(rows) <- c("0", fields) + rownames(rows) <- NULL + rows +} + dataview_apply_filter_model <- function(state, filter_model, row_idx) { if (is.null(filter_model) || !length(filter_model)) { return(row_idx) @@ -1178,17 +1307,43 @@ dataview_query_indices <- function(state, sort_model, filter_model) { return(NULL) } + if (dataview_is_arrow_lazy(state$data)) { + # Request-local only: share one projected scan across filtering and sorting. + state$query_columns <- dataview_arrow_query_columns(state, sort_model, filter_model) + } # NULL means the original order. Avoid materializing an identity index and # subsetting every full column just to reproduce its existing values. row_idx <- dataview_apply_filter_model(state, filter_model, NULL) dataview_apply_sort_model(state, sort_model, row_idx) } -dataview_rows <- function(state, row_idx) { - page <- dataview_slice(state$data, row_idx) +dataview_rows <- function( + state, + row_idx, + use_arrow_reader = FALSE, + use_arrow_query_cache = FALSE, + display_idx = row_idx +) { + page <- if (identical(state$column_fields, character())) { + if (use_arrow_reader && inherits(state$data, "arrow_dplyr_query")) { + data <- dataview_arrow_slice(state, row_idx) + data.frame(row.names = seq_len(nrow(data))) + } else { + data.frame(row.names = seq_along(row_idx)) + } + } else if (use_arrow_reader) { + dataview_arrow_slice(state, row_idx) + } else if (use_arrow_query_cache) { + dataview_arrow_query_slice(state, display_idx) + } else { + dataview_slice(state$data, row_idx) + } if (is.matrix(page)) { page <- as.data.frame(page, optional = TRUE) } + if (is.data.frame(page) && nrow(page) < length(row_idx)) { + row_idx <- row_idx[seq_len(nrow(page))] + } if (is.data.frame(page)) { for (position in seq_len(ncol(page))) { if (inherits(page[[position]], "POSIXct") || @@ -1201,25 +1356,23 @@ dataview_rows <- function(state, row_idx) { } } } - rows <- cbind( - data.frame( - dataview_column_values(state, 1L, row_idx), - check.names = FALSE - ), - page + dataview_bind_rows( + page, + dataview_column_values(state, 1L, row_idx), + state$column_fields %||% as.character(seq_len(ncol(page))) ) - names(rows) <- as.character(seq_len(ncol(rows)) - 1L) - rownames(rows) <- NULL - rows } handle_dataview_init <- function(params) { view_id <- as.character(params$view_id %||% "") state <- dataview_get_state(view_id, refresh = TRUE) - list( - columns = dataview_columns(state), - totalRows = state$total_rows, - live = isTRUE(state$live) + c( + list(columns = dataview_columns(state)), + if (!is.null(state$total_rows)) list(totalRows = state$total_rows), + list( + live = isTRUE(state$live), + columnProjection = !is.null(state$dbi) || dataview_is_arrow_lazy(state$data) + ) ) } @@ -1228,7 +1381,8 @@ handle_dataview_page <- function(params) { state <- dataview_get_state(view_id) start_row <- suppressWarnings(as.integer(params$startRow %||% 0L)) - end_row <- suppressWarnings(as.integer(params$endRow %||% min(state$total_rows, 500L))) + default_end_row <- if (is.null(state$total_rows)) 500L else min(state$total_rows, 500L) + end_row <- suppressWarnings(as.integer(params$endRow %||% default_end_row)) if (is.na(start_row) || start_row < 0L) { start_row <- 0L } @@ -1238,10 +1392,30 @@ handle_dataview_page <- function(params) { sort_model <- params$sortModel %||% list() filter_model <- params$filterModel %||% list() + if (!is.null(state$dbi)) { + return(dataview_dbi_page(state, start_row, end_row, sort_model, filter_model, params$fields)) + } + if (inherits(state$data, "arrow_dplyr_query")) { + page_idx <- dataview_lazy_page_indices(state, start_row, end_row, sort_model, filter_model) + page_state <- dataview_arrow_page_state(state, params$fields) + rows <- dataview_rows( + page_state, + page_idx, + use_arrow_reader = TRUE + ) + return(c(list(rows = rows), dataview_lazy_totals(state))) + } + query_key <- dataview_query_key(sort_model, filter_model) if (!identical(query_key, state$query_key)) { state$query_key <- query_key state$query_indices <- dataview_query_indices(state, sort_model, filter_model) + state$query_has_sort <- !is.null(sort_model) && length(sort_model) > 0L + if (dataview_is_arrow_lazy(state$data)) { + dataview_arrow_reader_reset(state) + state$arrow_reader$row_cache <- list() + state$arrow_reader$query_cache <- list() + } .sess_env$dataviews[[view_id]] <- state } @@ -1250,21 +1424,28 @@ handle_dataview_page <- function(params) { } else { length(state$query_indices) } - end_exclusive <- min(total, end_row) - if (start_row >= total || start_row >= end_exclusive) { - display_idx <- integer(0) - page_idx <- integer(0) + display_idx <- dataview_page_indices(start_row, end_row, total) + page_idx <- if (is.null(state$query_indices)) { + display_idx } else { - display_idx <- seq.int(start_row + 1L, end_exclusive) - page_idx <- if (is.null(state$query_indices)) { - display_idx - } else { - state$query_indices[display_idx] - } + state$query_indices[display_idx] } + page_state <- if (dataview_is_arrow_lazy(state$data)) { + dataview_arrow_page_state(state, params$fields) + } else { + state + } result <- list( - rows = dataview_rows(state, page_idx), + rows = dataview_rows( + page_state, + page_idx, + use_arrow_reader = dataview_is_arrow_lazy(state$data) && + is.null(state$query_indices), + use_arrow_query_cache = dataview_is_arrow_lazy(state$data) && + !is.null(state$query_indices), + display_idx = display_idx + ), totalRows = total, totalUnfiltered = state$total_rows, lastRow = total @@ -1272,9 +1453,10 @@ handle_dataview_page <- function(params) { if (isTRUE(params$formatNumbers)) { # Format only the requested page, using the session's digits, scipen and OutDec. # Keep raw rows unchanged for sorting, filtering and full-precision inspection. - numeric_fields <- vapply(state$columns, function(column) { + numeric_fields <- vapply(Filter(function(column) { column$type == "numericColumn" && column$field != "0" - }, logical(1)) + }, state$columns), function(column) as.character(column$field), character(1)) + numeric_fields <- intersect(numeric_fields, names(result$rows)) result$formattedColumns <- lapply(result$rows[numeric_fields], function(values) { I(format(values, trim = TRUE, justify = "none")) }) @@ -1287,6 +1469,12 @@ handle_dataview_dispose <- function(params) { state <- if (is.null(.sess_env$dataviews)) NULL else .sess_env$dataviews[[view_id]] state_generation <- suppressWarnings(as.integer(params$state_generation %||% NA_integer_)) if (!is.null(state) && identical(state$state_generation, state_generation)) { + if (!is.null(state$arrow_reader)) { + dataview_arrow_result_cache_cleanup(state) + } + if (!is.null(state$dbi_cache)) { + dataview_dbi_result_cache_cleanup(state) + } .sess_env$dataviews[[view_id]] <- NULL } TRUE diff --git a/sess/R/hooks.R b/sess/R/hooks.R index 0d207ea44..f7def34bd 100644 --- a/sess/R/hooks.R +++ b/sess/R/hooks.R @@ -103,7 +103,23 @@ runtime_start <- function(use_rstudioapi = TRUE, } if (view_type == "table") { - registration <- dataview_register(x, view_id = view_id) + total_rows <- NULL + workspace_dimensions <- .sess_env$workspace_dimensions + if (inherits(x, "arrow_dplyr_query") && + !is.null(workspace_dimensions) && + exists(title_key, envir = workspace_dimensions, inherits = FALSE)) { + workspace_dimension <- get( + title_key, envir = workspace_dimensions, inherits = FALSE + ) + if (identical(workspace_dimension$object, x) && + length(workspace_dimension$dim)) { + total_rows <- workspace_dimension$dim[[1L]] + } + } + + registration <- dataview_register( + x, view_id = view_id, total_rows = total_rows + ) notify_client("dataview", list( title = title, diff --git a/sess/R/interactive.R b/sess/R/interactive.R index 03e935025..5d4b63895 100644 --- a/sess/R/interactive.R +++ b/sess/R/interactive.R @@ -370,6 +370,9 @@ interactive_execute <- function(id, code, source = NULL) { } .interactive_rich_value <- function(value) { + # Lazy queries can have unknown totals and require their own paging/cache + # lifecycle. Let View() open the paged viewer instead of a live R snapshot. + if (inherits(value, "arrow_dplyr_query") || dataview_is_dbi_lazy(value)) return(FALSE) if (dataview_is_table(value)) { snapshot <- .interactive_table_snapshot(value) registration <- dataview_register(snapshot$value) diff --git a/sess/R/runtime.R b/sess/R/runtime.R index c62eb22fd..6b73d7471 100644 --- a/sess/R/runtime.R +++ b/sess/R/runtime.R @@ -17,7 +17,11 @@ .runtime_clear_viewer_state <- function() { # Open dataviews and title-to-id mappings belong to one runtime connection. - # Keep an empty registry available, but never restore data from a prior run. + # Release readers and disk caches before dropping their owning states. + for (view in .sess_env$dataviews) { + dataview_arrow_result_cache_cleanup(view) + dataview_dbi_result_cache_cleanup(view) + } .sess_env$dataviews <- list() .sess_env$dataview_registry <- new.env(parent = emptyenv()) invisible(NULL) diff --git a/sess/R/server.R b/sess/R/server.R index 0fb37b410..7f3705edd 100644 --- a/sess/R/server.R +++ b/sess/R/server.R @@ -67,7 +67,9 @@ connect <- function(endpoint = NULL, use_rstudioapi = TRUE, # Only follow discovery if it describes this connection, including explicit # endpoints from the generated attach script. Never redirect an unrelated client. discovery <- if (nzchar(discovery_file)) .read_discovery(discovery_file) else NULL + .sess_env$terminal_pid <- NULL if (!is.null(discovery) && identical(discovery$endpoint, endpoint)) { + .sess_env$terminal_pid <- discovery$terminalPid .configure_discovery_jgd(discovery, plot_backend %in% c("auto", "jgd")) .sess_env$reconnect <- list( path = discovery_file, endpoint = endpoint, @@ -277,6 +279,7 @@ connect <- function(endpoint = NULL, use_rstudioapi = TRUE, host = unname(host), version = sprintf("%s.%s", R.version$major, R.version$minor), pid = Sys.getpid(), + terminal_pid = .sess_env$terminal_pid, tempdir = .sess_env$tempdir, wd = getwd(), info = list( @@ -387,6 +390,11 @@ poll_connection <- function(generation = .sess_env$transport_generation) { !identical(.sess_env$interactive_pid, Sys.getpid())) return() con <- .sess_env$con if (is.null(con) || !identical(generation, .sess_env$transport_generation)) return() + on.exit(suspendInterrupts({ + if (!is.null(.sess_env$con) && identical(generation, .sess_env$transport_generation)) { + later::later(function() poll_connection(generation), 0.01) + } + }), add = TRUE) # Non-blocking poll: 0 ms timeout ready <- tryCatch( @@ -402,7 +410,7 @@ poll_connection <- function(generation = .sess_env$transport_generation) { return() } - # A NULL poll result is transient; keep the loop alive and reschedule below. + # A NULL poll result is transient; keep the loop alive. if (!is.null(ready) && length(ready) > 0 && identical(ready[[1]], "ready")) { chunk <- tryCatch( processx::conn_read_chars(con), @@ -417,7 +425,6 @@ poll_connection <- function(generation = .sess_env$transport_generation) { if (!has_data) { if (.transport_empty_read_is_eof(con)) { .transport_disconnect(silent = TRUE, reconnect = TRUE) - return() } } else { .sess_env$read_buffer <- paste0(.sess_env$read_buffer, paste0(chunk, collapse = "")) @@ -445,11 +452,6 @@ poll_connection <- function(generation = .sess_env$transport_generation) { } } else if (length(ready) > 0 && ready[[1]] %in% c("closed", "error")) { .transport_disconnect(silent = TRUE, reconnect = TRUE) - return() - } - - if (!is.null(.sess_env$con) && identical(generation, .sess_env$transport_generation)) { - later::later(function() poll_connection(generation), 0.01) } } @@ -494,7 +496,15 @@ dispatch_message <- function(line) { if (payload$method %in% names(handlers)) { tryCatch( { - res <- handlers[[payload$method]](payload$params) + if (identical(payload$method, "dataview_page") && isTRUE(payload$params$notify_busy)) { + on.exit( + suspendInterrupts(notify_client("dataview_busy", list(view_id = NULL))), add = TRUE + ) + notify_client("dataview_busy", list(view_id = payload$params$view_id)) + } + # Event-loop callbacks may inherit suspended interrupts from R's + # input processing. Explicitly allow cancellation during the work. + res <- allowInterrupts(handlers[[payload$method]](payload$params)) rpc_reply(payload$id, result = res) }, error = function(e) { @@ -502,9 +512,11 @@ dispatch_message <- function(line) { warning(sprintf("[sess] Error in handler for '%s': %s", payload$method, e$message)) }, interrupt = function(e) { - # Return a terminal reply without unwinding poll_connection(), which - # must reschedule itself and dispatch the remaining buffered requests. - rpc_reply(payload$id, error = list(code = -32000L, message = "R request interrupted")) + # Keep the polling loop alive after a cancelled page. Otherwise the + # queued disposal and all subsequent session requests stop being read. + suspendInterrupts( + rpc_reply(payload$id, error = list(code = -32000L, message = "R request interrupted")) + ) } ) } else { diff --git a/sess/inst/tinytest/test-dataview-cleanup.R b/sess/inst/tinytest/test-dataview-cleanup.R new file mode 100644 index 000000000..4b3321309 --- /dev/null +++ b/sess/inst/tinytest/test-dataview-cleanup.R @@ -0,0 +1,97 @@ +# Retain the resource handles so garbage collection cannot mask missing cleanup. +if (requireNamespace("arrow", quietly = TRUE) && + requireNamespace("dbplyr", quietly = TRUE) && + requireNamespace("RSQLite", quietly = TRUE)) local({ + env <- sess:::.sess_env + old_views <- env$dataviews + old_registry <- env$dataview_registry + env$dataviews <- list() + env$dataview_registry <- new.env(parent = emptyenv()) + on.exit({ + sess:::runtime_stop() + env$dataviews <- old_views + env$dataview_registry <- old_registry + }, add = TRUE) + data <- data.frame(id = seq_len(30000L)) + + for (action in c("idle-stop", "active-stop", "disconnect", "restart")) local({ + if (action != "idle-stop") { + sess:::runtime_start(use_rstudioapi = FALSE, plot_backend = "native") + } + con <- DBI::dbConnect(RSQLite::SQLite(), ":memory:") + unopened_con <- DBI::dbConnect(RSQLite::SQLite(), ":memory:") + states <- list() + on.exit({ + for (state in states) { + sess:::dataview_dbi_result_cache_cleanup(state) + sess:::dataview_arrow_result_cache_cleanup(state) + } + DBI::dbDisconnect(con) + DBI::dbDisconnect(unopened_con) + sess:::runtime_stop() + }, add = TRUE) + DBI::dbWriteTable(con, "fixture", data) + table <- dplyr::tbl(con, "fixture") + # Include a query whose cache exists but whose cursor has not been opened. + DBI::dbWriteTable(unopened_con, "fixture", data) + unopened <- sess:::dataview_register(dplyr::tbl(unopened_con, "fixture")) + states[[1L]] <- sess:::dataview_get_state(unopened$view_id) + view <- sess:::dataview_register(table) + states[[2L]] <- sess:::dataview_get_state(view$view_id) + sess:::handle_dataview_page(list( + view_id = view$view_id, startRow = 20000L, endRow = 20002L + )) + db_cache <- states[[2L]]$dbi_cache + result <- db_cache$result + expect_true(DBI::dbIsValid(result)) + expect_false(DBI::dbHasCompleted(result)) + expect_true(length(list.files(db_cache$result_cache_dir)) > 0L) + + view <- sess:::dataview_register(dplyr::filter(arrow::Table$create(data), id > 0L)) + states[[3L]] <- sess:::dataview_get_state(view$view_id) + arrow_cache <- states[[3L]]$arrow_reader + # Fixed batches ensure the reader is still open after spilling to disk. + reader <- arrow::RecordBatchReader$create(batches = lapply(0:5, function(i) { + arrow::as_record_batch(data[(i * 5000L + 1L):((i + 1L) * 5000L), , drop = FALSE]) + })) + closed <- FALSE + arrow_cache$reader <- list( + read_next_batch = reader$read_next_batch, + Close = function() { + closed <<- TRUE + reader$Close() + } + ) + sess:::handle_dataview_page(list( + view_id = view$view_id, startRow = 20000L, endRow = 20002L + )) + expect_false(closed) + expect_true(length(list.files(arrow_cache$result_cache_dir)) > 0L) + paths <- c(states[[1L]]$dbi_cache$result_cache_dir, + db_cache$result_cache_dir, arrow_cache$result_cache_dir) + # Ordinary viewers have no database or Arrow resources. + sess:::dataview_register(data.frame(id = 1L)) + assign("table:fixture", view$view_id, envir = env$dataview_registry) + + switch(action, + "idle-stop" = sess:::runtime_stop(), + "active-stop" = sess:::runtime_stop(), + "disconnect" = sess:::.transport_disconnect(silent = TRUE), + "restart" = sess:::runtime_start(use_rstudioapi = FALSE, plot_backend = "native") + ) + + expect_false(DBI::dbIsValid(result), info = action) + expect_null(db_cache$result, info = action) + expect_true(closed, info = action) + expect_null(arrow_cache$reader, info = action) + expect_false(any(dir.exists(paths)), info = action) + expect_length(db_cache$result_cache_blocks, 0L) + expect_length(arrow_cache$result_cache_batches, 0L) + expect_length(env$dataviews, 0L) + expect_length(ls(env$dataview_registry, all.names = TRUE), 0L) + expect_true(DBI::dbIsValid(con)) + expect_equal(DBI::dbGetQuery(con, "SELECT COUNT(*) AS n FROM fixture")$n, nrow(data)) + # Repeated teardown must also be safe with an empty registry. + expect_silent(sess:::runtime_stop()) + }) +}) diff --git a/sess/inst/tinytest/test-dataview-counts.R b/sess/inst/tinytest/test-dataview-counts.R new file mode 100644 index 000000000..c7adbcfb3 --- /dev/null +++ b/sess/inst/tinytest/test-dataview-counts.R @@ -0,0 +1,76 @@ +if (requireNamespace("arrow", quietly = TRUE) && + requireNamespace("dbplyr", quietly = TRUE) && + requireNamespace("RSQLite", quietly = TRUE)) local({ + con <- DBI::dbConnect(RSQLite::SQLite(), ":memory:") + path <- tempfile(fileext = ".parquet") + views <- list() + on.exit({ + for (view in views) sess:::handle_dataview_dispose(view) + DBI::dbDisconnect(con) + unlink(path) + }, add = TRUE) + + for (n in c(0L, 10L, 500L, 12010L)) { + data <- data.frame(id = seq_len(n)) + DBI::dbWriteTable(con, "fixture", data, overwrite = TRUE) + arrow::write_parquet(data, path, chunk_size = 5000L) + table <- dplyr::tbl(con, "fixture") + sources <- list( + arrow = dplyr::filter(arrow::open_dataset(path), id > 0L), + database_table = table, + database_query = dplyr::filter(table, id > 0L) + ) + for (backend in names(sources)) { + view <- sess:::dataview_register(sources[[backend]]) + views[[length(views) + 1L]] <- view + init <- sess:::handle_dataview_init(list(view_id = view$view_id)) + expect_null(init$totalRows, info = backend) + first <- sess:::handle_dataview_page(list( + view_id = view$view_id, startRow = 0L, endRow = if (n > 500L) 5000L else 500L + )) + if (n <= 500L) { + expect_equal(first$totalRows, n, info = backend) + expect_equal(first$lastRow, n, info = backend) + expect_null(first$fetchedRows, info = backend) + } else { + expect_equal(first$lastRow, -1L, info = backend) + expect_null(first$totalRows, info = backend) + expect_null(first$totalUnfiltered, info = backend) + cache <- sess:::dataview_lazy_cache(sess:::dataview_get_state(view$view_id)) + expect_equal(first$fetchedRows, cache$result_cache_rows, info = backend) + expect_true(first$fetchedRows >= 5000L, info = backend) + expect_true(first$fetchedRows < n, info = backend) + expect_false(cache$result_complete, info = backend) + second <- sess:::handle_dataview_page(list( + view_id = view$view_id, startRow = 5000L, endRow = 5500L, fields = "0" + )) + expect_true(second$fetchedRows >= first$fetchedRows, info = backend) + expect_equal(second$fetchedRows, cache$result_cache_rows, info = backend) + backward <- sess:::handle_dataview_page(list( + view_id = view$view_id, startRow = 0L, endRow = 2L + )) + expect_equal(backward$fetchedRows, second$fetchedRows, info = backend) + last <- sess:::handle_dataview_page(list( + view_id = view$view_id, startRow = 12000L, endRow = 12500L + )) + expect_equal(last$totalRows, n, info = backend) + expect_equal(last$totalUnfiltered, n, info = backend) + expect_equal(last$lastRow, n, info = backend) + expect_null(last$fetchedRows, info = backend) + expect_equal(nrow(last$rows), 10L, info = backend) + } + sess:::handle_dataview_dispose(view) + } + # Existing eager table responses retain exact totals and no estimate. + view <- sess:::dataview_register(data) + views[[length(views) + 1L]] <- view + expect_equal(sess:::handle_dataview_init(list(view_id = view$view_id))$totalRows, n) + page <- sess:::handle_dataview_page(list( + view_id = view$view_id, startRow = 0L, endRow = 500L + )) + expect_equal(page$totalRows, n) + expect_equal(page$lastRow, n) + expect_null(page$fetchedRows) + sess:::handle_dataview_dispose(view) + } +}) diff --git a/sess/inst/tinytest/test-dataview-dbi-connections.R b/sess/inst/tinytest/test-dataview-dbi-connections.R new file mode 100644 index 000000000..aedd9919e --- /dev/null +++ b/sess/inst/tinytest/test-dataview-dbi-connections.R @@ -0,0 +1,141 @@ +# Each viewer reserves its connection, including before its first page is fetched. +if (requireNamespace("RSQLite", quietly = TRUE) && + requireNamespace("arrow", quietly = TRUE) && + requireNamespace("dbplyr", quietly = TRUE)) local({ + con <- DBI::dbConnect(RSQLite::SQLite(), ":memory:") + other <- DBI::dbConnect(RSQLite::SQLite(), ":memory:") + views <- list() + on.exit({ + for (view in views) sess:::handle_dataview_dispose(view) + DBI::dbDisconnect(con) + DBI::dbDisconnect(other) + }, add = TRUE) + fixture <- data.frame(id = seq_len(25010L)) + DBI::dbWriteTable(con, "items", fixture, temporary = TRUE) + DBI::dbWriteTable(other, "items", data.frame(id = fixture$id + 100000L), temporary = TRUE) + table <- dplyr::tbl(con, "items") + ascending <- dplyr::arrange(table, id) + descending <- dplyr::arrange(table, dplyr::desc(id)) + other_query <- dplyr::arrange(dplyr::tbl(other, "items"), dplyr::desc(id)) + page <- function(view, start = 0L, end = start + 2L) { + sess:::handle_dataview_page(list( + view_id = view$view_id, startRow = start, endRow = end + )) + } + + a <- sess:::dataview_register(ascending) + views <- c(views, list(a)) + expect_error(sess:::dataview_register(descending), "separate connection") + first <- page(a) + expect_identical(first$rows[["1"]], 1:2) + expect_equal(first$fetchedRows, 5000L) + expect_equal(first$lastRow, -1L) + state <- sess:::dataview_get_state(a$view_id) + result <- state$dbi_cache$result + expect_error(sess:::dataview_register(descending), "separate connection") + expect_true(DBI::dbIsValid(result)) + expect_identical(state$dbi_cache$result, result) + expect_identical(page(a)$rows, first$rows) + + # Identical table names and SQL must not share data across connections. + expect_identical(DBI::dbGetInfo(con), DBI::dbGetInfo(other)) + b <- sess:::dataview_register(other_query) + views <- c(views, list(b)) + first_other <- page(b) + expect_identical(first_other$rows[["1"]], 125010:125009) + expect_equal(first_other$fetchedRows, 5000L) + + # Interleave requests past the first batch, then revisit both cached prefixes. + expect_identical(page(a, 6000L)$rows[["1"]], 6001:6002) + expect_identical(page(b, 6000L)$rows[["1"]], 119010:119009) + expect_identical(page(a)$rows[["1"]], 1:2) + expect_identical(page(b)$rows[["1"]], 125010:125009) + expect_identical(page(a, 12000L)$rows[["1"]], 12001:12002) + expect_identical(page(b, 12000L)$rows[["1"]], 113010:113009) + + # Finishing and disposing one viewer must leave the other viewer and the + # user's connection usable, including data read back from spilled blocks. + last <- page(a, 25000L, 25020L) + expect_identical(last$rows[["1"]], 25001:25010) + expect_equal(last$totalRows, 25010L) + cache_dir <- state$dbi_cache$result_cache_dir + expect_true(length(list.files(cache_dir)) > 0L) + expect_identical(page(a)$rows[["1"]], 1:2) + sess:::handle_dataview_dispose(a) + expect_false(dir.exists(cache_dir)) + expect_silent(count <- DBI::dbGetQuery(con, "SELECT COUNT(*) AS n FROM items")) + expect_equal(count$n, nrow(fixture)) + expect_identical(page(b, 20000L)$rows[["1"]], 105010:105009) + last <- page(b, 25000L, 25020L) + expect_identical(last$rows[["1"]], 100010:100001) + expect_equal(last$totalRows, 25010L) + expect_true(DBI::dbIsValid(con)) + expect_equal(DBI::dbGetQuery(con, "SELECT COUNT(*) AS n FROM items")$n, 25010L) + + # Closing releases the reservation. Reopening that same viewer must close + # its old cursor before the new metadata query, and remove its spilled cache. + a <- sess:::dataview_register(ascending) + views <- c(views, list(a)) + page(a, 20000L) + state <- sess:::dataview_get_state(a$view_id) + result <- state$dbi_cache$result + cache_dir <- state$dbi_cache$result_cache_dir + expect_true(DBI::dbIsValid(result)) + expect_true(length(list.files(cache_dir)) > 0L) + expect_silent(replacement <- sess:::dataview_register(descending, view_id = a$view_id)) + views <- c(views, list(replacement)) + expect_false(DBI::dbIsValid(result)) + expect_false(dir.exists(cache_dir)) + expect_true(replacement$state_generation > a$state_generation) + sess:::handle_dataview_dispose(a) + expect_identical(page(replacement)$rows[["1"]], 25010:25009) + + # Reusing a dedicated connection outside the viewer must not silently + # restart the query or append a new execution to previously cached rows. + expect_warning(DBI::dbGetQuery(con, "SELECT 1"), "Closing open result set") + cache <- sess:::dataview_get_state(replacement$view_id)$dbi_cache + user_result <- DBI::dbSendQuery(con, "SELECT id FROM items ORDER BY id") + expect_identical(DBI::dbFetch(user_result, n = 2L)$id, 1:2) + for (attempt in 1:2) { + expect_error(page(replacement, 6000L), "dedicated connection") + expect_equal(cache$result_cache_rows, 5000L) + expect_false(cache$result_complete) + expect_identical(page(replacement)$rows[["1"]], 25010:25009) + } + expect_identical(page(b)$rows[["1"]], 125010:125009) + expect_silent(sess:::handle_dataview_dispose(replacement)) + expect_true(DBI::dbIsValid(user_result)) + expect_identical(DBI::dbFetch(user_result, n = 2L)$id, 3:4) + DBI::dbClearResult(user_result) + expect_true(DBI::dbIsValid(con)) + expect_equal(DBI::dbGetQuery(con, "SELECT COUNT(*) AS n FROM items")$n, 25010L) +}) + +# Two dedicated connections to the SAME database must also remain independent. +if (requireNamespace("RSQLite", quietly = TRUE) && + requireNamespace("arrow", quietly = TRUE) && + requireNamespace("dbplyr", quietly = TRUE)) local({ + path <- tempfile(fileext = ".sqlite") + con <- DBI::dbConnect(RSQLite::SQLite(), path) + other <- DBI::dbConnect(RSQLite::SQLite(), path) + views <- list() + on.exit({ + for (view in views) sess:::handle_dataview_dispose(view) + DBI::dbDisconnect(con) + DBI::dbDisconnect(other) + unlink(path) + }, add = TRUE) + DBI::dbWriteTable(con, "items", data.frame(id = 1:12000)) + tables <- list(dplyr::tbl(con, "items"), dplyr::tbl(other, "items")) + for (table in tables) { + views[[length(views) + 1L]] <- sess:::dataview_register(dplyr::arrange(table, id)) + } + for (start in c(0L, 6000L, 0L, 11000L)) { + for (view in views) { + result <- sess:::handle_dataview_page(list( + view_id = view$view_id, startRow = start, endRow = start + 2L + )) + expect_identical(result$rows[["1"]], seq.int(start + 1L, start + 2L)) + } + } +}) diff --git a/sess/inst/tinytest/test-dataview-dbi-recovery.R b/sess/inst/tinytest/test-dataview-dbi-recovery.R new file mode 100644 index 000000000..74bb3b1b2 --- /dev/null +++ b/sess/inst/tinytest/test-dataview-dbi-recovery.R @@ -0,0 +1,99 @@ +# A successful fetch advances the real cursor before its rows enter the cache. +if (requireNamespace("RSQLite", quietly = TRUE) && + requireNamespace("arrow", quietly = TRUE) && + requireNamespace("dbplyr", quietly = TRUE)) local({ + con <- DBI::dbConnect(RSQLite::SQLite(), ":memory:") + on.exit(DBI::dbDisconnect(con), add = TRUE) + fixture <- data.frame(id = seq_len(35010L)) + DBI::dbWriteTable(con, "items", fixture) + query <- dplyr::arrange(dplyr::tbl(con, "items"), id) + + for (failure in c("error", "interrupt")) { + for (stage in c("cache", "completion", "spill")) { + for (prefill in if (stage == "spill") TRUE else c(FALSE, TRUE)) local({ + view <- sess:::dataview_register(query) + on.exit(sess:::handle_dataview_dispose(view), add = TRUE) + cache <- sess:::dataview_get_state(view$view_id)$dbi_cache + page <- function(start = 0L, end = start + 2L) { + sess:::handle_dataview_page(list( + view_id = view$view_id, startRow = start, endRow = end + )) + } + if (prefill) { + page(30000L) + expect_equal(cache$result_cache_rows, 35000L) + expect_true(length(list.files(cache$result_cache_dir)) > 0L) + } + + # Inject at cache entry, after the completion check, or after a spill + # rename succeeds but before the cache records the finalized file. + cursor <- NULL + disk_files <- character() + inject <- function() { + cursor <<- cache$result + disk_files <<- list.files(cache$result_cache_dir, full.names = TRUE) + stop(structure( + list(message = "injected after fetch", call = NULL), + class = c(failure, "condition") + )) + } + injection <- substitute(fail(), list(fail = inject)) + name <- switch(stage, + cache = "dataview_dbi_result_cache_add", + completion = "dbHasCompleted", + spill = "file.rename" + ) + where <- switch(stage, + cache = asNamespace("sess"), + completion = asNamespace("DBI"), + spill = baseenv() + ) + if (stage == "cache") { + suppressMessages(trace(name, where = where, tracer = injection, print = FALSE)) + } else { + suppressMessages(trace(name, where = where, exit = injection, print = FALSE)) + } + traced <- TRUE + on.exit({ + if (traced) suppressMessages(untrace(name, where = where)) + }, add = TRUE) + caught <- tryCatch( + page(if (prefill) 35000L else 0L), + error = identity, interrupt = identity + ) + suppressMessages(untrace(name, where = where)) + traced <- FALSE + + expect_true(inherits(caught, failure), info = stage) + expect_identical(conditionMessage(caught), "injected after fetch") + expect_false(DBI::dbIsValid(cursor)) + expect_null(cache$result) + expect_equal(cache$result_cache_rows, 0L) + expect_length(cache$result_cache_blocks, 0L) + expect_false(cache$result_complete) + expect_length(list.files(cache$result_cache_dir), 0L) + if (stage == "spill") { + # A newly finalized file exists before its path is published. + expect_true(length(disk_files) >= 2L) + expect_false(any(file.exists(disk_files))) + } + + # Retrying must start at row one and reproduce every row exactly once, + # including the spilled prefix and the final partial batch. + first <- page() + expect_identical(first$rows[["0"]], 1:2) + expect_identical(first$rows[["1"]], 1:2) + expect_equal(first$fetchedRows, 5000L) + expect_equal(first$lastRow, -1L) + all <- page(0L, nrow(fixture) + 2L) + expect_identical(all$rows[["0"]], fixture$id) + expect_identical(all$rows[["1"]], fixture$id) + expect_equal(all$totalRows, nrow(fixture)) + expect_equal(all$lastRow, nrow(fixture)) + expect_true(cache$result_complete) + expect_null(cache$result) + expect_true(DBI::dbIsValid(con)) + }) + } + } +}) diff --git a/sess/inst/tinytest/test-dataview-dbi.R b/sess/inst/tinytest/test-dataview-dbi.R new file mode 100644 index 000000000..65d64edec --- /dev/null +++ b/sess/inst/tinytest/test-dataview-dbi.R @@ -0,0 +1,240 @@ +# Exercise the DBI cursor contract without a database server or credentials. +if (requireNamespace("DBI", quietly = TRUE) && + requireNamespace("dbplyr", quietly = TRUE) && + requireNamespace("arrow", quietly = TRUE)) local({ + class_env <- environment() + methods::setClass("dataview_dbi_test", contains = "DBIConnection", where = class_env) + methods::setClass("dataview_dbi_test_result", contains = "DBIResult", where = class_env) + registerS3method( + "dbplyr_edition", "dataview_dbi_test", function(con) 2L, envir = asNamespace("dbplyr") + ) + registered <- list() + method <- function(name, signature, definition) { + assign(name, getExportedValue("DBI", name), envir = class_env) + methods::setMethod(name, signature, definition, where = class_env) + registered[[length(registered) + 1L]] <<- list(name, signature) + } + on.exit({ + for (entry in registered) methods::removeMethod(entry[[1L]], entry[[2L]], where = class_env) + methods::removeClass("dataview_dbi_test_result", where = class_env) + methods::removeClass("dataview_dbi_test", where = class_env) + rm( + list = "dbplyr_edition.dataview_dbi_test", + envir = get(".__S3MethodsTable__.", envir = asNamespace("dbplyr")) + ) + }, add = TRUE) + con <- methods::new("dataview_dbi_test") + calls <- character() + fail <- FALSE + result_statement <- NULL + result_next <- NULL + result_open <- FALSE + fixture <- data.frame( + id = seq_len(6010L), + " day " = rep(as.Date(c("2024-02-29", NA_character_)), 3005L), + event = rep(as.POSIXct("2024-02-29 12:34:56", tz = "UTC"), 6010L), + precise = "2024-02-29T23:59:59.1234567", + offset = "2024-02-29T12:34:56.1234567+11:00", + clock = "12:34:56.1234567", + note = "2024-02-29 is ordinary text", + flag = rep(c(TRUE, FALSE), 3005L), + check.names = FALSE + ) + fixture$event[[1L]] <- fixture$event[[1L]] + 0.125 + fixture$event[[2L]] <- NA + if (requireNamespace("bit64", quietly = TRUE)) { + fixture$big <- bit64::as.integer64(rep("9007199254740993", nrow(fixture))) + } + method("dbIsValid", "dataview_dbi_test", function(dbObj, ...) TRUE) + method("dbIsValid", "dataview_dbi_test_result", function(dbObj, ...) result_open) + method("dbGetInfo", "dataview_dbi_test", function(dbObj, ...) { + list(dbms.name = "Generic database") + }) + method("dbListFields", c("dataview_dbi_test", "character"), function(conn, name, ...) { + names(fixture) + }) + method("dbQuoteIdentifier", c("dataview_dbi_test", "character"), function(conn, x, ...) { + DBI::dbQuoteIdentifier(DBI::ANSI(), x) + }) + method("dbQuoteLiteral", c("dataview_dbi_test", "ANY"), function(conn, x, ...) { + DBI::dbQuoteLiteral(DBI::ANSI(), x) + }) + method("dbGetQuery", c("dataview_dbi_test", "character"), function(conn, statement, ...) { + calls <<- c(calls, statement) + if (isTRUE(fail)) stop("test database failure") + if (identical(fail, "interrupt")) { + stop(structure(list(message = "cancelled", call = NULL), class = c("interrupt", "condition"))) + } + if (grepl("WHERE (0 = 1)", statement, fixed = TRUE)) return(fixture[0L, , drop = FALSE]) + stop("unexpected dbGetQuery") + }) + method("dbSendQuery", c("dataview_dbi_test", "character"), function(conn, statement, ...) { + calls <<- c(calls, statement) + if (isTRUE(fail)) stop("test database failure") + if (identical(fail, "interrupt")) { + stop(structure(list(message = "cancelled", call = NULL), class = c("interrupt", "condition"))) + } + result_statement <<- statement + result_next <<- 1L + result_open <<- TRUE + methods::new("dataview_dbi_test_result") + }) + method("dbFetch", "dataview_dbi_test_result", function(res, n = -1, ...) { + if (isTRUE(fail)) stop("test database failure") + if (identical(fail, "interrupt")) { + stop(structure(list(message = "cancelled", call = NULL), class = c("interrupt", "condition"))) + } + positions <- seq_along(fixture) + last <- min(nrow(fixture), result_next + n - 1L) + rows <- if (result_next > last) integer() else seq.int(result_next, last) + result_next <<- last + 1L + fixture[rows, positions, drop = FALSE] + }) + method("dbHasCompleted", "dataview_dbi_test_result", function(res, ...) { + result_next > nrow(fixture) + }) + method("dbClearResult", "dataview_dbi_test_result", function(res, ...) { + result_statement <<- NULL + result_next <<- NULL + result_open <<- FALSE + TRUE + }) + + tbl <- dplyr::tbl(con, "fixture") + env <- sess:::.sess_env + old <- env$dataviews + on.exit(env$dataviews <- old, add = TRUE) + original_fixture <- fixture + fixture <- fixture[rep(seq_len(nrow(fixture)), length.out = 25010L), , drop = FALSE] + fixture$id <- seq_len(nrow(fixture)) + queries <- list(tbl, dplyr::filter(tbl, id > 0L), dplyr::arrange(tbl, dplyr::desc(id))) + for (query in queries) local({ + view <- sess:::dataview_register(query) + on.exit(sess:::handle_dataview_dispose(list( + view_id = view$view_id, state_generation = view$state_generation + )), add = TRUE) + init <- sess:::handle_dataview_init(list(view_id = view$view_id)) + expect_null(init$totalRows) + expect_true(init$columnProjection) + expect_true(all(vapply(init$columns, function(column) { + isFALSE(column$filter) && isFALSE(column$sortable) + }, logical(1)))) + page <- function(start = 0L, end = 2L, fields = "1", controls = list()) { + sess:::handle_dataview_page(c(list( + view_id = view$view_id, startRow = start, endRow = end, fields = fields + ), controls)) + } + before <- length(calls) + first <- page() + expect_identical(first$rows[["1"]], 1:2) + expect_identical(first$lastRow, -1L) + expect_identical(tail(calls, 1L), as.character(dbplyr::sql_render(query))) + expect_equal(length(calls), before + 1L) + controls <- list( + sortModel = list(list(colId = "1", sort = "desc")), + filterModel = list("1" = list(type = "lessThan", filter = 0L)) + ) + expect_identical(page(controls = controls), first) + expect_identical(page(100L, 102L, "7")$rows[["7"]], fixture$note[101:102]) + expect_equal(length(calls), before + 1L) + expect_identical(names(page(fields = "0")$rows), "0") + expect_equal(nrow(page(10L, 10L)$rows), 0L) + last <- page(25000L, 25020L) + expect_identical(last$rows[["1"]], 25001:25010) + expect_equal(last$totalRows, 25010L) + expect_equal(last$lastRow, 25010L) + expect_false(result_open) + cache <- sess:::dataview_get_state(view$view_id)$dbi_cache + expect_true(any(vapply(cache$result_cache_blocks, function(block) { + !is.null(block$path) && file.exists(block$path) + }, logical(1)))) + expect_identical(page()$rows, first$rows) + expect_equal(length(calls), before + 1L) + expect_equal(nrow(page(26000L, 26010L)$rows), 0L) + # Dates, timestamps, NULLs and exact bigint survive spilling and projection. + rows <- page(fields = c("2", "3", "9"))$rows + wire <- jsonlite::fromJSON(jsonlite::toJSON(rows, na = "null", digits = NA)) + expect_identical(wire[["2"]][[1L]], "2024-02-29") + expect_true(is.na(wire[["2"]][[2L]])) + expect_identical(wire[["3"]][[1L]], "2024-02-29T12:34:56.125") + if ("9" %in% names(wire)) expect_identical(wire[["9"]][[1L]], "9007199254740993") + cache_dir <- cache$result_cache_dir + sess:::handle_dataview_dispose(list( + view_id = view$view_id, state_generation = view$state_generation + )) + expect_false(dir.exists(cache_dir)) + expect_true(DBI::dbIsValid(con)) + }) + + for (failure in list(TRUE, "interrupt")) local({ + view <- sess:::dataview_register(tbl) + on.exit({ + fail <<- FALSE + sess:::handle_dataview_dispose(list( + view_id = view$view_id, state_generation = view$state_generation + )) + }, add = TRUE) + page <- function(start, end) { + sess:::handle_dataview_page(list( + view_id = view$view_id, startRow = start, endRow = end, fields = "1" + )) + } + page(20000L, 20002L) + cache <- sess:::dataview_get_state(view$view_id)$dbi_cache + spilled <- list.files(cache$result_cache_dir, full.names = TRUE) + expect_true(length(spilled) > 0L) + fail <<- failure + caught <- tryCatch(page(25000L, 25002L), error = identity, interrupt = identity) + fail <<- FALSE + expect_true(inherits(caught, if (isTRUE(failure)) "error" else "interrupt")) + expect_false(result_open) + expect_equal(cache$result_cache_rows, 0L) + expect_length(cache$result_cache_blocks, 0L) + expect_false(cache$result_complete) + expect_false(any(file.exists(spilled))) + expect_identical(page(6000L, 6002L)$rows[["1"]], 6001:6002) + expect_identical(page(0L, 2L)$rows[["1"]], 1:2) + expect_true(DBI::dbIsValid(con)) + }) +}) + +# Execute actual SQL through a non-SQL Server DBI driver. +if (requireNamespace("RSQLite", quietly = TRUE) && + requireNamespace("arrow", quietly = TRUE) && + requireNamespace("dbplyr", quietly = TRUE)) local({ + con <- DBI::dbConnect(RSQLite::SQLite(), ":memory:") + on.exit(DBI::dbDisconnect(con), add = TRUE) + fixture <- data.frame(id = 1:12, label = rep(c("b", "a", NA), 4L)) + DBI::dbWriteTable(con, "odd table", fixture) + tbl <- dplyr::tbl(con, "odd table") + queries <- list( + tbl, + tbl |> dplyr::filter(id > 3L) |> dplyr::arrange(dplyr::desc(id)), + tbl |> dplyr::group_by(label) |> dplyr::summarise(total = sum(id, na.rm = TRUE)) |> + dplyr::arrange(dplyr::desc(total)), + head(dplyr::arrange(tbl, dplyr::desc(id)), 3L), + dplyr::filter(tbl, id < 0L) + ) + for (query in queries) local({ + expected <- as.data.frame(dplyr::collect(query)) + view <- sess:::dataview_register(query) + on.exit(sess:::handle_dataview_dispose(list( + view_id = view$view_id, state_generation = view$state_generation + )), add = TRUE) + init <- sess:::handle_dataview_init(list(view_id = view$view_id)) + wire <- jsonlite::fromJSON(jsonlite::toJSON(init)) + expect_true(all(wire$columns$filter == FALSE)) + expect_true(all(wire$columns$sortable == FALSE)) + result <- sess:::handle_dataview_page(list( + view_id = view$view_id, startRow = 0L, endRow = 20L, + sortModel = list(list(colId = "1", sort = "desc")), + filterModel = list("1" = list(type = "lessThan", filter = -1L)) + )) + actual <- result$rows[, -1L, drop = FALSE] + names(actual) <- names(expected) + rownames(actual) <- rownames(expected) <- NULL + expect_equal(actual, expected) + expect_equal(result$lastRow, nrow(expected)) + expect_true(DBI::dbIsValid(con)) + }) +}) diff --git a/sess/inst/tinytest/test-dataview-fragments.R b/sess/inst/tinytest/test-dataview-fragments.R new file mode 100644 index 000000000..684442568 --- /dev/null +++ b/sess/inst/tinytest/test-dataview-fragments.R @@ -0,0 +1,414 @@ +# Struct children must be bound by row, recursively, across reader/cache pages. +local({ + make_page <- function(id) { + page <- data.frame(id = id) + page$struct <- data.frame(flag = id %% 2L == 0L) + page$struct$child <- data.frame(date = as.Date("2020-01-01") + id) + page$struct$child$items <- lapply(id, function(i) c(i, -i)) + page + } + expect_equal( + sess:::dataview_arrow_bind_pages(list(make_page(1:2), make_page(3:5))), + make_page(1:5) + ) +}) + +if (requireNamespace("arrow", quietly = TRUE)) local({ + path <- tempfile("dataview-cached-result-", fileext = ".parquet") + on.exit(unlink(path), add = TRUE) + ids <- seq_len(50000L) + df <- data.frame( + id = ids, label = rep(c("a", NA_character_), length.out = length(ids)), + date = as.Date("2020-01-01") + ids, + time = as.POSIXct("2020-01-01", tz = "Australia/Sydney") + ids + ) + if (requireNamespace("bit64", quietly = TRUE)) { + df$big <- bit64::as.integer64("9007199254740993") + bit64::as.integer64(ids) + } + arrow::write_parquet(df, path, chunk_size = 10000L) + query <- dplyr::filter(arrow::open_dataset(path), id > 0L) + view <- sess:::dataview_register(query) + on.exit(sess:::handle_dataview_dispose(list( + view_id = view$view_id, state_generation = view$state_generation + )), add = TRUE) + # An unordered query can deliver Parquet batches in a different order on + # each platform/execution. Observe this execution, not a second collect(). + arrow_state <- sess:::dataview_get_state(view$view_id) + reader <- sess:::dataview_arrow_reader_open(query) + observed <- list() + arrow_state$arrow_reader$reader <- list( + read_next_batch = function() { + batch <- reader$read_next_batch() + if (!is.null(batch)) { + observed[[length(observed) + 1L]] <<- sess:::dataview_arrow_data_frame(batch) + } + batch + }, + Close = reader$Close + ) + # Finish one execution and spill its earlier batches before reading backward. + last <- sess:::handle_dataview_page(list( + view_id = view$view_id, startRow = 49990L, endRow = 50010L + )) + expect_equal(last$lastRow, nrow(df)) + streamed <- sess:::dataview_arrow_bind_pages(observed) + contents <- streamed[order(streamed$id), , drop = FALSE] + rownames(contents) <- NULL + expect_equal(contents, df) # Check membership independently of arrival order. + expect_true(any(vapply(arrow_state$arrow_reader$result_cache_batches, function(batch) { + !is.null(batch$path) && file.exists(batch$path) + }, logical(1)))) + + # DBI uses the same IPC storage after fetching its result into R blocks. + dbi_state <- list( + uses_result_cache = TRUE, column_names = names(df), + dbi = list(schema = df[0L, ]), dbi_cache = sess:::dataview_dbi_cache_state() + ) + sess:::dataview_dbi_result_cache_init(dbi_state, "conversion-test") + on.exit(sess:::dataview_dbi_result_cache_cleanup(dbi_state), add = TRUE) + for (start in seq.int(1L, nrow(df), by = 5000L)) { + sess:::dataview_dbi_result_cache_add(dbi_state$dbi_cache, df[start:(start + 4999L), ]) + } + dbi_state$dbi_cache$result_complete <- TRUE + + # Instrument only a local copy of each reader, checking the R conversion + # boundary as well as values, types, duplicates, and cross-block ordering. + for (backend in c("arrow", "dbi")) { + slice <- if (backend == "arrow") { + sess:::dataview_arrow_result_cache_slice + } else { + sess:::dataview_dbi_result_cache_slice + } + probe <- new.env(parent = environment(slice)) + probe$converted <- integer() + probe$dataview_arrow_data_frame <- function(data, conversion = NULL) { + probe$converted <- c(probe$converted, nrow(data)) + sess:::dataview_arrow_data_frame(data, conversion) + } + environment(slice) <- probe + for (positions in list(seq_along(df), 3:4, integer())) { + for (rows in list(1:2, 4991:5010, c(20001L, 1L, 10000L, 10001L, 1L, 50000L))) { + probe$converted <- integer() + page <- if (backend == "arrow") { + arrow_state$arrow_reader$projection <- positions + slice(arrow_state, rows) + } else { + slice(dbi_state, rows, positions) + } + page <- as.data.frame(page) + source <- if (backend == "arrow") streamed else df + reference <- source[rows, positions, drop = FALSE] + rownames(page) <- rownames(reference) <- NULL + expect_equal(page, reference, info = backend) + expect_true(sum(probe$converted) <= length(rows), info = backend) + if (length(positions)) expect_true(length(probe$converted) > 0L, info = backend) + } + } + } +}) + +if (requireNamespace("arrow", quietly = TRUE)) local({ + root <- tempfile("dataview-row-group-selection-") + dir.create(root) + path <- file.path(root, "year=2025") + dir.create(path) + on.exit(unlink(root, recursive = TRUE), add = TRUE) + ids <- seq_len(60000L) + df <- data.frame( + id = ids, + date = as.Date("2020-01-01") + ids, + time = as.POSIXct("2020-01-01", tz = "Australia/Sydney") + ids, + label = rep(c("a", NA_character_), length.out = length(ids)) + ) + # Flat schemas use the direct row-group path. Nested schemas are covered by + # the scanner-fallback regressions in test-dataview-recovery.R. + if (requireNamespace("bit64", quietly = TRUE)) { + df$big <- bit64::as.integer64("9007199254740993") + bit64::as.integer64(ids) + } + arrow::write_parquet(df, file.path(path, "single.parquet"), chunk_size = 30000L) + data <- arrow::open_dataset(root) + expected <- sess:::dataview_arrow_data_frame(data) + state <- sess:::dataview_to_state(data) + on.exit(sess:::dataview_arrow_reader_reset(state), add = TRUE) + + # Observe the conversion boundary without changing the package namespace. + # A small page must not convert the entire 30,000-row group into R. + slice <- sess:::dataview_arrow_row_group_slice + probe <- new.env(parent = environment(slice)) + converted <- integer() + probe$dataview_arrow_data_frame <- function(data, conversion = NULL) { + converted <<- c(converted, nrow(data)) + sess:::dataview_arrow_data_frame(data, conversion) + } + environment(slice) <- probe + normalize <- function(data) { + if (!is.data.frame(data)) return(data) + data <- as.data.frame(data) + data[] <- lapply(data, normalize) + rownames(data) <- NULL + data + } + for (fields in list(names(expected), c("id", "year"), "year")) { + projected <- sess:::dataview_arrow_page_state( + state, as.character(match(fields, names(expected))) + ) + for (rows in list( + 50001:50030, 20001:21000, 29991:30010, + c(60000L, 1L, 30001L, 30000L, 1L) + )) { + converted <- integer() + page <- slice(projected, rows) + reference <- expected[rows, names(projected$data), drop = FALSE] + expect_equal(normalize(page), normalize(reference)) + expect_equal(converted, if (identical(fields, "year")) integer() else length(rows)) + } + } +}) + +if (requireNamespace("arrow", quietly = TRUE)) local({ + root <- tempfile("dataview-projection-") + dir.create(root) + view_ids <- character() + on.exit({ + for (id in view_ids) sess:::handle_dataview_dispose(list( + view_id = id, + state_generation = sess:::dataview_get_state(id)$state_generation + )) + unlink(root, recursive = TRUE) + }, add = TRUE) + for (year in 2024:2025) { + path <- file.path(root, paste0("year=", year)) + dir.create(path) + ids <- (year - 2024L) * 6010L + seq_len(6010L) + df <- data.frame( + id = ids, score = -ids, flag = ids %% 2L == 0L, + date = as.Date("2020-01-01") + ids, + time = as.POSIXct("2020-01-01", tz = "Australia/Sydney") + ids + ) + df$struct <- data.frame(value = ids, flag = df$flag) + if (requireNamespace("bit64", quietly = TRUE)) { + df$big <- bit64::as.integer64("9007199254740993") + bit64::as.integer64(ids) + } + arrow::write_parquet(df, file.path(path, "part.parquet"), chunk_size = 997L) + } + data <- arrow::open_dataset(root) + expected <- sess:::dataview_arrow_data_frame(data) + id <- sess:::dataview_register(data)$view_id + expected_id <- sess:::dataview_register(expected)$view_id + view_ids <- c(id, expected_id) + field <- function(name) as.character(match(name, names(data))) + model <- list( + sortModel = list(list(colId = field("score"), sort = "asc")), + filterModel = setNames(list(list(type = "true")), field("flag")) + ) + expect_true(sess:::handle_dataview_init(list(view_id = id))$columnProjection) + expect_false(sess:::handle_dataview_init(list(view_id = expected_id))$columnProjection) + + fragment_index <- NULL + # Begin with a narrow projection: the fragment index must still retain the + # full source schema so revealing more columns later can reuse that index. + projections <- list( + c("0", field("year"), field("id")), + c("0", field("date"), field("time")), + "0", character(), + c("0", field("struct")), + c("0", field("id"), field("big")), + NULL + ) + for (fields in projections) { + for (start in c(0L, 4990L, 0L, 12020L)) { + params <- c(list(startRow = start, endRow = start + 30L), model) + reference <- sess:::handle_dataview_page(c(list(view_id = expected_id), params)) + actual <- sess:::handle_dataview_page(c(list(view_id = id, fields = fields), params)) + selected <- if (is.null(fields)) { + names(reference$rows) + } else { + names(reference$rows)[names(reference$rows) %in% c("0", fields)] + } + reference$rows <- reference$rows[, selected, drop = FALSE] + expect_equal(actual, reference) + reader <- sess:::dataview_get_state(id)$arrow_reader + expect_equal(reader$projection, which(as.character(seq_along(expected)) %in% selected)) + if (is.null(fragment_index)) fragment_index <- reader$fragment_index + expect_identical(reader$fragment_index, fragment_index) + for (block in reader$query_cache) { + expect_equal(names(block$data), names(data)[reader$projection]) + } + } + } + state <- sess:::dataview_get_state(id) + expect_identical(state$data, data) + expect_null(state$query_columns) + + # The same key used by a filter and a sort is collected once, and another + # key joins the same projected scan; their original row ordering is retained. + columns <- sess:::dataview_arrow_query_columns( + state, list(list(colId = field("score"), sort = "asc")), + setNames(list(list(type = "lessThan", filter = -10L), list(type = "true")), + c(field("score"), field("flag"))) + ) + expect_equal(names(columns), as.character(match(c("score", "flag"), names(data)) + 1L)) + expect_equal(columns[[1L]], expected$score) + expect_equal(columns[[2L]], expected$flag) +}) + +if (requireNamespace("arrow", quietly = TRUE)) local({ + root <- tempfile("dataview-fragments-") + dir.create(root) + view_ids <- character() + on.exit({ + for (id in view_ids) sess:::handle_dataview_dispose(list( + view_id = id, + state_generation = sess:::dataview_get_state(id)$state_generation + )) + unlink(root, recursive = TRUE) + }, add = TRUE) + + # Unequal files/row groups, an empty file, nulls, schema evolution, and + # partition values which are not present in the physical Parquet columns. + counts <- c(0L, 7001L, 3999L, 11000L, 8000L, 15000L, 15000L) + ends <- cumsum(counts) + starts <- c(1L, head(ends, -1L) + 1L) + for (layout in c("hive", "directory")) { + path <- file.path(root, layout) + dir.create(path) + for (i in seq_along(counts)) { + partition <- if (layout == "hive") paste0("year=", 2020L + i) else as.character(2020L + i) + if (layout == "hive" && i == 7L) partition <- "year=__HIVE_DEFAULT_PARTITION__" + directory <- file.path(path, partition) + dir.create(directory) + ids <- if (counts[[i]]) seq.int(starts[[i]], ends[[i]]) else integer() + df <- data.frame( + id = ids, + flag = c(TRUE, FALSE, NA, TRUE)[ids %% 4L + 1L], + score = (ids * 7919L) %% 60000L + ) + if (i %% 2L == 0L) df$extra <- ids + arrow::write_parquet(df, file.path(directory, "data.parquet"), chunk_size = 997L) + } + data <- if (layout == "hive") { + arrow::open_dataset(path, unify_schemas = TRUE) + } else { + arrow::open_dataset(path, partitioning = "year", unify_schemas = TRUE) + } + expected <- sess:::dataview_arrow_data_frame(data) + id <- sess:::dataview_register(data)$view_id + expected_id <- sess:::dataview_register(expected)$view_id + view_ids <- c(view_ids, id, expected_id) + field <- function(name) as.character(match(name, names(expected))) + sort <- list(sortModel = list(list(colId = field("score"), sort = "desc"))) + true <- list(filterModel = setNames(list(list(type = "true")), field("flag"))) + false <- list(filterModel = setNames(list(list(type = "false")), field("flag"))) + year <- list(filterModel = setNames( + list(list(type = "equals", filter = 2024L)), field("year") + )) + for (model in list(true, c(true, sort), false, c(false, sort), year, sort, list())) { + # Cross cache blocks, evict earlier sorted blocks, and revisit old pages. + for (start in c(0L, 4990L, 10000L, 15000L, 20000L, 25000L, 0L, 59990L)) { + params <- c(list(startRow = start, endRow = start + 30L), model) + actual <- sess:::handle_dataview_page(c(list(view_id = id), params)) + reference <- sess:::handle_dataview_page(c(list(view_id = expected_id), params)) + expect_equal(actual, reference, info = paste(layout, start)) + reader <- sess:::dataview_get_state(id)$arrow_reader + expect_true(sum(vapply(reader$query_cache, function(x) nrow(x$data), integer(1))) <= 20000L) + } + } + + state <- sess:::dataview_get_state(id) + index <- sess:::dataview_arrow_fragment_index(state) + expect_equal(index$ends, ends[-1L]) + expect_equal(nrow(index$partition_values), 6L) + # Duplicates and interleaved fragments must retain the requested order. + rows <- c(60000L, 1L, 7001L, 7002L, 1L, 45001L, 11001L) + page <- as.data.frame(sess:::dataview_arrow_fragment_slice(state, rows)) + rownames(page) <- NULL + reference <- as.data.frame(expected[rows, , drop = FALSE]) + rownames(reference) <- NULL + expect_equal(page, reference) + + # Large jumps and backward misses use fragments without moving the + # sequential reader, which must remain usable for subsequent nearby pages. + state <- sess:::dataview_to_state(data) + sess:::dataview_arrow_slice(state, 1:30) + reader <- state$arrow_reader$reader + next_row <- state$arrow_reader$next_row + for (rows in list(59001:59030, 45001:45030)) { + expect_equal(sess:::dataview_arrow_slice(state, rows)$id, expected$id[rows]) + expect_identical(state$arrow_reader$reader, reader) + expect_equal(state$arrow_reader$next_row, next_row) + } + expect_equal(sess:::dataview_arrow_slice(state, 1001:1030)$id, expected$id[1001:1030]) + expect_identical(state$arrow_reader$reader, reader) + expect_equal(state$arrow_reader$next_row, next_row + 1000L) + sess:::dataview_arrow_reader_reset(state) + } + + # A single Parquet file uses row groups for distant/backward random access. + path <- file.path(root, "single.parquet") + ids <- seq_len(12020L) + df <- data.frame( + id = ids, + score = (ids * 7919L) %% 12020L, + flag = ids %% 2L == 0L + ) + arrow::write_parquet(df, path, chunk_size = 997L) + data <- arrow::open_dataset(path) + expected <- sess:::dataview_arrow_data_frame(data) + state <- sess:::dataview_to_state(data) + sess:::dataview_arrow_slice(state, 1:30) + reader <- state$arrow_reader$reader + next_row <- state$arrow_reader$next_row + for (rows in list(11001:11030, 7001:7030)) { + expect_equal(sess:::dataview_arrow_slice(state, rows)$id, expected$id[rows]) + expect_identical(state$arrow_reader$reader, reader) + expect_equal(state$arrow_reader$next_row, next_row) + } + index <- state$arrow_reader$row_group_index + expect_equal(index$ends, c(seq(997, 11964, by = 997), 12020)) + expect_equal(sess:::dataview_arrow_slice(state, 1001:1030)$id, expected$id[1001:1030]) + expect_identical(state$arrow_reader$reader, reader) + expect_equal(state$arrow_reader$next_row, next_row + 1000L) + sess:::dataview_arrow_reader_reset(state) + + id <- sess:::dataview_register(data)$view_id + expected_id <- sess:::dataview_register(df)$view_id + view_ids <- c(view_ids, id, expected_id) + filter <- list(filterModel = list("3" = list(type = "true"))) + sort <- list(sortModel = list(list(colId = "2", sort = "desc"))) + for (model in list(sort, filter, c(filter, sort))) { + for (start in c(0L, 4990L, 9000L, 0L)) { + params <- c(list(startRow = start, endRow = start + 30L), model) + expect_equal( + sess:::handle_dataview_page(c(list(view_id = id), params)), + sess:::handle_dataview_page(c(list(view_id = expected_id), params)) + ) + } + } + + # Multi-file structs also cross reader batches and query-cache boundaries. + path <- file.path(root, "struct") + dir.create(path) + for (i in 1:3) { + ids <- seq.int((i - 1L) * 2100L + 1L, i * 2100L) + struct <- data.frame(flag = ids %% 2L == 0L, date = as.Date("2020-01-01") + ids) + struct$child <- data.frame(value = ids) + tab <- arrow::Table$create(id = ids, struct = arrow::StructArray$create(struct)) + if (i == 2L) tab <- arrow::Table$create(id = ids) + arrow::write_parquet(tab, file.path(path, paste0(i, ".parquet")), chunk_size = 997L) + } + # Explicit nonlexical file order must also match the column/index scans. + data <- arrow::open_dataset(rev(list.files(path, full.names = TRUE)), unify_schemas = TRUE) + expected <- sess:::dataview_arrow_data_frame(data) + id <- sess:::dataview_register(data)$view_id + expected_id <- sess:::dataview_register(expected)$view_id + view_ids <- c(view_ids, id, expected_id) + for (model in list(list(), list(sortModel = list(list(colId = "1", sort = "desc"))))) { + for (start in c(990L, 2090L, 4990L, 0L)) { + params <- c(list(startRow = start, endRow = start + 30L), model) + expect_equal( + sess:::handle_dataview_page(c(list(view_id = id), params)), + sess:::handle_dataview_page(c(list(view_id = expected_id), params)) + ) + } + } +}) diff --git a/sess/inst/tinytest/test-dataview-lazy.R b/sess/inst/tinytest/test-dataview-lazy.R new file mode 100644 index 000000000..05cfb1317 --- /dev/null +++ b/sess/inst/tinytest/test-dataview-lazy.R @@ -0,0 +1,198 @@ +# Viewer controls operate on the result of one execution, not on the source. +if (requireNamespace("arrow", quietly = TRUE)) local({ + path <- tempfile("dataview-lazy-query-", fileext = ".parquet") + views <- list() + on.exit({ + for (view in views) sess:::handle_dataview_dispose(list( + view_id = view$view_id, state_generation = view$state_generation + )) + unlink(path) + }, add = TRUE) + register <- function(data) { + view <- sess:::dataview_register(data) + views[[length(views) + 1L]] <<- view + view + } + ids <- seq_len(50010L) + df <- data.frame( + id = ids, score = ids %% 7L, flag = ids %% 2L == 0L, + label = rep(c("Alpha", "beta", NA_character_), length.out = length(ids)), + date = as.Date("2024-01-01") + ids %% 5L, + time = as.POSIXct("2024-01-01", tz = "UTC") + ids + ) + if (requireNamespace("bit64", quietly = TRUE)) { + df$big <- bit64::as.integer64("9007199254740993") + bit64::as.integer64(ids) + } + arrow::write_parquet(df, path, chunk_size = 5000L) + source <- arrow::open_dataset(path) + query <- dplyr::filter(source, id > 10L) + expected <- df[df$id > 10L, ] + rownames(expected) <- NULL + view <- register(query) + stream_view_id <- view$view_id + stream_reader <- sess:::dataview_arrow_reader_open(query) + observed <- list() + observed_rows <- 0L + reference_rows <- -1L + stream_cache <- sess:::dataview_get_state(view$view_id)$arrow_reader + stream_cache$reader <- list( + read_next_batch = function() { + batch <- stream_reader$read_next_batch() + if (!is.null(batch)) { + observed[[length(observed) + 1L]] <<- sess:::dataview_arrow_data_frame(batch) + observed_rows <<- observed_rows + nrow(batch) + } + batch + }, + Close = stream_reader$Close + ) + reference <- register(expected) + columns <- sess:::handle_dataview_init(list(view_id = view$view_id))$columns + expect_true(all(vapply(columns[-1L], function(column) { + !isFALSE(column$filter) && !isFALSE(column$sortable) + }, logical(1)))) + field <- function(name) as.character(match(name, names(df))) + request <- function(model = list(), start = 0L, fields = NULL) { + params <- c(list(startRow = start, endRow = start + 20L, fields = fields), model) + actual <- sess:::handle_dataview_page(c(list(view_id = view$view_id), params)) + # Preserve the order of THIS execution, including stable sort ties and + # row labels. A separate collect() may deliver batches in another order. + if (identical(view$view_id, stream_view_id) && reference_rows != observed_rows) { + expected <<- sess:::dataview_arrow_bind_pages(observed) + reference <<- register(expected) + reference_rows <<- observed_rows + } + expected_page <- sess:::handle_dataview_page(c(list(view_id = reference$view_id), params)) + if (!is.null(fields)) { + selected <- names(expected_page$rows) %in% c("0", fields) + expected_page$rows <- expected_page$rows[, selected, drop = FALSE] + } + expect_equal(actual$rows, expected_page$rows) + if (actual$lastRow != -1L) { + expect_equal(actual$totalRows, expected_page$totalRows) + expect_equal(actual$totalUnfiltered, nrow(expected)) + expect_equal(actual$lastRow, expected_page$totalRows) + } + actual + } + expect_equal(request()$lastRow, -1L) + cache <- sess:::dataview_get_state(view$view_id)$arrow_reader + expect_false(cache$result_complete) + reader <- cache$reader + sort <- list(sortModel = list( + list(colId = field("score"), sort = "desc"), list(colId = field("label"), sort = "asc") + )) + request(sort, fields = field("id")) + expect_true(cache$result_complete) + contents <- expected[order(expected$id), , drop = FALSE] + rownames(contents) <- NULL + source_contents <- df[df$id > 10L, ] + rownames(source_contents) <- NULL + expect_equal(contents, source_contents) + expect_identical(cache$reader, reader) + expect_true(any(vapply(cache$result_cache_batches, function(batch) { + !is.null(batch$path) && file.exists(batch$path) + }, logical(1)))) + expect_true(sum(vapply(cache$result_cache_batches, function(batch) { + if (is.null(batch$data)) 0L else nrow(batch$data) + }, integer(1))) <= sess:::dataview_arrow_cache_rows) + indices <- cache$result_query_indices + request(sort, 4990L, c(field("date"), field("time"))) + expect_identical(cache$result_query_indices, indices) + + # TOP N and aggregation remain inside the original query's result boundary. + for (q in list( + head(dplyr::arrange(query, dplyr::desc(id)), 100L), + dplyr::arrange(dplyr::summarise(dplyr::group_by(query, flag), total = sum(id)), flag), + dplyr::filter(query, id < 0L) + )) { + saved_view <- view + saved_reference <- reference + saved_expected <- expected + expected <- as.data.frame(sess:::dataview_arrow_data_frame(q)) + rownames(expected) <- NULL + view <- register(q) + reference <- register(expected) + request(list(sortModel = list(list(colId = "1", sort = "asc")))) + request(list(filterModel = list("1" = list(type = "lessThan", filter = 0L)))) + request() + view <- saved_view + reference <- saved_reference + expected <- saved_expected + } + + # Forbid reopening the source after completion. This also works on Windows, + # where the finished reader can retain a file handle until disposal. + suppressMessages(trace( + "dataview_arrow_reader_open", where = asNamespace("sess"), + tracer = quote(stop("viewer controls must reuse the cached execution")), print = FALSE + )) + on.exit(suppressMessages(untrace( + "dataview_arrow_reader_open", where = asNamespace("sess") + )), add = TRUE) + models <- list( + c(list(filterModel = setNames(list(list(type = "true")), field("flag"))), sort), + list(filterModel = setNames(list(list(type = "true")), field("flag"))), + list(filterModel = setNames(list(list(type = "contains", filter = "ALP")), field("label"))), + list(filterModel = setNames(list(list(type = "blank")), field("label"))), + list(filterModel = setNames( + list(list(type = "equals", dateFrom = "2024-01-03")), field("date") + )), + list(filterModel = list("1" = list(type = "lessThan", filter = 0L))), + list(), sort + ) + if ("big" %in% names(df)) models <- c(models, list(list( + filterModel = setNames( + list(list(type = "equals", filter = as.character(df$big[[21L]]))), field("big") + ), + sortModel = list(list(colId = field("big"), sort = "desc")) + ))) + for (model in models) { + for (start in c(0L, 4990L, 49990L, 0L)) request(model, start, field("id")) + request(model, fields = character()) + expect_identical(cache$reader, reader) + } + request() + expect_null(cache$result_query_indices) +}) + +# Exercise a non-file-order stream on every OS, independent of thread timing. +if (requireNamespace("arrow", quietly = TRUE)) local({ + df <- data.frame(id = 1:25000, key = rep(1:5, 5000), text = as.character(1:25000)) + blocks <- lapply(c(3L, 1L, 5L, 2L, 4L), function(i) { + df[seq.int((i - 1L) * 5000L + 1L, i * 5000L), , drop = FALSE] + }) + expected <- do.call(rbind, blocks) + rownames(expected) <- NULL + query <- dplyr::filter(arrow::Table$create(df), id > 0L) + view <- sess:::dataview_register(query) + reference <- sess:::dataview_register(expected) + on.exit({ + sess:::handle_dataview_dispose(view) + sess:::handle_dataview_dispose(reference) + }, add = TRUE) + cache <- sess:::dataview_get_state(view$view_id)$arrow_reader + cache$reader <- arrow::RecordBatchReader$create( + batches = lapply(blocks, arrow::as_record_batch) + ) + request <- function(start, model = list(), fields = NULL) { + params <- c(list(startRow = start, endRow = start + 20L, fields = fields), model) + actual <- sess:::handle_dataview_page(c(list(view_id = view$view_id), params)) + wanted <- sess:::handle_dataview_page(c(list(view_id = reference$view_id), params)) + if (!is.null(fields)) wanted$rows <- wanted$rows[, c("0", fields), drop = FALSE] + expect_equal(actual$rows, wanted$rows) + if (actual$lastRow >= 0L) expect_equal(actual$totalRows, wanted$totalRows) + } + request(0L) + expect_false(cache$result_complete) + sort <- list(sortModel = list(list(colId = "2", sort = "asc"))) + filter <- list(filterModel = list("2" = list(type = "equals", filter = 3L))) + for (model in list(sort, filter, c(filter, sort), list())) { + for (start in c(0L, 4990L, 24980L, 0L)) request(start, model, c("1", "3")) + } + expect_true(cache$result_complete) + expect_true(any(vapply(cache$result_cache_batches, function(batch) { + !is.null(batch$path) && file.exists(batch$path) + }, logical(1)))) + expect_null(cache$result_query_indices) +}) diff --git a/sess/inst/tinytest/test-dataview-recovery.R b/sess/inst/tinytest/test-dataview-recovery.R new file mode 100644 index 000000000..7707b7c9d --- /dev/null +++ b/sess/inst/tinytest/test-dataview-recovery.R @@ -0,0 +1,137 @@ +# These regressions use actual Arrow execution and the public View() hook. +if (requireNamespace("arrow", quietly = TRUE) && + requireNamespace("dplyr", quietly = TRUE)) local({ + fixtures <- new.env(parent = emptyenv()) + fixtures$paths <- character() + fixtures$views <- list() + on.exit({ + for (view in fixtures$views) sess:::handle_dataview_dispose(view) + unlink(fixtures$paths) + }, add = TRUE) + dataset <- function(data) { + path <- tempfile(fileext = ".parquet") + fixtures$paths <- c(fixtures$paths, path) + arrow::write_parquet(data, path, chunk_size = 5000L) + arrow::open_dataset(path) + } + register <- function(data, ...) { + view <- sess:::dataview_register(data, ...) + fixtures$views[[length(fixtures$views) + 1L]] <- view + view + } + page <- function(view, start = 0L, fields = NULL, sort = list()) { + sess:::handle_dataview_page(list( + view_id = view$view_id, startRow = start, endRow = start + 2L, + fields = fields, sortModel = sort + )) + } + descending <- list(list(colId = "1", sort = "desc")) + df <- data.frame(id = 1:30000, text = as.character(1:30000)) + df$text[[25001L]] <- "oops" + query <- dplyr::mutate(dataset(df), number = as.integer(text)) + view <- register(query) + cache <- sess:::dataview_get_state(view$view_id)$arrow_reader + for (attempt in 1:2) { + # Previously the second request succeeded with a false total of 25000. + expect_error(page(view, sort = descending), "Failed to parse string") + expect_false(cache$result_complete) + expect_equal(cache$result_cache_rows, 0L) + expect_null(cache$reader) + expect_null(cache$result_query_key) + expect_null(cache$result_query_indices) + expect_length(list.files(cache$result_cache_dir), 0L) + } + + # Cancellation and transient errors must restart from row one without mixing + # executions, and must remove any previously spilled blocks. + df$text[[25001L]] <- "25001" + query <- dplyr::arrange(dplyr::mutate(dataset(df), number = as.integer(text)), id) + for (condition in c("error", "interrupt")) { + view <- register(query) + state <- sess:::dataview_get_state(view$view_id) + cache <- state$arrow_reader + # Sorting may coalesce the entire result into one batch. Use fixed batches + # for the spill/cancellation setup; retry still executes the real query. + cache$reader <- arrow::RecordBatchReader$create(batches = lapply(0:5, function(i) { + rows <- seq.int(i * 5000L + 1L, (i + 1L) * 5000L) + arrow::as_record_batch(data.frame(df[rows, ], number = rows)) + })) + sess:::dataview_arrow_result_cache_slice(state, 25000L) + expect_true(length(list.files(cache$result_cache_dir)) > 0L) + reader <- cache$reader + cache$reader <- list( + read_next_batch = function() { + stop(structure( + list(message = "injected read failure", call = NULL), + class = c(condition, "condition") + )) + }, + Close = reader$Close + ) + caught <- tryCatch(page(view, sort = descending), + error = identity, interrupt = identity) + expect_true(inherits(caught, condition)) + expect_equal(cache$result_cache_rows, 0L) + expect_length(list.files(cache$result_cache_dir), 0L) + result <- page(view, sort = descending) + expect_equal(result$totalRows, 30000L) + expect_identical(result$rows[["1"]], 30000:29999) + expect_identical(page(view)$rows[["1"]], 1:2) + } + + # A struct contributes multiple Parquet leaves. Scalars following it must + # remain available on distant reads and when only that scalar is projected. + df <- data.frame(id = 1:20000) + df$nested <- data.frame(a = 1:20000, b = 20001:40000) + df$text <- as.character(df$id) + source <- dataset(df) + view <- register(source) + reference <- register(df) + for (fields in list(NULL, "3", c("2", "3"))) { + for (start in c(0L, 15000L, 10000L, 0L)) { + actual <- page(view, start, fields)$rows + expected <- page(reference, start)$rows + if (!is.null(fields)) expected <- expected[, c("0", fields), drop = FALSE] + expect_equal(actual, expected) + } + } + projected <- register(source$WithSchema(source$schema[c("id", "text")])) + expect_identical(page(projected, 15000L)$rows[["2"]], as.character(15001:15002)) +}) + +if (requireNamespace("arrow", quietly = TRUE) && + requireNamespace("dplyr", quietly = TRUE)) local({ + env <- sess:::.sess_env + old_interactive <- env$interactive_connected + name <- "sess_review_aggregate_query" + on.exit({ + env$interactive_connected <- old_interactive + sess:::runtime_stop() + rm(list = name, envir = .GlobalEnv) + }, add = TRUE) + sess:::runtime_start(use_rstudioapi = FALSE, plot_backend = "native") + query <- arrow::Table$create(data.frame(id = 1:10, k = rep(1:2, 5))) |> + dplyr::group_by(k) |> + dplyr::summarise(n = dplyr::n()) + expect_true(is.na(nrow(query))) + assign(name, query, envir = .GlobalEnv) + invisible(sess:::get_workspace_data()) + expect_true(is.na(get(name, env$workspace_dimensions)$dim[[1L]])) + for (interactive in c(FALSE, TRUE)) { + env$interactive_connected <- interactive + utils::View(query, title = name) + id <- get(paste0("table:", name), env$dataview_registry) + metadata <- sess:::handle_dataview_init(list(view_id = id)) + expect_null(metadata$totalRows) + result <- sess:::handle_dataview_page(list(view_id = id, startRow = 0L, endRow = 10L)) + expect_equal(result$totalRows, 2L) + expect_identical(sort(result$rows[["1"]]), 1:2) + expect_equal(as.numeric(result$rows[["2"]]), c(5, 5)) + } + # The same route applies to dbplyr queries: no eager snapshot or row count. + if (requireNamespace("dbplyr", quietly = TRUE)) { + db_query <- dbplyr::tbl_lazy(data.frame(id = 1L), con = dbplyr::simulate_mssql()) + class(db_query) <- c("tbl_sql", class(db_query)) + expect_false(sess:::.interactive_rich_value(db_query)) + } +}) diff --git a/sess/inst/tinytest/test-dataview-reproducibility.R b/sess/inst/tinytest/test-dataview-reproducibility.R index 38481842a..34936eed8 100644 --- a/sess/inst/tinytest/test-dataview-reproducibility.R +++ b/sess/inst/tinytest/test-dataview-reproducibility.R @@ -24,6 +24,9 @@ local({ sess:::handle_dataview_init(list(view_id = view)) sess:::handle_dataview_page(list(view_id = view, startRow = 20L, endRow = 40L)) expect_identical(.Random.seed, before) - sess:::handle_dataview_dispose(list(view_id = view)) + sess:::handle_dataview_dispose(list( + view_id = view, + state_generation = sess:::dataview_get_state(view)$state_generation + )) expect_identical(.Random.seed, before) }) diff --git a/sess/inst/tinytest/test-dataview.R b/sess/inst/tinytest/test-dataview.R index 512998d20..c9b76e246 100644 --- a/sess/inst/tinytest/test-dataview.R +++ b/sess/inst/tinytest/test-dataview.R @@ -116,7 +116,7 @@ local({ table = data.frame(value = 1:2), value = list(1 + 2i)) formatted <- sess:::dataview_format_column(list(nested))[[1L]] expect_identical(formatted$date, nested$date) - expect_identical(formatted$table, nested$table) + expect_equal(formatted$table, nested$table) expect_equal(formatted$value, list("1+2i")) }) @@ -288,6 +288,607 @@ local({ } }) + +# A Dataset stand-in verifies row-fetch counts without scanning real files. +if (requireNamespace("arrow", quietly = TRUE)) local({ + row_takes <- 0L + subset_arrow_lazy <- function(x, i, j, ..., drop = FALSE) { + if (!missing(i) && length(i)) { + row_takes <<- row_takes + 1L + } + class_name <- class(x)[[1L]] + data <- structure(x, class = "data.frame") + if (missing(i)) { + return(structure(data[, j, drop = drop], class = c(class_name, "data.frame"))) + } + page <- if (missing(j)) data[i, , drop = FALSE] else data[i, j, drop = drop] + page + } + registerS3method("[", "dataview_dataset_test", subset_arrow_lazy) + registerS3method("dim", "dataview_dataset_test", function(x) { + dim(structure(x, class = "data.frame")) + }) + registerS3method("names", "dataview_dataset_test", function(x) { + names(structure(x, class = "data.frame")) + }) + on.exit({ + rm( + list = c( + "[.dataview_dataset_test", "dim.dataview_dataset_test", "names.dataview_dataset_test" + ), + envir = get(".__S3MethodsTable__.", envir = asNamespace("base")) + ) + }, add = TRUE) + + for (class_name in "dataview_dataset_test") { + data <- structure( + data.frame( + id = 1:5, + value = c("a", "b", "c", "d", "e"), + score = c(5, 4, 3, 2, 1) + ), + class = c(class_name, "Dataset", "data.frame") + ) + + expect_true(sess:::dataview_is_table(data), info = class_name) + state <- sess:::dataview_to_state(data) + expect_equal(state$total_rows, 5L, info = class_name) + expect_equal( + vapply(state$columns[-1L], function(x) as.character(x$headerName), ""), + names(data), + info = class_name + ) + + page <- sess:::dataview_rows(state, 2:3) + expect_equal(page[["1"]], 2:3, info = class_name) + expect_equal(page[["2"]], c("b", "c"), info = class_name) + expect_equal(page[["3"]], c(4, 3), info = class_name) + + filtered <- sess:::dataview_query_indices( + state, NULL, list("2" = list(type = "equals", filter = "c")) + ) + expect_equal(filtered, 3L, info = class_name) + + sorted <- sess:::dataview_query_indices( + state, list(list(colId = "3", sort = "asc")), NULL + ) + expect_equal(sorted, 5:1, info = class_name) + } + + data <- structure( + data.frame(id = 1:6000, score = 6000:1), + class = c("dataview_dataset_test", "Dataset", "data.frame") + ) + state <- sess:::dataview_to_state(data) + # This data-frame stand-in has R row names; real Dataset objects do not. + state$row_index <- NULL + sort_model <- list(list(colId = "2", sort = "asc")) + state$query_key <- sess:::dataview_query_key(sort_model, NULL) + state$query_indices <- sess:::dataview_query_indices(state, sort_model, NULL) + state$query_has_sort <- TRUE + row_takes <- 0L + + page1 <- sess:::dataview_rows( + state, + state$query_indices[1:500], + use_arrow_query_cache = TRUE, + display_idx = 1:500 + ) + page2 <- sess:::dataview_rows( + state, + state$query_indices[501:1000], + use_arrow_query_cache = TRUE, + display_idx = 501:1000 + ) + + expect_equal(row_takes, 1L) + expect_equal(page1[["0"]], 6000:5501) + expect_equal(page1[["2"]], 1:500) + expect_equal(page2[["0"]], 5500:5001) + expect_equal(page2[["2"]], 501:1000) + + page11 <- sess:::dataview_rows( + state, + state$query_indices[5001:5500], + use_arrow_query_cache = TRUE, + display_idx = 5001:5500 + ) + expect_equal(row_takes, 2L) + expect_equal(page11[["0"]], 1000:501) + expect_equal(page11[["2"]], 5001:5500) + + # Returning to an earlier sorted block does not rescan the source. + again <- sess:::dataview_rows( + state, state$query_indices[1:500], use_arrow_query_cache = TRUE, display_idx = 1:500 + ) + expect_equal(again, page1) + expect_equal(row_takes, 2L) + + # The first 5,000 display rows were prefetched by the first query-cache miss. + for (start in c(1001L, 2001L, 3001L, 4001L)) { + sess:::dataview_rows( + state, + state$query_indices[start:(start + 999L)], + use_arrow_query_cache = TRUE, + display_idx = start:(start + 999L) + ) + } + expect_equal(row_takes, 2L) + expect_length(state$arrow_reader$row_cache, 0L) + expect_equal( + sum(vapply(state$arrow_reader$query_cache, function(x) nrow(x$data), integer(1))), + 6000L + ) + again11 <- sess:::dataview_rows( + state, + state$query_indices[5001:5500], + use_arrow_query_cache = TRUE, + display_idx = 5001:5500 + ) + expect_equal(again11, page11) + expect_equal(row_takes, 2L) +}) + + +# Real Arrow Dataset pages reuse one forward-only reader when Arrow is available. +if (requireNamespace("arrow", quietly = TRUE)) { + local({ + Array <- getExportedValue("arrow", "Array") + Table <- getExportedValue("arrow", "Table") + int32 <- getExportedValue("arrow", "int32") + list_of <- getExportedValue("arrow", "list_of") + + data <- Table$create( + list_col = Array$create( + list(c(1L, 2L), c(3L, 4L, 5L), integer()), + type = list_of(int32()) + ) + ) + state <- sess:::dataview_to_state(data) + page <- sess:::dataview_rows(state, 1:3) + + expect_true(is.list(page[["1"]])) + expect_false(inherits(page[["1"]], "vctrs_list_of")) + expect_equal(page[["1"]][[1L]], c(1L, 2L)) + expect_silent(jsonlite::toJSON( + page, auto_unbox = TRUE, null = "null", force = TRUE, digits = NA + )) + }) + + local({ + schema <- getExportedValue("arrow", "schema") + Table <- getExportedValue("arrow", "Table") + InMemoryDataset <- getExportedValue("arrow", "InMemoryDataset") + int32 <- getExportedValue("arrow", "int32") + utf8 <- getExportedValue("arrow", "utf8") + list_of <- getExportedValue("arrow", "list_of") + fixed_size_list_of <- getExportedValue("arrow", "fixed_size_list_of") + map_of <- getExportedValue("arrow", "map_of") + struct <- getExportedValue("arrow", "struct") + + data <- InMemoryDataset$create(Table$create(schema = schema( + list_col = list_of(int32()), + struct_col = struct(value = int32()), + map_col = map_of(utf8(), int32()), + fixed_list_col = fixed_size_list_of(int32(), 2L), + int_col = int32() + ))) + state <- sess:::dataview_to_state(data) + + for (position in 2:5) { + expect_false(state$columns[[position]]$filter) + expect_false(state$columns[[position]]$sortable) + } + expect_true(state$columns[[6L]]$sortable) + }) + + local({ + Table <- getExportedValue("arrow", "Table") + StructArray <- getExportedValue("arrow", "StructArray") + InMemoryDataset <- getExportedValue("arrow", "InMemoryDataset") + + data <- InMemoryDataset$create(Table$create( + id = 1:4, + struct_col = StructArray$create(data.frame( + value = 11:14, + label = c("a", "b", "c", "d") + )) + )) + state <- sess:::dataview_to_state(data) + page <- sess:::dataview_rows(state, 1:4, use_arrow_reader = TRUE) + + expect_true(is.list(page[["2"]])) + expect_false(is.data.frame(page[["2"]])) + expect_equal(page[["2"]][[1L]], list(value = 11L, label = "a")) + expect_silent(jsonlite::toJSON( + page, auto_unbox = TRUE, null = "null", force = TRUE, digits = NA + )) + }) + + local({ + Array <- getExportedValue("arrow", "Array") + Table <- getExportedValue("arrow", "Table") + InMemoryDataset <- getExportedValue("arrow", "InMemoryDataset") + int32 <- getExportedValue("arrow", "int32") + list_of <- getExportedValue("arrow", "list_of") + + data <- InMemoryDataset$create(Table$create( + id = 1:1200, + list_col = Array$create( + rep(list(c(1L, 2L), c(3L, 4L)), 600L), + type = list_of(int32()) + ) + )) + state <- sess:::dataview_to_state(data) + sess:::dataview_rows(state, 501:1000, use_arrow_reader = TRUE) + + state$query_indices <- 1:100 + page <- sess:::dataview_arrow_query_slice(state, 1:100) + + expect_true(is.list(page[["list_col"]])) + expect_false(inherits(page[["list_col"]], "vctrs_list_of")) + expect_silent(jsonlite::toJSON( + page, auto_unbox = TRUE, null = "null", force = TRUE, digits = NA + )) + state$arrow_reader$reader$Close() + }) + + local({ + sess_env <- sess:::.sess_env + original_views <- sess_env$dataviews + on.exit(sess_env$dataviews <- original_views, add = TRUE) + + InMemoryDataset <- getExportedValue("arrow", "InMemoryDataset") + data <- InMemoryDataset$create(data.frame( + id = 1:12, + logical = rep(c(TRUE, FALSE, NA), 4L) + )) + view_id <- sess:::dataview_register(data)$view_id + + true_page <- sess:::handle_dataview_page(list( + view_id = view_id, startRow = 0L, endRow = 20L, + filterModel = list("2" = list(type = "true")) + )) + expect_true(all(true_page$rows[["2"]] %in% TRUE)) + + false_page <- sess:::handle_dataview_page(list( + view_id = view_id, startRow = 0L, endRow = 20L, + filterModel = list("2" = list(type = "false")) + )) + expect_true(all(false_page$rows[["2"]] %in% FALSE)) + + true_again <- sess:::handle_dataview_page(list( + view_id = view_id, startRow = 0L, endRow = 20L, + filterModel = list("2" = list(type = "true")) + )) + expect_true(all(true_again$rows[["2"]] %in% TRUE)) + }) + + local({ + sess_env <- sess:::.sess_env + original_views <- sess_env$dataviews + path <- tempfile("dataview-arrow-") + dir.create(path) + view_ids <- character() + on.exit({ + unlink(path, recursive = TRUE) + for (view_id in view_ids) { + sess:::handle_dataview_dispose(list( + view_id = view_id, + state_generation = sess:::dataview_get_state(view_id)$state_generation + )) + } + sess_env$dataviews <- original_views + }, add = TRUE) + + n <- 300L + df <- data.frame( + id = seq_len(n), + logical_col = rep(c(TRUE, FALSE, NA), length.out = n), + date_col = as.Date("2015-01-01") + rep(0:4, length.out = n), + datetime_col = as.POSIXct("2015-01-01", tz = "Australia/Sydney") + + rep(0:4, length.out = n), + difftime_col = as.difftime(rep(0:4, length.out = n), units = "secs") + ) + if (requireNamespace("bit64", quietly = TRUE)) { + df$int64_col <- bit64::as.integer64("9007199254740993") + + bit64::as.integer64(seq_len(n)) + } + for (i in seq_len(10L)) { + rows <- seq.int((i - 1L) * 30L + 1L, i * 30L) + arrow::write_parquet( + df[rows, , drop = FALSE], + file.path(path, sprintf("part-%02d.parquet", i)) + ) + } + + data <- arrow::open_dataset(path, format = "parquet") + view_id <- sess:::dataview_register(data)$view_id + expected_id <- sess:::dataview_register(df)$view_id + view_ids <- c(view_id, expected_id) + + state <- sess:::dataview_get_state(view_id) + expect_equal(as.character(state$columns[[4L]]$type), "dateColumn") + expect_equal(as.character(state$columns[[5L]]$type), "datetimeColumn") + expect_equal(as.character(state$columns[[6L]]$type), "textColumn") + if ("int64_col" %in% names(df)) { + expect_equal(as.character(state$columns[[7L]]$type), "bigintColumn") + expect_true(inherits( + sess:::dataview_arrow_column(data, match("int64_col", names(data))), + "integer64" + )) + } + expect_true(inherits( + sess:::dataview_arrow_column(data, match("date_col", names(data))), + "Date" + )) + + compare_page <- function(params) { + actual <- sess:::handle_dataview_page(c(list(view_id = view_id), params)) + expected <- sess:::handle_dataview_page(c(list(view_id = expected_id), params)) + expect_equal(actual, expected) + expect_silent(jsonlite::toJSON( + actual, auto_unbox = TRUE, null = "null", force = TRUE, digits = NA + )) + actual + } + + true_page <- compare_page(list( + startRow = 0L, endRow = n, + filterModel = list("2" = list(type = "true")) + )) + expect_true(all(true_page$rows[["2"]] %in% TRUE)) + + target_date <- as.Date("2015-01-03") + date_page <- compare_page(list( + startRow = 0L, endRow = n, + filterModel = list("3" = list( + type = "equals", dateFrom = as.character(target_date) + )) + )) + expect_equal(date_page$totalRows, sum(df$date_col == target_date)) + + true_again <- compare_page(list( + startRow = 0L, endRow = n, + filterModel = list("2" = list(type = "true")) + )) + expect_true(all(true_again$rows[["2"]] %in% TRUE)) + + sorted_page <- compare_page(list( + startRow = 0L, endRow = n, + sortModel = list(list(colId = "1", sort = "desc")) + )) + expect_equal(sorted_page$rows[["1"]], rev(df$id)) + + reader_state <- sess:::dataview_get_state(view_id)$arrow_reader + expect_null(reader_state$reader) + expect_equal(length(reader_state$fragment_index$datasets), 10L) + expect_equal(reader_state$fragment_index$ends, seq(30, n, by = 30)) + }) + + local({ + InMemoryDataset <- getExportedValue("arrow", "InMemoryDataset") + data <- InMemoryDataset$create(data.frame(id = 1:1200, value = sprintf("v%04d", 1:1200))) + state <- sess:::dataview_to_state(data) + + page1 <- sess:::dataview_rows(state, 1:500, use_arrow_reader = TRUE) + reader <- state$arrow_reader$reader + page2 <- sess:::dataview_rows(state, 501:1000, use_arrow_reader = TRUE) + + expect_identical(state$arrow_reader$reader, reader) + expect_equal(state$arrow_reader$next_row, 1001L) + expect_equal(page1[["1"]], 1:500) + expect_equal(page2[["1"]], 501:1000) + + cached <- sess:::dataview_rows(state, 101:200, use_arrow_reader = TRUE) + expect_identical(state$arrow_reader$reader, reader) + expect_equal(state$arrow_reader$next_row, 1001L) + expect_equal(cached[["1"]], 101:200) + + sess:::dataview_arrow_reader_reset(state) + expect_null(state$arrow_reader$reader) + state$arrow_reader$row_cache <- list() + filtered_idx <- seq.int(2L, 1000L, by = 2L) + filtered <- sess:::dataview_rows(state, filtered_idx, use_arrow_reader = TRUE) + expect_equal(state$arrow_reader$next_row, 1001L) + expect_equal(filtered[["1"]], filtered_idx) + + filtered_reader <- state$arrow_reader$reader + filtered_next <- sess:::dataview_rows( + state, seq.int(1002L, 1200L, by = 2L), use_arrow_reader = TRUE + ) + expect_identical(state$arrow_reader$reader, filtered_reader) + expect_equal(state$arrow_reader$next_row, 1201L) + expect_equal(filtered_next[["1"]], seq.int(1002L, 1200L, by = 2L)) + + state$arrow_reader$reader$Close() + }) + + local({ + InMemoryDataset <- getExportedValue("arrow", "InMemoryDataset") + data <- InMemoryDataset$create(data.frame(id = 1:25000)) + state <- sess:::dataview_to_state(data) + + sess:::dataview_rows(state, 1:500, use_arrow_reader = TRUE) + sess:::dataview_rows(state, 5001:5500, use_arrow_reader = TRUE) + sess:::dataview_rows(state, 10001:10500, use_arrow_reader = TRUE) + sess:::dataview_rows(state, 15001:15500, use_arrow_reader = TRUE) + + expect_equal( + vapply(state$arrow_reader$row_cache, function(x) x$first_row, integer(1)), + c(1L, 5001L, 10001L, 15001L) + ) + + reader <- state$arrow_reader$reader + next_row <- state$arrow_reader$next_row + sess:::dataview_rows(state, 1:500, use_arrow_reader = TRUE) + expect_identical(state$arrow_reader$reader, reader) + expect_equal(state$arrow_reader$next_row, next_row) + + sess:::dataview_rows(state, 20001:20500, use_arrow_reader = TRUE) + expect_equal( + vapply(state$arrow_reader$row_cache, function(x) x$first_row, integer(1)), + c(5001L, 10001L, 15001L, 1L, 20001L) + ) + + reader <- state$arrow_reader$reader + next_row <- state$arrow_reader$next_row + sess:::dataview_rows(state, 1:500, use_arrow_reader = TRUE) + expect_identical(state$arrow_reader$reader, reader) + expect_equal(state$arrow_reader$next_row, next_row) + + sess:::dataview_rows(state, 5001:5500, use_arrow_reader = TRUE) + expect_identical(state$arrow_reader$reader, reader) + expect_equal(state$arrow_reader$next_row, next_row) + + state$arrow_reader$reader$Close() + }) + + # Filtered and sorted display blocks use a query-local display cache and + # preserve row order across page boundaries and model changes. + local({ + sess_env <- sess:::.sess_env + original_views <- sess_env$dataviews + view_ids <- character() + on.exit({ + for (view_id in view_ids) { + sess:::handle_dataview_dispose(list( + view_id = view_id, + state_generation = sess:::dataview_get_state(view_id)$state_generation + )) + } + sess_env$dataviews <- original_views + }, add = TRUE) + df <- data.frame( + id = seq_len(50000L), + score = (seq_len(50000L) * 7919L) %% 50000L, + group = rep(c("a", "b"), 25000L), + rank = ordered(rep(c("low", "high", NA, "low"), 12500L), levels = c("low", "high")) + ) + data <- getExportedValue("arrow", "InMemoryDataset")$create(df) + view_id <- sess:::dataview_register(data)$view_id + expected_id <- sess:::dataview_register(df)$view_id + view_ids <- c(view_id, expected_id) + compare_page <- function(start, model = list(), size = 500L) { + params <- c(list(startRow = start, endRow = start + size), model) + actual <- sess:::handle_dataview_page(c(list(view_id = view_id), params)) + expected <- sess:::handle_dataview_page(c(list(view_id = expected_id), params)) + expect_equal(actual, expected) + reader_state <- sess:::dataview_get_state(view_id)$arrow_reader + expect_true(sum(vapply( + reader_state$query_cache, + function(x) nrow(x$data), + integer(1) + )) <= 20000L) + } + filter <- list(filterModel = list("3" = list(type = "equals", filter = "b"))) + sort <- list(sortModel = list(list(colId = "2", sort = "desc"))) + compare_page(0L, filter) + reader_state <- sess:::dataview_get_state(view_id)$arrow_reader + reader <- reader_state$reader + expect_false(is.null(reader)) + expect_equal(reader_state$next_row, 2001L) + expect_equal(length(reader_state$query_cache), 1L) + compare_page(500L, filter) + expect_identical(reader_state$reader, reader) + expect_equal(reader_state$next_row, 2001L) + expect_equal(length(reader_state$query_cache), 1L) + compare_page(5000L, filter) + expect_identical(reader_state$reader, reader) + expect_equal(reader_state$next_row, 12001L) + expect_equal(length(reader_state$query_cache), 2L) + compare_page(0L, filter) + expect_equal(reader_state$next_row, 12001L) + expect_equal(length(reader_state$query_cache), 2L) + + for (model in list(sort, c(filter, sort), list(), filter, + list(sortModel = list(list(colId = "4", sort = "asc"))))) { + for (start in c(0L, 4990L, 10000L, 20000L, 5000L, 0L)) { + compare_page(start, model) + } + } + compare_page(0L, sort, size = 21000L) + compare_page(50000L) + compare_page(0L, list(filterModel = list("1" = list(type = "lessThan", filter = 0L)))) + }) + + if (requireNamespace("dplyr", quietly = TRUE)) { + local({ + sess_env <- sess:::.sess_env + original_views <- sess_env$dataviews + view_ids <- character() + on.exit({ + for (view_id in view_ids) { + sess:::handle_dataview_dispose(list( + view_id = view_id, + state_generation = sess:::dataview_get_state(view_id)$state_generation + )) + } + sess_env$dataviews <- original_views + }, add = TRUE) + data <- getExportedValue("arrow", "InMemoryDataset")$create(data.frame( + id = seq_len(12020L), group = rep(c("a", "b"), 6010L) + )) + query <- dplyr::arrange(data, dplyr::desc(id)) + score_data <- getExportedValue("arrow", "InMemoryDataset")$create(data.frame( + id = seq_len(12020L), + score = (seq_len(12020L) * 7919L) %% 12020L + )) + score_query <- dplyr::arrange(score_data, dplyr::desc(score)) + score_view_id <- sess:::dataview_register(score_query)$view_id + score_expected_id <- sess:::dataview_register( + sess:::dataview_arrow_data_frame(score_query) + )$view_id + view_ids <- c(view_ids, score_view_id, score_expected_id) + # Viewer operations apply to the original result, including its source order. + for (direction in c("asc", "desc")) { + params <- list( + startRow = 0L, endRow = 30L, + sortModel = list(list(colId = "1", sort = direction)) + ) + actual <- sess:::handle_dataview_page(c(list(view_id = score_view_id), params)) + expected <- sess:::handle_dataview_page(c(list(view_id = score_expected_id), params)) + expect_equal(actual$rows, expected$rows) + } + grouped <- dplyr::group_by(query, group) + aggregated <- dplyr::arrange( + dplyr::summarise(dplyr::group_by(data, group), total = sum(id)), group + ) + for (data in list(query, grouped, aggregated)) { + view_id <- sess:::dataview_register(data)$view_id + expected_df <- sess:::dataview_arrow_data_frame(data) + expected_id <- sess:::dataview_register(expected_df)$view_id + view_ids <- c(view_ids, view_id, expected_id) + expect_true(sess:::handle_dataview_init(list(view_id = view_id))$columnProjection) + for (model in list( + list(), + list(sortModel = list(list(colId = "1", sort = "asc"))), + list(filterModel = list("1" = list(type = "equals", filter = expected_df[[1L]][1L]))), + list( + filterModel = list("2" = list(type = "equals", filter = expected_df[[2L]][1L])), + sortModel = list(list(colId = "1", sort = "asc")) + ) + )) { + for (start in c(0L, 5000L, 4990L, 0L, 12020L)) { + params <- c(list(startRow = start, endRow = start + 20L), model) + actual <- sess:::handle_dataview_page(c(list(view_id = view_id), params)) + expected <- sess:::handle_dataview_page(c(list(view_id = expected_id), params)) + expect_equal(actual$rows, expected$rows) + expect_true(actual$lastRow %in% c(-1L, expected$totalRows)) + } + } + + params <- list(startRow = 0L, endRow = 20L, fields = "2") + actual <- sess:::handle_dataview_page(c(list(view_id = view_id), params)) + expected <- sess:::handle_dataview_page(c(list(view_id = expected_id), params)) + expected$rows <- expected$rows[, c("0", "2"), drop = FALSE] + expect_equal(actual$rows, expected$rows) + expect_true(actual$lastRow %in% c(-1L, nrow(expected_df))) + expect_equal(sess:::dataview_get_state(view_id)$arrow_reader$projection, 2L) + } + }) + } +} # aggregate() can return a matrix per column. Its cells are tuples, not scalars. local({ .sess_env <- sess:::.sess_env diff --git a/src/dataViewer.ts b/src/dataViewer.ts index 525e9a811..5234aed38 100644 --- a/src/dataViewer.ts +++ b/src/dataViewer.ts @@ -93,6 +93,9 @@ export function getDataViewerScript(): string { const saved = vscode.getState?.(); viewerInitialState = saved?.schema === viewerSchema ? saved.gridState : undefined; viewerFloatingFilters = saved?.schema === viewerSchema && saved.floatingFilters === true; + if (!columns.some(column => column.field !== '0' && column.filter !== false)) { + viewerFloatingFilters = false; + } viewerSizingMode = saved?.schema === viewerSchema && saved.sizingMode === 'content' ? 'content' : 'fit'; if (viewerInitialState && !saved.sizingMode) { // Older views saved fixed content widths. Let the new default flex @@ -110,12 +113,16 @@ export function getDataViewerScript(): string { }; if (column.type === 'dateColumn' || column.type === 'datetimeColumn') { column.cellDataType = column.type === 'dateColumn' ? 'dateString' : 'dateTimeString'; - column.filter = 'agDateColumnFilter'; - column.filterParams = { browserDatePicker: true }; + if (column.filter !== false) { + column.filter = 'agDateColumnFilter'; + column.filterParams = { browserDatePicker: true }; + } column.initialWidth = 200; } else if (column.type === 'bigintColumn') { column.cellDataType = 'bigint'; - column.filter = 'agBigIntColumnFilter'; + if (column.filter !== false) { + column.filter = 'agBigIntColumnFilter'; + } bigintFields.push(column.field); } else if (column.type === 'numericColumn') { column.cellDataType = 'number'; @@ -209,8 +216,10 @@ export function getDataViewerScript(): string { function updateViewerRowCount(filtered, total) { const format = value => new Intl.NumberFormat().format(value); - document.querySelector('#viewerRowCount').textContent = filtered === total - ? format(total) + ' rows' : format(filtered) + ' of ' + format(total) + ' rows'; + document.querySelector('#viewerRowCount').textContent = !Number.isFinite(total) + ? format(filtered) + '+ rows' + : filtered === total + ? format(total) + ' rows' : format(filtered) + ' of ' + format(total) + ' rows'; } function updateViewerSizingButtons() { @@ -236,12 +245,16 @@ export function getDataViewerScript(): string { function initializeViewerToolbar() { const filters = document.querySelector('#viewerFilters'); const clear = document.querySelector('#viewerClearFilters'); + const filterable = gridApi.getColumns().some(column => + column.getColId() !== '0' && column.getColDef().filter !== false + ); + filters.disabled = !filterable; const updateFilters = () => { - filters.setAttribute('aria-pressed', String(viewerFloatingFilters)); - clear.disabled = Object.keys(gridApi.getFilterModel()).length === 0; + filters.setAttribute('aria-pressed', String(filterable && viewerFloatingFilters)); + clear.disabled = !filterable || Object.keys(gridApi.getFilterModel()).length === 0; }; const setFloatingFilters = visible => { - viewerFloatingFilters = visible; + viewerFloatingFilters = filterable && visible; gridApi.setGridOption('defaultColDef', { ...gridApi.getGridOption('defaultColDef'), floatingFilter: visible }); diff --git a/src/session.ts b/src/session.ts index 0305f652c..37867d034 100644 --- a/src/session.ts +++ b/src/session.ts @@ -96,6 +96,7 @@ interface IpcSocket extends net.Socket { } export class Session { + public busyDataView: string | undefined; private exited = false; public get processExited(): boolean { return this.exited; @@ -686,13 +687,16 @@ interface DataViewColumnDef { interface DataViewInitResult { columns: DataViewColumnDef[]; - totalRows: number; + totalRows?: number; + columnProjection?: boolean; + live?: boolean; } interface DataViewPageResult { rows: Record[]; - totalRows: number; - totalUnfiltered: number; + totalRows?: number; + totalUnfiltered?: number; + fetchedRows?: number; lastRow: number; } @@ -705,6 +709,7 @@ interface DataViewRequestMessage { endRow?: number; sortModel?: unknown[]; filterModel?: Record; + fields?: string[]; } const dynamicDataViewPanels = new Map(); @@ -744,6 +749,10 @@ function registerDataViewPanel( return; } dynamicDataViewPanels.delete(key); + const owner = sessions.get(sessionId ?? ''); + if (owner) { + void interruptDataView(owner, viewId); + } // Interactive transcripts retain this handle after the expanded viewer closes. if (currentStateGeneration === undefined && sessions.get(sessionId ?? '')?.requester) { return; @@ -764,6 +773,11 @@ function attachDynamicDataViewBridge( sessionId: string | null, ): void { const webview = panel.webview; + let disposed = false; + let pageQueue = Promise.resolve(); + panel.onDidDispose(() => { + disposed = true; + }); const postResponse = ( documentGeneration: number, requestId: number, @@ -771,7 +785,7 @@ function attachDynamicDataViewBridge( result?: unknown, error?: string, ) => { - if (documentGeneration !== documentGenerations.get(webview)) { + if (disposed || documentGeneration !== documentGenerations.get(webview)) { return; } void webview.postMessage({ @@ -794,6 +808,9 @@ function attachDynamicDataViewBridge( ) { return; } + if (disposed) { + return; + } try { if (msg.action === 'init') { @@ -804,10 +821,13 @@ function attachDynamicDataViewBridge( }, sessionId, )) as DataViewInitResult | undefined; + if (disposed) { + return; + } if ( !result || !Array.isArray(result.columns) || - typeof result.totalRows !== 'number' + (result.totalRows !== undefined && typeof result.totalRows !== 'number') ) { throw new Error('Invalid dataview_init response'); } @@ -816,28 +836,55 @@ function attachDynamicDataViewBridge( } if (msg.action === 'page') { - const result = (await sessionRequest( - { - method: 'dataview_page', - params: { - view_id: viewId, - startRow: Number(msg.startRow ?? 0), - endRow: Number(msg.endRow ?? 0), - sortModel: Array.isArray(msg.sortModel) ? msg.sortModel : [], - filterModel: msg.filterModel ?? {}, + // Send one page at a time, so closing can discard queued pages + // without leaving additional expensive reads in R's IPC queue. + const previous = pageQueue; + let release!: () => void; + pageQueue = new Promise((resolve) => { + release = resolve; + }); + try { + await previous; + if (disposed || msg.documentGeneration !== documentGenerations.get(webview)) { + return; + } + const result = (await sessionRequest( + { + method: 'dataview_page', + params: { + view_id: viewId, + startRow: Number(msg.startRow ?? 0), + endRow: Number(msg.endRow ?? 0), + sortModel: Array.isArray(msg.sortModel) ? msg.sortModel : [], + filterModel: msg.filterModel ?? {}, + fields: Array.isArray(msg.fields) ? msg.fields : undefined, + notify_busy: true, + }, }, - }, - sessionId, - )) as DataViewPageResult | undefined; - if ( - !result || - !Array.isArray(result.rows) || - typeof result.totalRows !== 'number' || - typeof result.totalUnfiltered !== 'number' - ) { - throw new Error('Invalid dataview_page response'); + sessionId, + 0, + )) as DataViewPageResult | undefined; + if ( + !result || + !Array.isArray(result.rows) || + typeof result.lastRow !== 'number' || + (result.totalRows !== undefined && typeof result.totalRows !== 'number') || + (result.lastRow < 0 && + (typeof result.fetchedRows !== 'number' || + !Number.isSafeInteger(result.fetchedRows) || + result.fetchedRows < 0)) || + (result.totalUnfiltered !== undefined && + typeof result.totalUnfiltered !== 'number') + ) { + throw new Error('Invalid dataview_page response'); + } + if (disposed || msg.documentGeneration !== documentGenerations.get(webview)) { + return; + } + postResponse(msg.documentGeneration, msg.requestId, true, result); + } finally { + release(); } - postResponse(msg.documentGeneration, msg.requestId, true, result); return; } @@ -860,6 +907,26 @@ function attachDynamicDataViewBridge( }); } +async function interruptDataView(session: Session, viewId: string): Promise { + const isReading = () => session.busyDataView === viewId && !session.socket.destroyed; + if (!isReading()) { + return; + } + for (const terminal of window.terminals) { + const terminalPid = await terminal.processId; + if (!isReading()) { + return; + } + if ( + terminalPid !== undefined && + (terminalRegistry.ownerOf(terminal) === session || String(terminalPid) === session.pid) + ) { + terminal.sendText('\x03', false); + return; + } + } +} + export function deploySessionWatcher(extensionPath: string): void { if (attachmentLifecycleStopping) { // A new activation must not inherit endpoints being torn down by the prior host lifecycle. @@ -2807,6 +2874,7 @@ export async function getTableHtml( let longFetchTimer; let filteredRows = 0; let totalRows = 0; + let rowCountKnown = false; let isFiltered = false; let verticalScrollbar; let scrollbarPositionAttached = false; @@ -2900,7 +2968,7 @@ export async function getTableHtml( const positionEl = document.querySelector('#scrollPosition'); const containerEl = document.querySelector('#gridContainer'); - if (!positionEl || !containerEl || filteredRows < 1) { + if (!positionEl || !containerEl) { return; } @@ -2910,26 +2978,33 @@ export async function getTableHtml( ? Math.max(0, Math.min(1, verticalScrollbar.scrollTop / maximumScroll)) : 0; - let pageStart = 0; - let rowsInView = filteredRows; - if (${pageSize > 0 ? 'true' : 'false'} && - typeof gridApi.paginationGetCurrentPage === 'function' && - typeof gridApi.paginationGetPageSize === 'function') { - pageStart = - gridApi.paginationGetCurrentPage() * gridApi.paginationGetPageSize(); - pageStart = Math.min(pageStart, Math.max(0, filteredRows - 1)); - rowsInView = Math.min( - gridApi.paginationGetPageSize(), - filteredRows - pageStart + if (rowCountKnown && filteredRows > 0) { + let pageStart = 0; + let rowsInView = filteredRows; + if (${pageSize > 0 ? 'true' : 'false'} && + typeof gridApi.paginationGetCurrentPage === 'function' && + typeof gridApi.paginationGetPageSize === 'function') { + pageStart = + gridApi.paginationGetCurrentPage() * gridApi.paginationGetPageSize(); + pageStart = Math.min(pageStart, Math.max(0, filteredRows - 1)); + rowsInView = Math.min( + gridApi.paginationGetPageSize(), + filteredRows - pageStart + ); + } + const currentRow = Math.min( + filteredRows, + pageStart + Math.round(scrollRatio * Math.max(0, rowsInView - 1)) + 1 ); + positionEl.textContent = + rowNumberFormatter.format(currentRow) + + ' of ' + rowNumberFormatter.format(filteredRows); + } else { + const firstDisplayed = typeof gridApi.getFirstDisplayedRowIndex === 'function' + ? gridApi.getFirstDisplayedRowIndex() : 0; + positionEl.textContent = + rowNumberFormatter.format(Math.max(1, firstDisplayed + 1)); } - const currentRow = Math.min( - filteredRows, - pageStart + Math.round(scrollRatio * Math.max(0, rowsInView - 1)) + 1 - ); - positionEl.textContent = - rowNumberFormatter.format(currentRow) + - ' of ' + rowNumberFormatter.format(filteredRows); const scrollbarRect = verticalScrollbar.getBoundingClientRect(); const containerRect = containerEl.getBoundingClientRect(); @@ -3030,8 +3105,11 @@ export async function getTableHtml( } const columns = Array.isArray(init.columns) ? init.columns : []; - filteredRows = init.totalRows; - totalRows = init.totalRows; + rowCountKnown = Number.isFinite(init.totalRows); + if (rowCountKnown) { + filteredRows = init.totalRows; + totalRows = init.totalRows; + } if (init.live) { const info = document.createElement('span'); info.textContent = 'Full data'; @@ -3039,7 +3117,9 @@ export async function getTableHtml( document.querySelector('#viewerControls').append(info); } const bigintFields = prepareViewerColumns(columns); - updateViewerRowCount(filteredRows, totalRows); + if (rowCountKnown) { + updateViewerRowCount(filteredRows, totalRows); + } const rowIndexColumn = columns.find(column => column.field === '0'); if (rowIndexColumn) { rowIndexColumn.headerValueGetter = () => isFiltered @@ -3048,9 +3128,41 @@ export async function getTableHtml( } const blockSize = ${pageSize > 0 ? pageSize : 500}; + let cachedFieldsKey; + let projectionRefreshTimer; + function pageFields() { + if (!init.columnProjection) { + return undefined; + } + const displayed = gridApi?.getAllDisplayedColumns(); + return (displayed + ? displayed.map(column => column.getColId()) + : columns.filter(column => !column.hide).map(column => column.field) + ).sort(); + } + function displayedColumnsChanged() { + updateFetchStatusPosition(); + if (!init.columnProjection) { + return; + } + clearTimeout(projectionRefreshTimer); + // Column events can fire during grid setup or a batch of visibility + // changes. Refresh once, after the column model has settled. + projectionRefreshTimer = setTimeout(() => { + const key = JSON.stringify(pageFields()); + if (gridApi && cachedFieldsKey !== undefined && key !== cachedFieldsKey) { + cachedFieldsKey = key; + gridApi.purgeInfiniteCache(); + } + }, 0); + } const datasource = { getRows: async function(params) { + const fields = pageFields(); + if (cachedFieldsKey === undefined) { + cachedFieldsKey = JSON.stringify(fields); + } beginFetch('Fetching rows from R session...'); try { const result = await request('page', { @@ -3058,16 +3170,33 @@ export async function getTableHtml( endRow: params.endRow, sortModel: params.sortModel, filterModel: params.filterModel, + fields, }); - filteredRows = result.totalRows; - totalRows = result.totalUnfiltered; - isFiltered = Object.keys(params.filterModel || {}).length > 0; + if (JSON.stringify(fields) !== JSON.stringify(pageFields())) { + params.failCallback(); + finishFetch(true); + return; + } + if (Number.isFinite(result.totalRows)) { + filteredRows = result.totalRows; + totalRows = Number.isFinite(result.totalUnfiltered) + ? result.totalUnfiltered : result.totalRows; + rowCountKnown = true; + } else if (Number.isFinite(result.lastRow) && result.lastRow >= 0) { + filteredRows = result.lastRow; + totalRows = result.lastRow; + rowCountKnown = true; + } else { + filteredRows = result.fetchedRows; + totalRows = undefined; + rowCountKnown = false; + } + isFiltered = rowCountKnown && Object.keys(params.filterModel || {}).length > 0; updateViewerRowCount(filteredRows, totalRows); gridApi?.refreshHeader(); updateScrollPosition(); - const resolvedLastRow = Number.isFinite(result.totalRows) ? result.totalRows : result.lastRow; const rows = prepareViewerRows(result.rows || [], bigintFields); - params.successCallback(rows, resolvedLastRow); + params.successCallback(rows, result.lastRow); finishFetch(true); } catch (e) { console.error('[dataview] Failed to load page', e); @@ -3086,7 +3215,7 @@ export async function getTableHtml( cacheBlockSize: blockSize, onPaginationChanged: updateScrollPosition, onGridSizeChanged: updateFetchStatusPosition, - onDisplayedColumnsChanged: updateFetchStatusPosition, + onDisplayedColumnsChanged: displayedColumnsChanged, onFirstDataRendered: function() { updateFetchStatusPosition(); attachScrollbarPositionIndicator(); @@ -3581,6 +3710,23 @@ async function handleNotification(message: Record, socket: IpcS const params = (message.params as Record) || {}; switch (method) { + case 'dataview_busy': { + const session = sessions.get(socket._sessionId ?? ''); + if (session) { + session.busyDataView = + typeof params.view_id === 'string' ? params.view_id : undefined; + // A page may begin after its panel was closed while R was busy. + if ( + session.busyDataView && + !dynamicDataViewPanels.has( + JSON.stringify([session.sessionId, session.busyDataView]), + ) + ) { + void interruptDataView(session, session.busyDataView); + } + } + break; + } case 'attach': { if (!isLiveSessionSocket(socket)) { socket.destroy(); @@ -4089,6 +4235,7 @@ export async function cleanupSession(sessionId: string, closingSocket?: IpcSocke export async function sessionRequest( data: Record, target: Session | string | null | undefined = activeSession, + timeout = 5000, ): Promise { try { const owner = typeof target === 'string' ? sessions.get(target) : target; @@ -4121,12 +4268,14 @@ export async function sessionRequest( reject(e); } - setTimeout(() => { - if (pendingRequests.has(id)) { - pendingRequests.delete(id); - reject(new Error('Request timed out')); - } - }, 5000); + if (timeout > 0) { + setTimeout(() => { + if (pendingRequests.has(id)) { + pendingRequests.delete(id); + reject(new Error('Request timed out')); + } + }, timeout); + } }); } catch (error) { if (error instanceof Error) { diff --git a/src/test/suite/dataViewerSessions.test.ts b/src/test/suite/dataViewerSessions.test.ts index 729265cd6..cdbfab223 100644 --- a/src/test/suite/dataViewerSessions.test.ts +++ b/src/test/suite/dataViewerSessions.test.ts @@ -5,6 +5,8 @@ import * as os from 'os'; import * as path from 'path'; import * as sinon from 'sinon'; import * as vscode from 'vscode'; +import * as vm from 'vm'; +import type { GridApi, GridOptions, IDatasource } from 'ag-grid-community'; import * as extension from '../../extension'; import * as session from '../../session'; import * as util from '../../util'; @@ -31,7 +33,13 @@ interface Client { interface Panel { panel: vscode.WebviewPanel; receive: (message: unknown) => Promise; - replies: Array<{ ok?: boolean; error?: string; message?: string; text?: string }>; + replies: Array<{ + ok?: boolean; + error?: string; + message?: string; + text?: string; + result?: unknown; + }>; } async function waitFor(condition: () => boolean): Promise { @@ -182,7 +190,12 @@ suite('Viewer session ownership', () => { result = { columns: [], totalRows: 1 }; break; case 'dataview_page': - result = { rows: [{ owner: id }], totalRows: 1, totalUnfiltered: 1 }; + result = { + rows: [{ owner: id }], + totalRows: 1, + totalUnfiltered: 1, + lastRow: 1, + }; break; case 'workspace_children': result = { children: [], next_start: null }; @@ -426,6 +439,129 @@ suite('Viewer session ownership', () => { const viewerRequests = (client: Client) => client.requests.filter((request) => request.method !== 'workspace'); + test('lazy row counts use fetched rows until EOF, while known counts stay exact', async () => { + const webview = { asWebviewUri: (uri: vscode.Uri) => uri } as vscode.Webview; + const html = await session.getTableHtml(webview, undefined, 'lazy counts'); + const script = /