Update buffer element query functions

NAMESPACE

@@ -151,6 +151,7 @@ export(tiledb_query_get_range)
 export(tiledb_query_get_range_num)
 export(tiledb_query_import_buffer)
 export(tiledb_query_result_buffer_elements)
+export(tiledb_query_result_buffer_elements_vec)
 export(tiledb_query_set_buffer)
 export(tiledb_query_set_buffer_ptr)
 export(tiledb_query_set_buffer_ptr_char)

R/Query.R

@@ -101,7 +101,8 @@ tiledb_query_set_subarray <- function(query, subarray, type) {
 #'
 #' This function allocates query buffers directly from R vectors in
 #' case the types match: `integer`, `double`, `logical`. For more
-#' general types see \code{tiledb_query_buffer_alloc_ptr}.
+#' general types see \code{tiledb_query_buffer_alloc_ptr} and
+#' \code{tiledb_query_buffer_assign_ptr}
 #' @param query A TileDB Query object
 #' @param attr A character value containing the attribute
 #' @param buffer A vector providing the query buffer
@@ -298,17 +299,56 @@ tiledb_query_status <- function(query) {
   libtiledb_query_status(query@ptr)
 }
 
-#' Get TileDB Query result buffer elements
+#' Get TileDB Query result buffer element size
+#'
+#' The underlying library functions returns a pair of values as a vector
+#' of length two. The first number is the number of element offsets for variable
+#' size attributes (and always zero for fixed-sized attributes and coordinates).
+#' The second is the number of elements in the data buffer. For variable-sized
+#' attributes the first number is the number of cells read (and hence the number
+#' of offsets), the second number is the number of elements in the data buffer.
+#'
+#' As this function was first made available when only a scalar (corresponding to
+#' the second result) was returned, we still return that value.
 #'
 #' @param query A TileDB Query object
 #' @param attr A character value containing the attribute
 #' @return A integer with the number of elements in the results buffer
 #' for the given attribute
+#' @seealso tiledb_query_result_buffer_elements_vec
 #' @export
 tiledb_query_result_buffer_elements <- function(query, attr) {
   stopifnot(query_object=is(query, "tiledb_query"),
             attr_variable=is.character(attr))
-  libtiledb_query_result_buffer_elements(query@ptr, attr)
+  libtiledb_query_result_buffer_elements(query@ptr, attr, 1) # request 2nd el in pair
+}
+
+#' Get TileDB Query result buffer element size pair as vector
+#'
+#' The underlying library functions returns a pair of values as a vector
+#' of length two. The first number is the number of element offsets for variable
+#' size attributes (and always zero for fixed-sized attributes and coordinates).
+#' The second is the number of elements in the data buffer. For variable-sized
+#' attributes the first number is the number of cells read (and hence the number
+#' of offsets), the second number is the number of elements in the data buffer.
+#' In the case of a nullable attribute, a third element is returned with the size of
+#' the validity buffer.
+#'
+#' @param query A TileDB Query object
+#' @param attr A character value containing the attribute
+#' @param nullable A logical variable that is \sQuote{TRUE} to signal that the attribute
+#' is nullable, and \sQuote{FALSE} otherwise
+#' @return A vector with the number of elements in the offsets buffer (and zero
+#' for fixed-size attribute or dimensions), the number elements in the results
+#' buffer for the given attribute, and (if nullable) a third element with the validity
+#' buffer size.
+#' @seealso tiledb_query_result_buffer_elements
+#' @export
+tiledb_query_result_buffer_elements_vec <- function(query, attr, nullable = FALSE) {
+  stopifnot(`query must be a query_object` = is(query, "tiledb_query"),
+            `attr must be a string` = is.character(attr),
+            `nullable must be a logical` = is.logical(nullable))
+  libtiledb_query_result_buffer_elements_vec(query@ptr, attr, nullable)
 }
 
 #' Set a range for a given query

R/RcppExports.R

@@ -609,6 +609,10 @@ libtiledb_query_result_buffer_elements <- function(query, attribute, which = 1L)
     .Call(`_tiledb_libtiledb_query_result_buffer_elements`, query, attribute, which)
 }
 
+libtiledb_query_result_buffer_elements_vec <- function(query, attribute, nullable = FALSE) {
+    .Call(`_tiledb_libtiledb_query_result_buffer_elements_vec`, query, attribute, nullable)
+}
+
 libtiledb_query_get_fragment_num <- function(query) {
     .Call(`_tiledb_libtiledb_query_get_fragment_num`, query)
 }
@@ -661,6 +665,14 @@ libtiledb_query_set_condition <- function(query, query_cond) {
     .Call(`_tiledb_libtiledb_query_set_condition`, query, query_cond)
 }
 
+libtiledb_query_get_array <- function(query, ctx) {
+    .Call(`_tiledb_libtiledb_query_get_array`, query, ctx)
+}
+
+libtiledb_query_get_schema <- function(query, ctx) {
+    .Call(`_tiledb_libtiledb_query_get_schema`, query, ctx)
+}
+
 libtiledb_query_condition <- function(ctx) {
     .Call(`_tiledb_libtiledb_query_condition`, ctx)
 }

R/TileDBArray.R

@@ -665,9 +665,8 @@ setMethod("[", "tiledb_array",
   ## get results
   getResult <- function(buf, name, varnum, resrv, qryptr) {
     if (is.na(varnum)) {
-      sz <- libtiledb_query_result_buffer_elements(qryptr, name, 0)
-      len <- libtiledb_query_result_buffer_elements(qryptr, name, 1)
-      libtiledb_query_get_buffer_var_char(buf, sz, len)[,1]
+      vec <- libtiledb_query_result_buffer_elements_vec(qryptr, name)
+      libtiledb_query_get_buffer_var_char(buf, vec[1], vec[2])[,1]
     } else {
       libtiledb_query_get_buffer_ptr(buf, asint64)
     }

inst/tinytest/test_query.R

@@ -156,7 +156,8 @@ dir.create(tmp)
 
 dom <- tiledb_domain(dims = tiledb_dim("rows", c(1L, 10L), 1L, type = "INT32"))
 schema <- tiledb_array_schema(dom,
-                              attrs = tiledb_attr("vals", type = "INT32"),
+                              attrs = c(tiledb_attr("vals", type = "INT32"),
+                                        tiledb_attr("keys", type = "INT32", nullable = TRUE)),
                               sparse=FALSE)
 tiledb_array_create(tmp, schema)
 arr <- tiledb_array(tmp)
@@ -169,6 +170,11 @@ qry <- tiledb_query_set_buffer(qry, "rows", rows)
 vals <- seq(101,110)
 qry <- tiledb_query_set_buffer(qry, "vals", vals)
 
+keys <- c(201:204, NA_integer_, 206L, NA_integer_, 208:210)
+buf <- tiledb:::libtiledb_query_buffer_alloc_ptr(arr@ptr, "INT32", 10, TRUE)
+buf <- tiledb:::libtiledb_query_buffer_assign_ptr(buf, "INT32", keys, FALSE)
+qry@ptr <- tiledb:::libtiledb_query_set_buffer_ptr(qry@ptr, "keys", buf)
+
 tiledb_query_set_layout(qry, "UNORDERED")
 tiledb_query_submit(qry)
 tiledb_query_finalize(qry)
@@ -183,15 +189,46 @@ qry <- tiledb_query_set_buffer(qry, "rows", rowdat)
 valdat <- integer(10)
 qry <- tiledb_query_set_buffer(qry, "vals", valdat)
 
+buf <- tiledb:::libtiledb_query_buffer_alloc_ptr(arr@ptr, "INT32", 10, TRUE)
+buf <- tiledb:::libtiledb_query_buffer_assign_ptr(buf, "INT32", keys, FALSE)
+qry@ptr <- tiledb:::libtiledb_query_set_buffer_ptr(qry@ptr, "keys", buf)
+
 qry <- tiledb_query_set_subarray(qry, c(4L,7L))
 tiledb_query_submit(qry)
 tiledb_query_finalize(qry)
 expect_equal(tiledb_query_status(qry), "COMPLETE")
 
+keydat <- tiledb:::libtiledb_query_get_buffer_ptr(buf, FALSE)
+
 n <- tiledb_query_result_buffer_elements(qry, "rows")
 expect_equal(n, 4L)
 expect_equal(rowdat[1:n], rows[4:7])
 expect_equal(valdat[1:n], vals[4:7])
+expect_equal(keydat[1:n], keys[4:7])
+n2 <- tiledb:::libtiledb_query_result_buffer_elements(qry@ptr, "rows", 0)
+expect_equal(n2, 0)                     # first element can be requested, is zero for fixed-sized
+
+## not as streamlined as it could, may need a wrapper for schema-from-query
+arrschptr <- tiledb:::libtiledb_query_get_schema(qry@ptr, tiledb_get_context()@ptr)
+sch <- tiledb:::tiledb_array_schema.from_ptr(arrschptr)
+
+attrs <- attrs(sch)
+isnullable <- tiledb:::libtiledb_attribute_get_nullable( attrs[["vals"]]@ptr )
+
+nv <- tiledb_query_result_buffer_elements_vec(qry, "vals", isnullable)
+expect_equal(length(nv), 2)             # vector accessors have two elements (or three if nullable)
+expect_equal(nv[1], 0)                  # first is zero for fixed-sized attributes
+expect_equal(nv[2], n)                  # second is what tiledb_query_result_buffer_elements has
+
+attrs <- attrs(sch)
+isnullable <- tiledb:::libtiledb_attribute_get_nullable( attrs[["keys"]]@ptr )
+
+nv <- tiledb_query_result_buffer_elements_vec(qry, "keys", isnullable)
+expect_equal(length(nv), 3)             # vector accessors have two elements (or three if nullable)
+expect_equal(nv[1], 0)                  # first is zero for fixed-sized attributes
+expect_equal(nv[2], n)                  # second is what tiledb_query_result_buffer_elements has
+expect_equal(nv[3], n)                  # third is length of validity buffer (if nullable)
+
 #})
 
 ## check for warning in insufficient memory

src/RcppExports.cpp

@@ -1806,6 +1806,19 @@ BEGIN_RCPP
     return rcpp_result_gen;
 END_RCPP
 }
+// libtiledb_query_result_buffer_elements_vec
+NumericVector libtiledb_query_result_buffer_elements_vec(XPtr<tiledb::Query> query, std::string attribute, bool nullable);
+RcppExport SEXP _tiledb_libtiledb_query_result_buffer_elements_vec(SEXP querySEXP, SEXP attributeSEXP, SEXP nullableSEXP) {
+BEGIN_RCPP
+    Rcpp::RObject rcpp_result_gen;
+    Rcpp::RNGScope rcpp_rngScope_gen;
+    Rcpp::traits::input_parameter< XPtr<tiledb::Query> >::type query(querySEXP);
+    Rcpp::traits::input_parameter< std::string >::type attribute(attributeSEXP);
+    Rcpp::traits::input_parameter< bool >::type nullable(nullableSEXP);
+    rcpp_result_gen = Rcpp::wrap(libtiledb_query_result_buffer_elements_vec(query, attribute, nullable));
+    return rcpp_result_gen;
+END_RCPP
+}
 // libtiledb_query_get_fragment_num
 int libtiledb_query_get_fragment_num(XPtr<tiledb::Query> query);
 RcppExport SEXP _tiledb_libtiledb_query_get_fragment_num(SEXP querySEXP) {
@@ -1970,6 +1983,30 @@ BEGIN_RCPP
     return rcpp_result_gen;
 END_RCPP
 }
+// libtiledb_query_get_array
+XPtr<tiledb::Array> libtiledb_query_get_array(XPtr<tiledb::Query> query, XPtr<tiledb::Context> ctx);
+RcppExport SEXP _tiledb_libtiledb_query_get_array(SEXP querySEXP, SEXP ctxSEXP) {
+BEGIN_RCPP
+    Rcpp::RObject rcpp_result_gen;
+    Rcpp::RNGScope rcpp_rngScope_gen;
+    Rcpp::traits::input_parameter< XPtr<tiledb::Query> >::type query(querySEXP);
+    Rcpp::traits::input_parameter< XPtr<tiledb::Context> >::type ctx(ctxSEXP);
+    rcpp_result_gen = Rcpp::wrap(libtiledb_query_get_array(query, ctx));
+    return rcpp_result_gen;
+END_RCPP
+}
+// libtiledb_query_get_schema
+XPtr<tiledb::ArraySchema> libtiledb_query_get_schema(XPtr<tiledb::Query> query, XPtr<tiledb::Context> ctx);
+RcppExport SEXP _tiledb_libtiledb_query_get_schema(SEXP querySEXP, SEXP ctxSEXP) {
+BEGIN_RCPP
+    Rcpp::RObject rcpp_result_gen;
+    Rcpp::RNGScope rcpp_rngScope_gen;
+    Rcpp::traits::input_parameter< XPtr<tiledb::Query> >::type query(querySEXP);
+    Rcpp::traits::input_parameter< XPtr<tiledb::Context> >::type ctx(ctxSEXP);
+    rcpp_result_gen = Rcpp::wrap(libtiledb_query_get_schema(query, ctx));
+    return rcpp_result_gen;
+END_RCPP
+}
 // libtiledb_query_condition
 XPtr<tiledb::QueryCondition> libtiledb_query_condition(XPtr<tiledb::Context> ctx);
 RcppExport SEXP _tiledb_libtiledb_query_condition(SEXP ctxSEXP) {
@@ -2574,6 +2611,7 @@ static const R_CallMethodDef CallEntries[] = {
     {"_tiledb_libtiledb_query_finalize", (DL_FUNC) &_tiledb_libtiledb_query_finalize, 1},
     {"_tiledb_libtiledb_query_status", (DL_FUNC) &_tiledb_libtiledb_query_status, 1},
     {"_tiledb_libtiledb_query_result_buffer_elements", (DL_FUNC) &_tiledb_libtiledb_query_result_buffer_elements, 3},
+    {"_tiledb_libtiledb_query_result_buffer_elements_vec", (DL_FUNC) &_tiledb_libtiledb_query_result_buffer_elements_vec, 3},
     {"_tiledb_libtiledb_query_get_fragment_num", (DL_FUNC) &_tiledb_libtiledb_query_get_fragment_num, 1},
     {"_tiledb_libtiledb_query_get_fragment_uri", (DL_FUNC) &_tiledb_libtiledb_query_get_fragment_uri, 2},
     {"_tiledb_libtiledb_query_get_fragment_timestamp_range", (DL_FUNC) &_tiledb_libtiledb_query_get_fragment_timestamp_range, 2},
@@ -2587,6 +2625,8 @@ static const R_CallMethodDef CallEntries[] = {
     {"_tiledb_libtiledb_query_get_range", (DL_FUNC) &_tiledb_libtiledb_query_get_range, 3},
     {"_tiledb_libtiledb_query_get_range_var", (DL_FUNC) &_tiledb_libtiledb_query_get_range_var, 3},
     {"_tiledb_libtiledb_query_set_condition", (DL_FUNC) &_tiledb_libtiledb_query_set_condition, 2},
+    {"_tiledb_libtiledb_query_get_array", (DL_FUNC) &_tiledb_libtiledb_query_get_array, 2},
+    {"_tiledb_libtiledb_query_get_schema", (DL_FUNC) &_tiledb_libtiledb_query_get_schema, 2},
     {"_tiledb_libtiledb_query_condition", (DL_FUNC) &_tiledb_libtiledb_query_condition, 1},
     {"_tiledb_libtiledb_query_condition_init", (DL_FUNC) &_tiledb_libtiledb_query_condition_init, 5},
     {"_tiledb_libtiledb_query_condition_combine", (DL_FUNC) &_tiledb_libtiledb_query_condition_combine, 3},

src/libtiledb.cpp

@@ -2912,12 +2912,30 @@ std::string libtiledb_query_status(XPtr<tiledb::Query> query) {
 // [[Rcpp::export]]
 R_xlen_t libtiledb_query_result_buffer_elements(XPtr<tiledb::Query> query,
                                                 std::string attribute,
-                                                int32_t which=1) {
+                                                int32_t which = 1) {
     auto nelements = query->result_buffer_elements()[attribute];
     R_xlen_t nelem = (which == 0 ? nelements.first : nelements.second);
     return nelem;
 }
 
+// [[Rcpp::export]]
+NumericVector libtiledb_query_result_buffer_elements_vec(XPtr<tiledb::Query> query,
+                                                         std::string attribute,
+                                                         bool nullable = false) {
+    if (nullable) {
+        auto nelem = query->result_buffer_elements_nullable()[attribute];
+        auto vec = NumericVector( {
+                static_cast<double>(std::get<0>(nelem)),
+                static_cast<double>(std::get<1>(nelem)),
+                static_cast<double>(std::get<2>(nelem))
+            });
+        return vec;
+    } else {
+        auto nelem = query->result_buffer_elements()[attribute];
+        return NumericVector( { static_cast<double>(nelem.first), static_cast<double>(nelem.second) });
+    }
+}
+
 // [[Rcpp::export]]
 int libtiledb_query_get_fragment_num(XPtr<tiledb::Query> query) {
   if (query->query_type() != TILEDB_WRITE) {
@@ -3208,6 +3226,23 @@ XPtr<tiledb::Query> libtiledb_query_set_condition(XPtr<tiledb::Query> query,
     return query;
 }
 
+// Note that the Array pointer returned here from a Query object is not owned and will not
+// outlive the query object
+//
+// [[Rcpp::export]]
+XPtr<tiledb::Array> libtiledb_query_get_array(XPtr<tiledb::Query> query, XPtr<tiledb::Context> ctx) {
+    auto arr = query->array();
+    auto cptr = arr.ptr().get();
+    return XPtr<tiledb::Array>(new tiledb::Array(*ctx.get(), cptr, false));
+}
+
+// [[Rcpp::export]]
+XPtr<tiledb::ArraySchema> libtiledb_query_get_schema(XPtr<tiledb::Query> query,
+                                                     XPtr<tiledb::Context> ctx) {
+    auto arr = query->array();
+    return libtiledb_array_schema_load(ctx, arr.uri()); // returns an XPtr<tiledb::ArraySchema>
+}
+
 /**
  * Query Condition
  */