Support pre-filtering RagnarStores in ragnar_retrieve() (#29)

* use new duckdb array->matrix support

* Add support for tbl_sql in ragnar_retrieve functions

* accept store-tbl in `ragnar_retrieve()`; support pre-filtering

* redocument

* `R CMD check` fixes

* use mock `embed` function in example

Author: t-kalinowski

DESCRIPTION

@@ -15,9 +15,11 @@ License: MIT + file LICENSE
 Encoding: UTF-8
 Roxygen: list(markdown = TRUE)
 RoxygenNote: 7.3.2
+Depends:
+    R (>= 4.3.0)
 Imports:
     DBI,
-    duckdb,
+    duckdb (>= 1.2.2),
     glue,
     rlang (>= 1.1.0),
     dplyr,
@@ -44,6 +46,7 @@ Suggests:
     readr,
     rmarkdown,
     stringr,
+    dbplyr,
     testthat (>= 3.0.0),
     paws.common,
     shiny

NAMESPACE

@@ -34,7 +34,14 @@ importFrom(DBI,dbQuoteString)
 importFrom(DBI,dbReadTable)
 importFrom(DBI,dbWriteTable)
 importFrom(dotty,.)
+importFrom(dplyr,arrange)
 importFrom(dplyr,bind_rows)
+importFrom(dplyr,collect)
+importFrom(dplyr,filter)
+importFrom(dplyr,mutate)
+importFrom(dplyr,select)
+importFrom(dplyr,sql)
+importFrom(dplyr,tbl)
 importFrom(glue,as_glue)
 importFrom(glue,glue)
 importFrom(glue,glue_data)
@@ -80,6 +87,7 @@ importFrom(stringi,stri_trim_both)
 importFrom(tibble,as_tibble)
 importFrom(tibble,tibble)
 importFrom(tidyr,unchop)
+importFrom(utils,head)
 importFrom(vctrs,data_frame)
 importFrom(vctrs,list_unchop)
 importFrom(vctrs,new_data_frame)

R/embed-bedrock.R

@@ -6,6 +6,7 @@
 #'   There are no guardarails for the kind of model that is used, but the model
 #'   must be available in the AWS region specified by the profile.
 #'   You may look for available models in the Bedrock Model Catalog
+#' @param profile AWS profile to use.
 #' @param api_args Additional arguments to pass to the Bedrock API. Dependending
 #'   on the `model`, you might be able to provide different parameters. Check
 #'   the documentation for the model you are using in the

R/ragnar-package.R

@@ -16,6 +16,7 @@ dotty::.
 
 .onLoad <- function(libname, pkgname) {
   Sys.setenv(RETICULATE_PYTHON = "managed")
+  S7::methods_register()
   reticulate::py_require(c(
     "markitdown[all]"
   ))

R/retrieve.R

@@ -38,6 +38,10 @@ ragnar_retrieve_vss <- function(
   check_string(text)
   check_number_whole(top_k)
   method <- rlang::arg_match(method)
+  if (inherits(store, "tbl_sql")) {
+    tbl <- store
+    return(ragnar_retrieve_vss_tbl(tbl, text, top_k, method))
+  }
 
   cols <- names(store@schema) |>
     stringi::stri_subset_regex("^embedding$", negate = TRUE) |>
@@ -66,22 +70,22 @@ ragnar_retrieve_vss <- function(
 # store |> dplyr::mutate(score = calculate_vss(store, text))
 # using dbplyr
 calculate_vss <- function(store, text, method) {
-  if (is.null(store@embed)) {
+  embed <- get_store_embed(store)
+  if (is.null(embed)) {
     cli::cli_abort("Store must have an embed function but got {.code NULL}")
   }
 
-  embedded_text <- store@embed(text)
+  embedded_text <- embed(text)
   embedding_size <- ncol(embedded_text)
 
   .[method_function, ..] <- method_to_info(method)
-
   glue::glue(
     r"---(
     {method_function}(
       embedding,
       [{stri_flatten(embedded_text, ", ")}]::FLOAT[{embedding_size}]
     )
-  )---"
+    )---"
   )
 }
 
@@ -94,10 +98,71 @@ method_to_info <- function(method) {
     euclidean_distance = c("array_distance", "ASC"),
     negative_dot_product = c("array_negative_dot_product", "ASC"),
     cosine_similarity = c("array_cosine_similarity", "DESC"),
-    dot_product = c("array_dot_product", "DESC")
+    dot_product = c("array_dot_product", "DESC"),
+    stop("Unknown method")
   )
 }
 
+
+get_store_embed <- function(x) {
+  if (S7_inherits(x, RagnarStore)) {
+    return(x@embed)
+  }
+
+  if (inherits(x, "tbl_sql")) {
+    con <- dbplyr::remote_con(x)
+    ptr <- con@conn_ref
+    embed <- attr(ptr, "embed_function", exact = TRUE)
+    if (!is.null(embed)) {
+      return(embed)
+    }
+
+    # Attribute missing: reread from db and cache on ptr
+    embed_blob <- DBI::dbGetQuery(
+      con,
+      "SELECT embed_func FROM metadata LIMIT 1"
+    )$embed_func[[1]]
+    embed <- unserialize(embed_blob)
+    attr(ptr, "embed_function") <- embed
+    return(embed)
+  }
+
+  cli::cli_abort("`store` must be a RagnarStore or a dplyr::tbl()")
+}
+
+
+ragnar_retrieve_vss_tbl <- function(tbl, text, top_k, method) {
+  .[.., order_key] <- method_to_info(method)
+  tbl |>
+    mutate(
+      metric_value = sql(calculate_vss(tbl, text, method)),
+      metric_name = method
+    ) |>
+    select(-"embedding") |>
+    arrange(sql(glue("metric_value {order_key}"))) |>
+    head(n = top_k) |>
+    collect()
+}
+
+ragnar_retrieve_bm25_tbl_sql <- function(tbl, text, top_k) {
+  con <- dbplyr::remote_con(tbl)
+  text_quoted <- DBI::dbQuoteString(con, text)
+
+  tbl |>
+    mutate(
+      metric_value = sql(glue::glue(
+        "fts_main_chunks.match_bm25(id, {text_quoted})"
+      )),
+      metric_name = "bm25"
+    ) |>
+    filter(sql('metric_value IS NOT NULL')) |>
+    arrange(.data$metric_value) |>
+    select(-"embedding") |>
+    head(n = top_k) |>
+    collect()
+}
+
+
 #' Retrieves chunks using the BM25 score
 #'
 #' BM25 refers to Okapi Best Matching 25. See \doi{10.1561/1500000019} for more information.
@@ -108,6 +173,9 @@ method_to_info <- function(method) {
 ragnar_retrieve_bm25 <- function(store, text, top_k = 3L) {
   check_string(text)
   check_number_whole(top_k)
+  if (inherits(store, "tbl_sql")) {
+    return(ragnar_retrieve_bm25_tbl_sql(store, text, top_k))
+  }
 
   cols <- names(store@schema) |>
     stringi::stri_subset_regex("^embedding$", negate = TRUE) |>
@@ -168,7 +236,6 @@ ragnar_retrieve_vss_and_bm25 <- function(store, text, top_k = 3, ...) {
   )
 
   # TODO: come up with a nice reordering that doesn't involve too much compute.
-
   as_tibble(out)
 }
 
@@ -178,18 +245,54 @@ ragnar_retrieve_vss_and_bm25 <- function(store, text, top_k = 3, ...) {
 #' [ragnar_retrieve()] is a thin wrapper around [ragnar_retrieve_vss_and_bm25()]
 #' using the recommended best practices.
 #'
-#' @param store A `RagnarStore` object.
+#' @param store A `RagnarStore` object or a `dplyr::tbl()` derived from
+#'   it. When you pass a `tbl`, you may use usual dplyr verbs (e.g.
+#'   `filter()`, `slice()`) to restrict the rows examined before
+#'   similarity scoring. Avoid dropping essential columns such as
+#'   `text`, `embedding`, `origin`, and `hash`.
 #' @param text A string to find the nearest match too
 #' @param top_k Integer, the number of nearest entries to find *per method*.
 #'
 #' @returns A dataframe of retrieved chunks. Each row corresponds to an
 #'   individual chunk in the store. It always contains a column named `text`
 #'   that contains the chunks.
 #'
+#' @section Pre-filtering with dplyr:
+#' The store behaves like a lazy table backed by DuckDB, so row‑wise
+#' filtering is executed directly in the database. This lets you narrow the
+#' search space efficiently without pulling data into R.
+#'
 #' @family ragnar_retrieve
 #' @export
+#' @examples
+#' # Basic usage
+#' mock_embed <- function(x) matrix(stats::runif(10), nrow = length(x), ncol = 10)
+#' store <- ragnar_store_create(embed = mock_embed)
+#' ragnar_store_insert(store, data.frame(text = c("foo", "bar")))
+#' ragnar_store_build_index(store)
+#' ragnar_retrieve(store, "foo")
+#'
+#' # More Advanced: store metadata, retrieve with pre-filtering
+#' store <- ragnar_store_create(
+#'   embed = mock_embed,
+#'   extra_cols = data.frame(category = character())
+#' )
+#' ragnar_store_insert(
+#'   store,
+#'   data.frame(
+#'     category = c("desert", "desert", "desert", "meal", "meal", "meal"),
+#'     text = c("ice cream", "cake", "cookies", "pasta", "burger", "salad")
+#'   )
+#' )
+#' ragnar_store_build_index(store)
+#'
+#' # simple retrieve
+#' ragnar_retrieve(store, "yummy")
+#'
+#' # retrieve with pre-filtering
+#' dplyr::tbl(store) |>
+#'   dplyr::filter(category == "meal") |>
+#'   ragnar_retrieve("yummy")
 ragnar_retrieve <- function(store, text, top_k = 3L) {
   ragnar_retrieve_vss_and_bm25(store, text, top_k)
 }
-
-# TODO: re-ranking.

R/store.R

@@ -69,7 +69,7 @@ ragnar_store_create <- function(
       stop("File already exists: ", location)
     }
   }
-  con <- dbConnect(duckdb::duckdb(), dbdir = location)
+  con <- dbConnect(duckdb::duckdb(), dbdir = location, array = "matrix")
 
   default_schema <- vctrs::vec_ptype(data_frame(
     origin = character(0),
@@ -152,6 +152,10 @@ ragnar_store_create <- function(
   schema <- unserialize(metadata$schema[[1]])
   name <- metadata$name
 
+  # attach function to externalptr, so we can retreive it from just the connection.
+  ptr <- con@conn_ref
+  attr(ptr, "embed_function") <- embed
+
   # duckdb R interface does not support array columns yet,
   # so we hand-write the sql.
   columns <- map2(names(schema), schema, function(nm, type) {
@@ -219,7 +223,12 @@ ragnar_store_connect <- function(
   # mode <- match.arg(mode)
   # read_only <- mode == "retrieve"
 
-  con <- dbConnect(duckdb::duckdb(), dbdir = location, read_only = read_only)
+  con <- dbConnect(
+    duckdb::duckdb(),
+    dbdir = location,
+    read_only = read_only,
+    array = "matrix"
+  )
 
   # can't use dbExistsTable() because internally it runs:
   # > dbGetQuery(conn, sqlInterpolate(conn, "SELECT * FROM ? WHERE FALSE", dbQuoteIdentifier(conn, name)))
@@ -237,6 +246,10 @@ ragnar_store_connect <- function(
   schema <- unserialize(metadata$schema[[1L]])
   name <- metadata$name %||% unique_store_name()
 
+  # attach function to externalptr, so we can retreive it from just the connection.
+  ptr <- con@conn_ref
+  attr(ptr, "embed_function") <- embed
+
   if (build_index) ragnar_store_build_index(con)
 
   DuckDBRagnarStore(embed = embed, schema = schema, .con = con, name = name)
@@ -398,13 +411,13 @@ ragnar_store_insert <- function(store, chunks) {
 
   # Ideally this would use dbWriteTable, but we can't really because it currently
   # doesn't support array columns.
-  cols <- map2(names(schema), schema, function(nm, ptype) {
+  cols <- imap(schema, function(ptype, name) {
     # Ensures that the column in chunks has the expected ptype. (or at least
     # something that can be cast to the correct ptype with no loss)
     col <- vctrs::vec_cast(
-      chunks[[nm]],
+      chunks[[name]],
       ptype,
-      x_arg = glue::glue("chunks${nm}")
+      x_arg = glue::glue("chunks${name}")
     )
 
     if (is.matrix(col) && is.numeric(col)) {
@@ -418,7 +431,7 @@ ragnar_store_insert <- function(store, chunks) {
     } else if (is.numeric(col)) {
       DBI::dbQuoteLiteral(store@.con, col)
     } else {
-      cli::cli_abort("Unsupported type {.cls {class(col)}")
+      cli::cli_abort("Unsupported type {.cls {class(col)}}")
     }
   })
 
@@ -518,3 +531,9 @@ ragnar_store_inspect <- function(store, ...) {
   })
   invisible(NULL)
 }
+
+#' @importFrom dplyr tbl sql arrange collect
+method(tbl, ragnar:::DuckDBRagnarStore) <- function(src, from = "chunks", ...) {
+  tbl(src@.con, from)
+}
+rm(tbl)

R/utils.R

@@ -10,7 +10,7 @@
 #' @importFrom xml2 xml_add_sibling xml_find_all xml_name xml_attr xml_text
 #'   xml_url url_absolute xml_contents xml_find_first
 #' @importFrom tibble tibble as_tibble
-#' @importFrom dplyr bind_rows
+#' @importFrom dplyr bind_rows select mutate filter
 #' @importFrom tidyr unchop
 #' @importFrom vctrs data_frame vec_split vec_rbind vec_cbind vec_locate_matches
 #'   vec_fill_missing vec_unique vec_slice vec_c list_unchop new_data_frame
@@ -21,6 +21,7 @@
 #'   dbWriteTable dbListTables dbReadTable
 #' @importFrom glue glue glue_data as_glue
 #' @importFrom methods is
+#' @importFrom utils head
 #' @useDynLib ragnar, .registration = TRUE
 NULL
 
@@ -57,6 +58,12 @@ map3 <- function(.x, .y, .z, .f, ...) {
   out
 }
 
+imap <- function(.x, .f, ...) {
+  out <- .mapply(.f, list(.x, names(.x) %||% seq_along(.x)), list(...))
+  names(out) <- names(.x)
+  out
+}
+
 `%""%` <- function(x, y) {
   stopifnot(is.character(x), is.character(y))
   if (length(y) != length(x)) {