[feature](lance) Extend WARM UP SELECT with index prewarm (branch-4.2)

be/src/format_v2/lance/lance_session_manager.cpp

@@ -215,6 +215,12 @@ Status LanceSessionManager::_initialize() {
 
 Status LanceSessionManager::open_dataset(const char* uri, const char* const* storage_options,
                                          uint64_t version, LanceDataset** dataset) {
+    return _open_dataset(uri, storage_options, version, dataset, false);
+}
+
+Status LanceSessionManager::_open_dataset(const char* uri, const char* const* storage_options,
+                                          uint64_t version, LanceDataset** dataset,
+                                          bool redact_errors) {
     if (uri == nullptr || dataset == nullptr) {
         return Status::InvalidArgument("Lance dataset URI and output must not be null");
     }
@@ -231,9 +237,46 @@ Status LanceSessionManager::open_dataset(const char* uri, const char* const* sto
 
     *dataset = lance_dataset_open_with_session(uri, storage_options, version, _session);
     if (*dataset == nullptr) {
+        // Constructing an SDK error Status can itself log the provider message. Redact before
+        // that construction, not only when returning the prewarm RPC response.
+        if (redact_errors) {
+            return Status::InternalError(
+                    "Unable to open the pinned Lance snapshot for index prewarm");
+        }
         return lance_error("open Lance dataset with shared session");
     }
     return Status::OK();
 }
 
+Status LanceSessionManager::prewarm_index(const char* uri, const char* const* storage_options,
+                                          uint64_t version, const char* index_name) {
+    if (uri == nullptr || *uri == '\0' || version == 0 || index_name == nullptr ||
+        *index_name == '\0') {
+        return Status::InvalidArgument(
+                "Lance index prewarm requires URI, index name and a fixed version");
+    }
+    // Bound concurrent SDK prewarm IO per BE without blocking all RPC workers behind a mutex.
+    // A failed or timed-out caller may retry after the existing SDK operation finishes.
+    std::unique_lock lock(_prewarm_mutex, std::try_to_lock);
+    if (!lock.owns_lock()) {
+        return Status::Cancelled("Another Lance index prewarm is running on this backend");
+    }
+    LanceDataset* raw_dataset = nullptr;
+    auto status = _open_dataset(uri, storage_options, version, &raw_dataset, true);
+    std::unique_ptr<LanceDataset, decltype(&lance_dataset_close)> dataset(raw_dataset,
+                                                                          lance_dataset_close);
+    if (!status.ok()) {
+        // SDK storage errors can include credentials or signed URLs. Return operation context only.
+        return Status::InternalError("Unable to open the pinned Lance snapshot for index prewarm");
+    }
+    if (lance_dataset_version(dataset.get()) != version) {
+        return Status::InternalError("Lance index prewarm opened an unexpected snapshot");
+    }
+    if (lance_dataset_prewarm_index(dataset.get(), index_name) != 0) {
+        return Status::InternalError(
+                "Lance index prewarm failed: index unavailable, unsupported, or storage IO failed");
+    }
+    return Status::OK();
+}
+
 } // namespace doris::format::lance

be/src/format_v2/lance/lance_session_manager.h

@@ -58,11 +58,17 @@ class LanceSessionManager final {
     Status open_dataset(const char* uri, const char* const* storage_options, uint64_t version,
                         LanceDataset** dataset);
 
+    Status prewarm_index(const char* uri, const char* const* storage_options, uint64_t version,
+                         const char* index_name);
+
 private:
     Status _initialize();
+    Status _open_dataset(const char* uri, const char* const* storage_options, uint64_t version,
+                         LanceDataset** dataset, bool redact_errors);
 
     Config _config;
     std::once_flag _initialize_once;
+    std::mutex _prewarm_mutex;
     LanceSession* _session = nullptr;
     std::unique_ptr<LanceSessionMetrics> _metrics;
     Status _initialize_status = Status::OK();

be/src/service/internal_service.cpp

@@ -42,6 +42,7 @@
 #include <sys/stat.h>
 
 #include <algorithm>
+#include <chrono>
 #include <exception>
 #include <filesystem>
 #include <memory>
@@ -78,6 +79,7 @@
 #include "format/json/new_json_reader.h"
 #include "format/text/text_reader.h"
 #include "format_v2/column_data.h"
+#include "format_v2/lance/lance_session_manager.h"
 #include "format_v2/schema_reader.h"
 #include "format_v2/table/lance_reader.h"
 #include "io/fs/local_file_system.h"
@@ -230,6 +232,9 @@ PInternalService::PInternalService(ExecEnv* exec_env)
                                    ? config::brpc_heavy_work_pool_max_queue_size
                                    : std::max(10240, CpuInfo::num_cores() * 320),
                            "brpc_heavy"),
+          // The synchronous SDK cannot be interrupted. Keep abandoned prewarm IO off the
+          // shared RPC workers, with at most one running and one queued request per service.
+          _lance_index_prewarm_pool(1, 1, "lance_prewarm"),
           // peer fetch threadpool isolates fetch_peer_data from heavy load traffic to avoid peer reads starving imports.
           _peer_fetch_pool(resolved_brpc_peer_fetch_pool_threads(),
                            resolved_brpc_peer_fetch_pool_max_queue_size(), "brpc_peer_fetch"),
@@ -811,6 +816,60 @@ void PInternalService::outfile_write_success(google::protobuf::RpcController* co
     }
 }
 
+void PInternalService::prewarm_lance_index(google::protobuf::RpcController* controller,
+                                           const PLanceIndexPrewarmRequest* request,
+                                           PLanceIndexPrewarmResponse* response,
+                                           google::protobuf::Closure* done) {
+    const auto received = std::chrono::steady_clock::now();
+    bool offered = _lance_index_prewarm_pool.try_offer([request, response, done, received]() {
+        brpc::ClosureGuard closure_guard(done);
+        Status status;
+        try {
+            auto run = [&]() -> Status {
+                DBUG_EXECUTE_IF("PInternalService.prewarm_lance_index.fail", {
+                    return Status::InternalError("Injected Lance index prewarm failure");
+                });
+                const auto queued_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
+                                               std::chrono::steady_clock::now() - received)
+                                               .count();
+                if (!request->has_timeout_ms() || request->timeout_ms() <= queued_ms) {
+                    return Status::TimedOut("Lance index prewarm expired before execution");
+                }
+                // C strings cannot represent embedded NUL bytes; reject rather than silently
+                // opening a different path/index or truncating a vended credential.
+                if (request->dataset_uri().find('\0') != std::string::npos ||
+                    request->index_name().find('\0') != std::string::npos) {
+                    return Status::InvalidArgument("Invalid Lance prewarm URI or index name");
+                }
+                std::vector<const char*> options;
+                options.reserve(request->storage_options_size() * 2 + 1);
+                for (const auto& [key, value] : request->storage_options()) {
+                    if (key.find('\0') != std::string::npos ||
+                        value.find('\0') != std::string::npos) {
+                        return Status::InvalidArgument("Invalid Lance prewarm storage option");
+                    }
+                    options.push_back(key.c_str());
+                    options.push_back(value.c_str());
+                }
+                options.push_back(nullptr);
+                RETURN_IF_ERROR(format::lance::LanceSessionManager::instance().prewarm_index(
+                        request->dataset_uri().c_str(), options.data(), request->dataset_version(),
+                        request->index_name().c_str()));
+                response->set_dataset_version(request->dataset_version());
+                return Status::OK();
+            };
+            status = run();
+        } catch (const std::exception&) {
+            // Provider exception messages may contain the storage credentials in the request.
+            status = Status::InternalError("Lance index prewarm failed on this backend");
+        }
+        status.to_protobuf(response->mutable_status());
+    });
+    if (!offered) {
+        offer_failed(response, done, _lance_index_prewarm_pool);
+    }
+}
+
 void PInternalService::fetch_table_schema(google::protobuf::RpcController* controller,
                                           const PFetchTableSchemaRequest* request,
                                           PFetchTableSchemaResult* result,

be/src/service/internal_service.h

@@ -75,6 +75,11 @@ class PInternalService : public PBackendService {
                                POutfileWriteSuccessResult* result,
                                google::protobuf::Closure* done) override;
 
+    void prewarm_lance_index(google::protobuf::RpcController* controller,
+                             const PLanceIndexPrewarmRequest* request,
+                             PLanceIndexPrewarmResponse* response,
+                             google::protobuf::Closure* done) override;
+
     void fetch_table_schema(google::protobuf::RpcController* controller,
                             const PFetchTableSchemaRequest* request,
                             PFetchTableSchemaResult* result,
@@ -274,6 +279,7 @@ class PInternalService : public PBackendService {
     // define the interface for reading and writing data as heavy interface
     // otherwise as light interface
     FifoThreadPool _heavy_work_pool;
+    FifoThreadPool _lance_index_prewarm_pool;
     FifoThreadPool _peer_fetch_pool;
     FifoThreadPool _light_work_pool;
     FifoThreadPool _arrow_flight_work_pool;

be/test/format_v2/lance/lance_session_manager_test.cpp

@@ -22,12 +22,14 @@
 
 #include <chrono>
 #include <filesystem>
+#include <lance/lance.hpp>
 #include <memory>
 #include <string>
 #include <utility>
 
 #include "common/metrics/doris_metrics.h"
 #include "common/metrics/metrics.h"
+#include "util/defer_op.h"
 
 namespace doris::format::lance {
 namespace {
@@ -99,6 +101,59 @@ TEST(LanceSessionManagerTest, PublishesSessionCacheMetrics) {
     EXPECT_NE(entity->get_metric("lance_session_metadata_cache_misses_total"), nullptr);
 }
 
+TEST(LanceSessionManagerTest, PrewarmRequiresPinnedVersionAndIndex) {
+    LanceSessionManager manager(
+            {.lance_index_cache_size_bytes = 8 * 1024 * 1024, .lance_data_cache_path = ""});
+    EXPECT_FALSE(manager.prewarm_index(nullptr, nullptr, 1, "idx").ok());
+    EXPECT_FALSE(manager.prewarm_index("dataset.lance", nullptr, 0, "idx").ok());
+    EXPECT_FALSE(manager.prewarm_index("dataset.lance", nullptr, 1, nullptr).ok());
+    EXPECT_FALSE(manager.prewarm_index("dataset.lance", nullptr, 1, "").ok());
+}
+
+TEST(LanceSessionManagerTest, PrewarmErrorsDoNotExposeStorageDetails) {
+    LanceSessionManager manager(
+            {.lance_index_cache_size_bytes = 8 * 1024 * 1024, .lance_data_cache_path = ""});
+    const auto status =
+            manager.prewarm_index("/nonexistent/private-dataset.lance", nullptr, 1, "idx");
+    EXPECT_FALSE(status.ok());
+    EXPECT_EQ(std::string::npos, status.to_string().find("private-dataset"));
+    EXPECT_FALSE(
+            manager.prewarm_index(lance_fixture_path().c_str(), nullptr, 1, "missing_idx").ok());
+}
+
+TEST(LanceSessionManagerTest, PrewarmUsesTheQuerySessionAndFixedSnapshot) {
+    const auto path = unique_cache_path();
+    std::filesystem::copy(lance_fixture_path(), path, std::filesystem::copy_options::recursive);
+    Defer cleanup {[&] {
+        std::error_code error;
+        std::filesystem::remove_all(path, error);
+    }};
+    auto writer = ::lance::Dataset::open(path.string());
+    writer.create_scalar_index("row_id", LANCE_SCALAR_BTREE, "value_idx");
+    const uint64_t version = writer.version();
+    LanceSessionManager manager({.lance_index_cache_size_bytes = 8 * 1024 * 1024,
+                                 .lance_metadata_cache_size_bytes = 4 * 1024 * 1024,
+                                 .lance_data_cache_path = ""});
+    ASSERT_TRUE(manager.prewarm_index(path.c_str(), nullptr, version, "value_idx").ok());
+    auto* entity = DorisMetrics::instance()->server_entity();
+    entity->trigger_hook_unlocked(true);
+    auto* hits =
+            static_cast<IntCounter*>(entity->get_metric("lance_session_index_cache_hits_total"));
+    const auto before = hits->value();
+
+    // A separate query handle must see the cache populated by the completed prewarm call.
+    LanceDataset* raw_dataset = nullptr;
+    ASSERT_TRUE(manager.open_dataset(path.c_str(), nullptr, version, &raw_dataset).ok());
+    LanceDatasetPtr query(raw_dataset, lance_dataset_close);
+    ASSERT_EQ(0, lance_dataset_prewarm_index(query.get(), "value_idx"));
+    entity->trigger_hook_unlocked(true);
+    EXPECT_GT(hits->value(), before);
+
+    writer.drop_index("value_idx");
+    EXPECT_FALSE(manager.prewarm_index(path.c_str(), nullptr, writer.version(), "value_idx").ok());
+    EXPECT_TRUE(manager.prewarm_index(path.c_str(), nullptr, version, "value_idx").ok());
+}
+
 TEST(LanceSessionManagerTest, CreatesFoyerBackedSession) {
     const auto cache_path = unique_cache_path();
     std::filesystem::create_directories(cache_path);

be/test/service/lance_index_prewarm_test.cpp

@@ -0,0 +1,123 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gtest/gtest.h>
+
+#include <chrono>
+#include <future>
+#include <utility>
+
+#include "common/config.h"
+#include "load/channel/load_stream_mgr.h"
+#include "runtime/exec_env.h"
+#include "service/internal_service.h"
+
+namespace doris {
+namespace {
+
+class PrewarmClosure : public google::protobuf::Closure {
+public:
+    void Run() override { completed.set_value(); }
+    std::promise<void> completed;
+};
+
+class LanceIndexPrewarmServiceTest : public testing::Test {
+protected:
+    void SetUp() override {
+        auto heavy = std::exchange(config::brpc_heavy_work_pool_threads, 1);
+        auto light = std::exchange(config::brpc_light_work_pool_threads, 1);
+        auto peer = std::exchange(config::brpc_peer_fetch_pool_threads, 1);
+        auto flight = std::exchange(config::brpc_arrow_flight_work_pool_threads, 1);
+        auto* env = ExecEnv::GetInstance();
+        previous_load_mgr = std::move(env->_load_stream_mgr);
+        env->_load_stream_mgr = std::make_unique<LoadStreamMgr>(1);
+        service = std::make_unique<PInternalService>(env);
+        config::brpc_heavy_work_pool_threads = heavy;
+        config::brpc_light_work_pool_threads = light;
+        config::brpc_peer_fetch_pool_threads = peer;
+        config::brpc_arrow_flight_work_pool_threads = flight;
+    }
+
+    void TearDown() override {
+        service.reset();
+        ExecEnv::GetInstance()->_load_stream_mgr = std::move(previous_load_mgr);
+    }
+
+    std::unique_ptr<PInternalService> service;
+    std::unique_ptr<LoadStreamMgr> previous_load_mgr;
+};
+
+TEST_F(LanceIndexPrewarmServiceTest, PrewarmDoesNotUseTheSharedHeavyWorker) {
+    std::promise<void> entered;
+    std::promise<void> release;
+    auto released = release.get_future().share();
+    service->_heavy_work_pool.try_offer([&, released] {
+        entered.set_value();
+        released.wait();
+    });
+    entered.get_future().wait();
+    PLanceIndexPrewarmRequest request;
+    request.set_timeout_ms(0);
+    PLanceIndexPrewarmResponse response;
+    PrewarmClosure done;
+    auto completed = done.completed.get_future();
+    service->prewarm_lance_index(nullptr, &request, &response, &done);
+    auto result = completed.wait_for(std::chrono::seconds(2));
+    release.set_value();
+    completed.wait();
+    EXPECT_EQ(std::future_status::ready, result);
+    EXPECT_FALSE(Status::create(response.status()).ok());
+}
+
+TEST_F(LanceIndexPrewarmServiceTest, StalledPrewarmHasBoundedAdmissionAndExpiredQueueIsRejected) {
+    std::promise<void> entered;
+    std::promise<void> release;
+    auto released = release.get_future().share();
+    service->_lance_index_prewarm_pool.try_offer([&, released] {
+        entered.set_value();
+        released.wait();
+    });
+    entered.get_future().wait();
+    PLanceIndexPrewarmRequest request;
+    request.set_timeout_ms(0);
+    PLanceIndexPrewarmResponse queued_response;
+    PLanceIndexPrewarmResponse rejected_response;
+    PrewarmClosure queued_done;
+    PrewarmClosure rejected_done;
+    auto queued = queued_done.completed.get_future();
+    auto rejected = rejected_done.completed.get_future();
+    service->prewarm_lance_index(nullptr, &request, &queued_response, &queued_done);
+    service->prewarm_lance_index(nullptr, &request, &rejected_response, &rejected_done);
+    auto rejection = rejected.wait_for(std::chrono::seconds(2));
+
+    // Stalled native work must not occupy the shared RPC worker or create more prewarm workers.
+    std::promise<void> heavy_done;
+    service->_heavy_work_pool.try_offer([&] { heavy_done.set_value(); });
+    auto heavy_future = heavy_done.get_future();
+    auto heavy = heavy_future.wait_for(std::chrono::seconds(2));
+    release.set_value();
+    queued.wait();
+    rejected.wait();
+    heavy_future.wait();
+    EXPECT_EQ(std::future_status::ready, rejection);
+    EXPECT_EQ(std::future_status::ready, heavy);
+    EXPECT_FALSE(Status::create(rejected_response.status()).ok());
+    EXPECT_TRUE(Status::create(queued_response.status()).is<ErrorCode::TIMEOUT>());
+}
+
+} // namespace
+} // namespace doris

fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4

@@ -542,7 +542,8 @@ supportedOtherStatement
             onTablesClause?
             properties=propertyClause?                                              #warmUpCluster
     | explain? WARM UP SELECT namedExpressionSeq
-      FROM warmUpSingleTableRef whereClause?                                        #warmUpSelect
+      FROM warmUpSingleTableRef whereClause?
+      properties=propertyClause?                                                   #warmUpSelect
     | BACKUP SNAPSHOT label=multipartIdentifier TO repo=identifier
         ((ON | EXCLUDE) LEFT_PAREN baseTableRef (COMMA baseTableRef)* RIGHT_PAREN)?
         properties=propertyClause?                                                  #backup