Change query plan capnp to use structured data (#4611)

This PR changes the CAPNP format not to serialize the JSON dump of the
Query Plan, but to serialize just the required info from the Cloud
server (only `Query strategy` for now).

It also fixes a couple of issues found during testing, the most
ambiguous being:
- Part of the information Query Plan prints to the user is `VFS.Backend`
. For local queries that’s trivial to get from parsing the array URI.
For remote ones, it’s feasible again to get it from URI if it’s of the
form `tiledb://<namespace>/<backend>://<bucket>/<array_name>` but not if
it’s in form `tiledb://<namespace>/<array-name>`.
To address this for now I am printing `Unknown` as `VFS.Backend` for
such array URIs but I am open to better ideas.

example QueryPlan output: 
```
"TileDB Query Plan": {
        "Array.Type": "sparse",
        "Array.URI": "tiledb://unit/s3://tiledb-ypatia/scratch/aov2w20",
        "Query.Attributes": [
            "value"
        ],
        "Query.Dimensions": [
            "__tiledb_rows"
        ],
        "Query.Layout": "global-order",
        "Query.Strategy.Name": "SparseGlobalOrderReader",
        "VFS.Backend": "s3"
    }
```

---
TYPE: NO_HISTORY
DESC: Change query plan capnp to use structured data

Author: ypatia

test/src/unit-request-handlers.cc

@@ -85,7 +85,7 @@ struct HandleQueryPlanRequestFx : RequestHandlerFx {
   }
 
   virtual shared_ptr<ArraySchema> create_schema() override;
-  std::string call_handler(SerializationType stype, Query& query);
+  QueryPlan call_handler(SerializationType stype, Query& query);
 };
 
 struct HandleConsolidationPlanRequestFx : RequestHandlerFx {
@@ -229,7 +229,7 @@ TEST_CASE_METHOD(
   auto query_plan_ser_deser = call_handler(stype, *query->query_);
 
   // Compare the two query plans
-  REQUIRE(query_plan->view() == query_plan_ser_deser);
+  REQUIRE(query_plan->view() == query_plan_ser_deser.dump_json());
 
   // Clean up
   REQUIRE(tiledb_array_close(ctx, array) == TILEDB_OK);
@@ -466,7 +466,7 @@ shared_ptr<ArraySchema> HandleQueryPlanRequestFx::create_schema() {
   return schema;
 }
 
-std::string HandleQueryPlanRequestFx::call_handler(
+QueryPlan HandleQueryPlanRequestFx::call_handler(
     SerializationType stype, Query& query) {
   auto ctx = tiledb::Context();
   auto array = tiledb::Array(ctx, uri_.to_string(), TILEDB_READ);
@@ -483,8 +483,8 @@ std::string HandleQueryPlanRequestFx::call_handler(
       resp_buf);
   REQUIRE(rval == TILEDB_OK);
 
-  auto query_plan =
-      serialization::deserialize_query_plan_response(stype, resp_buf->buffer());
+  auto query_plan = serialization::deserialize_query_plan_response(
+      query, stype, resp_buf->buffer());
 
   tiledb_buffer_handle_t::break_handle(req_buf);
   tiledb_buffer_handle_t::break_handle(resp_buf);

tiledb/api/c_api/query_plan/query_plan_api.cc

@@ -48,6 +48,12 @@ capi_return_t tiledb_query_get_plan(
     throw CAPIStatusException("argument `query` may not be nullptr");
   }
 
+  if ((*query->query_).array()->is_remote()) {
+    throw std::logic_error(
+        "Failed to create a query plan; Remote arrays"
+        "are not currently supported.");
+  }
+
   sm::QueryPlan plan(*query->query_);
 
   *rv = tiledb_string_handle_t::make_handle(plan.dump_json());

tiledb/sm/c_api/tiledb.cc

@@ -4347,7 +4347,7 @@ capi_return_t tiledb_handle_query_plan_request(
     const tiledb_buffer_t* request,
     tiledb_buffer_t* response) {
   if (sanity_check(ctx, array) == TILEDB_ERR) {
-    throw std::invalid_argument("Array paramter must be valid.");
+    throw std::invalid_argument("Array parameter must be valid.");
   }
 
   api::ensure_buffer_is_valid(request);
@@ -4362,7 +4362,7 @@ capi_return_t tiledb_handle_query_plan_request(
   sm::QueryPlan plan(query);
 
   tiledb::sm::serialization::serialize_query_plan_response(
-      plan.dump_json(),
+      plan,
       static_cast<tiledb::sm::SerializationType>(serialization_type),
       response->buffer());
 

tiledb/sm/filesystem/uri.cc

@@ -326,7 +326,17 @@ std::string URI::to_path(const std::string& uri) {
 }
 
 std::string URI::backend_name() const {
-  return uri_.substr(0, uri_.find_first_of(':'));
+  if (is_tiledb(uri_)) {
+    std::string array_ns, array_uri;
+    throw_if_not_ok(URI(uri_).get_rest_components(&array_ns, &array_uri));
+    auto prefix = array_uri.substr(0, array_uri.find_first_of(':'));
+    if (prefix == array_uri) {  // no `:` separator found in URI
+      prefix = "Unknown";
+    }
+    return prefix;
+  } else {
+    return uri_.substr(0, uri_.find_first_of(':'));
+  }
 }
 
 std::string URI::to_path() const {

tiledb/sm/query_plan/query_plan.cc

@@ -48,20 +48,18 @@ namespace sm {
 /* ********************************* */
 /*     CONSTRUCTORS & DESTRUCTORS    */
 /* ********************************* */
-QueryPlan::QueryPlan(Query& query) {
-  if (query.array()->is_remote()) {
-    throw std::logic_error(
-        "Failed to create a query plan; Remote arrays"
-        "are not currently supported.");
-  }
-
+QueryPlan::QueryPlan(Query& query, std::string strategy) {
   array_uri_ = query.array()->array_uri().to_string();
   vfs_backend_ = URI(array_uri_).backend_name();
   query_layout_ = query.layout();
 
-  // This most probably ends up creating the strategy on the query
-  auto strategy_ptr = query.strategy();
-  strategy_name_ = strategy_ptr->name();
+  if (strategy.empty()) {
+    // This most probably ends up creating the strategy on the query
+    auto strategy_ptr = query.strategy();
+    strategy_name_ = strategy_ptr->name();
+  } else {
+    strategy_name_ = strategy;
+  }
 
   array_type_ = query.array()->array_schema_latest().array_type();
 

tiledb/sm/query_plan/query_plan.h

@@ -65,8 +65,9 @@ class QueryPlan {
    * Constructor
    *
    * @param query A query object for which we want to calculate the plan
+   * @param strategy The strategy of the query, required only for remote queries
    */
-  QueryPlan(Query& query);
+  QueryPlan(Query& query, const std::string strategy = "");
 
   /* ****************************** */
   /*              API               */
@@ -81,6 +82,13 @@ class QueryPlan {
    */
   std::string dump_json(uint32_t indent = 4);
 
+  /*
+   * Get the strategy name stored in the query plan
+   */
+  inline std::string strategy() const {
+    return strategy_name_;
+  }
+
  private:
   /* ****************************** */
   /*       PRIVATE ATTRIBUTES       */

tiledb/sm/serialization/query.cc

@@ -1613,7 +1613,8 @@ Status query_from_capnp(
     void* buffer_start,
     CopyState* const copy_state,
     Query* const query,
-    ThreadPool* compute_tp) {
+    ThreadPool* compute_tp,
+    const bool query_plan) {
   using namespace tiledb::sm;
 
   auto array = query->array();
@@ -2114,12 +2115,16 @@ Status query_from_capnp(
           attr_state->validity_len_data.swap(validity_buff);
 
           throw_if_not_ok(query->set_data_buffer(
-              name, varlen_data, &attr_state->var_len_size, true, true));
+              name, varlen_data, &attr_state->var_len_size, !query_plan, true));
           throw_if_not_ok(query->set_offsets_buffer(
-              name, offsets, &attr_state->fixed_len_size, true, true));
+              name, offsets, &attr_state->fixed_len_size, !query_plan, true));
           if (nullable) {
             throw_if_not_ok(query->set_validity_buffer(
-                name, validity, &attr_state->validity_len_size, true, true));
+                name,
+                validity,
+                &attr_state->validity_len_size,
+                !query_plan,
+                true));
           }
         } else {
           auto* data = attribute_buffer_start;
@@ -2142,10 +2147,14 @@ Status query_from_capnp(
           attr_state->validity_len_data.swap(validity_buff);
 
           throw_if_not_ok(query->set_data_buffer(
-              name, data, &attr_state->fixed_len_size, true, true));
+              name, data, &attr_state->fixed_len_size, !query_plan, true));
           if (nullable) {
             throw_if_not_ok(query->set_validity_buffer(
-                name, validity, &attr_state->validity_len_size, true, true));
+                name,
+                validity,
+                &attr_state->validity_len_size,
+                !query_plan,
+                true));
           }
         }
       } else {

tiledb/sm/serialization/query.h

@@ -275,7 +275,8 @@ Status query_from_capnp(
     void* buffer_start,
     CopyState* const copy_state,
     Query* const query,
-    ThreadPool* compute_tp);
+    ThreadPool* compute_tp,
+    const bool query_plan = false);
 
 #endif
 

tiledb/sm/serialization/query_plan.cc

@@ -75,23 +75,24 @@ void query_plan_request_from_capnp(
         nullptr,
         nullptr,
         &query,
-        &compute_tp));
+        &compute_tp,
+        true));
   }
 }
 
 void query_plan_response_to_capnp(
-    capnp::QueryPlanResponse::Builder& builder, const std::string& query_plan) {
-  builder.setQueryPlan(query_plan);
+    capnp::QueryPlanResponse::Builder& builder, const QueryPlan& query_plan) {
+  builder.setStrategy(query_plan.strategy());
 }
 
 std::string query_plan_response_from_capnp(
     const capnp::QueryPlanResponse::Reader& reader) {
-  std::string query_plan;
-  if (reader.hasQueryPlan()) {
-    query_plan = reader.getQueryPlan().cStr();
+  std::string strategy;
+  if (reader.hasStrategy()) {
+    strategy = reader.getStrategy().cStr();
   }
 
-  return query_plan;
+  return strategy;
 }
 
 void serialize_query_plan_request(
@@ -193,7 +194,7 @@ void deserialize_query_plan_request(
 }
 
 void serialize_query_plan_response(
-    const std::string& query_plan,
+    const QueryPlan& query_plan,
     SerializationType serialization_type,
     Buffer& response) {
   try {
@@ -243,8 +244,10 @@ void serialize_query_plan_response(
   }
 }
 
-std::string deserialize_query_plan_response(
-    SerializationType serialization_type, const Buffer& response) {
+QueryPlan deserialize_query_plan_response(
+    Query& query,
+    SerializationType serialization_type,
+    const Buffer& response) {
   try {
     switch (serialization_type) {
       case SerializationType::JSON: {
@@ -255,7 +258,7 @@ std::string deserialize_query_plan_response(
         json.decode(
             kj::StringPtr(static_cast<const char*>(response.data())), builder);
         capnp::QueryPlanResponse::Reader reader = builder.asReader();
-        return query_plan_response_from_capnp(reader);
+        return QueryPlan(query, query_plan_response_from_capnp(reader));
       }
       case SerializationType::CAPNP: {
         const auto mBytes = reinterpret_cast<const kj::byte*>(response.data());
@@ -264,7 +267,7 @@ std::string deserialize_query_plan_response(
             response.size() / sizeof(::capnp::word)));
         capnp::QueryPlanResponse::Reader reader =
             array_reader.getRoot<capnp::QueryPlanResponse>();
-        return query_plan_response_from_capnp(reader);
+        return QueryPlan(query, query_plan_response_from_capnp(reader));
       }
       default: {
         throw Status_SerializationError(
@@ -298,13 +301,13 @@ void deserialize_query_plan_request(
 }
 
 void serialize_query_plan_response(
-    const std::string&, const SerializationType, Buffer&) {
+    const QueryPlan&, const SerializationType, Buffer&) {
   throw Status_SerializationError(
       "Cannot serialize; serialization not enabled.");
 }
 
-std::string deserialize_query_plan_response(
-    const SerializationType, const Buffer&) {
+QueryPlan deserialize_query_plan_response(
+    Query&, const SerializationType, const Buffer&) {
   throw Status_SerializationError(
       "Cannot serialize; serialization not enabled.");
 }

tiledb/sm/serialization/query_plan.h

@@ -84,19 +84,22 @@ void deserialize_query_plan_request(
  * @param response Buffer to store serialized bytes in.
  */
 void serialize_query_plan_response(
-    const std::string& query_plan,
+    const QueryPlan& query_plan,
     const SerializationType serialization_type,
     Buffer& response);
 
 /**
- * Deserialize a Query Plan request to cap'n proto object
+ * Deserialize a Query Plan response from cap'n proto object
  *
+ * @param query The query the plan is requested for.
  * @param serialization_type Format to serialize from: Cap'n Proto or JSON.
  * @param response Buffer to read serialized bytes from.
- * @return The requested query plan as a string.
+ * @return The requested query plan.
  */
-std::string deserialize_query_plan_response(
-    const SerializationType serialization_type, const Buffer& response);
+QueryPlan deserialize_query_plan_response(
+    Query& query,
+    const SerializationType serialization_type,
+    const Buffer& response);
 
 }  // namespace serialization
 }  // namespace tiledb::sm