Merge pull request #3130 from activeloopai/stateless-v2

khustup2 · web-flow · commit 41e606eb2fb6 · 2026-02-07T09:36:02.000-08:00
Stateless support and switch to new deeplake-api.
diff --git a/DEEPLAKE_API_VERSION b/DEEPLAKE_API_VERSION
@@ -1 +1 @@
-4.5.0
+4.5.1
diff --git a/cpp/deeplake_pg/dl_catalog.cpp b/cpp/deeplake_pg/dl_catalog.cpp
diff --git a/cpp/deeplake_pg/dl_catalog.hpp b/cpp/deeplake_pg/dl_catalog.hpp
@@ -42,6 +42,10 @@ std::vector<table_meta> load_tables(const std::string& root_path, icm::string_ma
 std::vector<column_meta> load_columns(const std::string& root_path, icm::string_map<> creds);
 std::vector<index_meta> load_indexes(const std::string& root_path, icm::string_map<> creds);
 
+// Load tables and columns in parallel for better performance
+std::pair<std::vector<table_meta>, std::vector<column_meta>>
+load_tables_and_columns(const std::string& root_path, icm::string_map<> creds);
+
 void upsert_table(const std::string& root_path, icm::string_map<> creds, const table_meta& meta);
 void upsert_columns(const std::string& root_path, icm::string_map<> creds, const std::vector<column_meta>& columns);
 
diff --git a/cpp/deeplake_pg/duckdb_deeplake_convert.cpp b/cpp/deeplake_pg/duckdb_deeplake_convert.cpp
@@ -4,6 +4,7 @@
 #include "utils.hpp"
 
 #include <codecs/compression.hpp>
+#include <icm/vector.hpp>
 #include <nd/adapt.hpp>
 #include <nd/none.hpp>
 
@@ -85,7 +86,7 @@ T to_cpp_value(const duckdb::Value& val)
 
 nd::array to_deeplake_value_as_array_list(const duckdb::vector<duckdb::Value>& values)
 {
-    std::vector<nd::array> arr;
+    icm::vector<nd::array> arr;
     arr.reserve(values.size());
     for (const auto& v : values) {
         arr.push_back(pg::to_deeplake_value(v));
@@ -105,7 +106,7 @@ nd::array to_deeplake_value(const duckdb::LogicalType& duckdb_type, const duckdb
     }
     return switch_duckdb_type(duckdb_type, [&values]<typename T>() {
         if constexpr (std::is_same_v<T, bytea_type>) {
-            std::vector<nd::array> arr;
+            icm::vector<nd::array> arr;
             arr.reserve(values.size());
             for (const auto& val : values) {
                 duckdb::string_t blob_data = duckdb::StringValue::Get(val);
diff --git a/cpp/deeplake_pg/duckdb_deeplake_scan.cpp b/cpp/deeplake_pg/duckdb_deeplake_scan.cpp
@@ -43,7 +43,7 @@ struct deeplake_scan_bind_data final : public duckdb::TableFunctionData
 struct deeplake_scan_global_state final : public duckdb::GlobalTableFunctionState
 {
     duckdb::vector<duckdb::column_t> column_ids;
-    std::vector<base::function<async::promise<std::vector<icm::roaring>>()>> index_searchers;
+    icm::vector<base::function<async::promise<icm::vector<icm::roaring>>()>> index_searchers;
     duckdb::unique_ptr<duckdb::Expression> filter_expr;
     std::mutex index_search_mutex;
     heimdall::dataset_view_ptr index_search_result;
@@ -205,10 +205,10 @@ duckdb::unique_ptr<duckdb::FunctionData> deeplake_scan_bind(duckdb::ClientContex
     return duckdb::make_uniq<deeplake_scan_bind_data>(td, return_types);
 }
 
-base::function<async::promise<std::vector<icm::roaring>>()>
+base::function<async::promise<icm::vector<icm::roaring>>()>
 try_get_index_searcher(heimdall::column_view_ptr column_view, const duckdb::ConstantFilter& filter)
 {
-    base::function<async::promise<std::vector<icm::roaring>>()> result;
+    base::function<async::promise<icm::vector<icm::roaring>>()> result;
     auto index_holder = column_view->index_holder();
     ASSERT(index_holder != nullptr);
     auto constant = pg::to_deeplake_value(filter.constant);
@@ -250,7 +250,7 @@ try_get_index_searcher(heimdall::column_view_ptr column_view, const duckdb::Cons
             query_core::text_search_info info;
             info.column_name = column_view->name();
             info.type = query_core::text_search_info::search_type::equals;
-            info.search_values.push_back(std::vector<std::string>{filter.constant.ToString()});
+            info.search_values.push_back(icm::vector<std::string>{filter.constant.ToString()});
             if (index_holder->can_run_query(info)) {
                 result = [index_holder, si = std::move(info)]() {
                     return index_holder->run_query(si);
@@ -261,7 +261,7 @@ try_get_index_searcher(heimdall::column_view_ptr column_view, const duckdb::Cons
     return result;
 }
 
-base::function<async::promise<std::vector<icm::roaring>>()>
+base::function<async::promise<icm::vector<icm::roaring>>()>
 try_get_index_searcher(heimdall::column_view_ptr column_view, const duckdb::InFilter& filter)
 {
     query_core::inverted_index_search_info info;
@@ -277,10 +277,10 @@ try_get_index_searcher(heimdall::column_view_ptr column_view, const duckdb::InFi
     };
 }
 
-base::function<async::promise<std::vector<icm::roaring>>()>
+base::function<async::promise<icm::vector<icm::roaring>>()>
 try_get_index_searcher(heimdall::column_view_ptr column_view, const duckdb::TableFilter& filter)
 {
-    base::function<async::promise<std::vector<icm::roaring>>()> result;
+    base::function<async::promise<icm::vector<icm::roaring>>()> result;
     ASSERT(column_view != nullptr);
     if (column_view->index_holder() == nullptr) {
         return result;
@@ -887,14 +887,14 @@ class deeplake_scan_function_helper
         if (is_index_search_done()) {
             return;
         }
-        std::vector<async::promise<icm::roaring>> promises;
+        icm::vector<async::promise<icm::roaring>> promises;
         for (auto& is : global_state_.index_searchers) {
-            promises.push_back(is().then_any([](std::vector<icm::roaring>&& results) {
+            promises.push_back(is().then_any([](icm::vector<icm::roaring>&& results) {
                 ASSERT(results.size() == 1);
                 return std::move(results.front());
             }));
         }
-        auto combined_promise = async::combine(std::move(promises)).then_any([](std::vector<icm::roaring>&& results) {
+        auto combined_promise = async::combine(std::move(promises)).then_any([](icm::vector<icm::roaring>&& results) {
             ASSERT(!results.empty());
             icm::roaring& combined = results[0];
             for (size_t i = 1; i < results.size(); ++i) {
@@ -903,7 +903,7 @@ class deeplake_scan_function_helper
             return std::move(combined);
         });
         auto indices = combined_promise.get_future().get();
-        std::vector<int64_t> indices_vec;
+        icm::vector<int64_t> indices_vec;
         indices_vec.reserve(indices.cardinality());
         for (auto x : indices) {
             indices_vec.push_back(x);
@@ -955,7 +955,7 @@ class deeplake_scan_function_helper
         }
 
         ASSERT(output_.ColumnCount() == global_state_.column_ids.size());
-        std::vector<async::promise<void>> column_promises;
+        icm::vector<async::promise<void>> column_promises;
         // Fill output vectors column by column using table_data streamers
         for (unsigned i = 0; i < global_state_.column_ids.size(); ++i) {
             const auto col_idx = global_state_.column_ids[i];
diff --git a/cpp/deeplake_pg/extension_init.cpp b/cpp/deeplake_pg/extension_init.cpp
@@ -1197,10 +1197,18 @@ static void process_utility(PlannedStmt* pstmt,
         }
         // When root_path is set, auto-discover tables from the deeplake catalog
         if (vstmt->name != nullptr && pg_strcasecmp(vstmt->name, "deeplake.root_path") == 0) {
-            // Reload table metadata from the catalog at the new root_path
-            // This enables stateless multi-instance support where tables are
-            // auto-discovered when pointing to a shared root_path
-            pg::table_storage::instance().force_load_table_metadata();
+            // Track the previous root_path to detect actual changes
+            static thread_local std::string last_root_path;
+            auto current_root_path = pg::session_credentials::get_root_path();
+
+            if (current_root_path != last_root_path) {
+                // Path changed - force full reload
+                last_root_path = current_root_path;
+                pg::table_storage::instance().force_load_table_metadata();
+            } else {
+                // Same path - just check for catalog updates (fast path)
+                pg::table_storage::instance().load_table_metadata();
+            }
         }
     }
 }
diff --git a/cpp/deeplake_pg/hybrid_query_merge.hpp b/cpp/deeplake_pg/hybrid_query_merge.hpp
@@ -126,7 +126,7 @@ inline query_core::query_result merge_query_results(
     // Take top_k results
     size_t result_size = std::min(top_k, final_scores.size());
     std::vector<float> top_scores;
-    std::vector<int64_t> top_indices;
+    icm::vector<int64_t> top_indices;
     top_scores.reserve(result_size);
     top_indices.reserve(result_size);
 
diff --git a/cpp/deeplake_pg/index_search.cpp b/cpp/deeplake_pg/index_search.cpp
@@ -108,7 +108,7 @@ struct scan_opaque
 
 query_core::query_result run_index_search(nd::array input_array, std::string func_name, const std::string& column_name, pg::index_info& idx_info)
 {
-    std::vector<query_core::expr> args;
+    icm::vector<query_core::expr> args;
     args.emplace_back(query_core::expr::make_column_ref(column_name, std::string{}));
     args.emplace_back(query_core::expr::make_literal_array(std::move(input_array)));
     const bool is_cosine_similarity = (func_name == "COSINE_SIMILARITY");
@@ -194,7 +194,7 @@ icm::roaring run_exact_text_search(std::string text_value, StrategyNumber strate
         return {};
     }
     info.column_name = idx_info.column_name();
-    info.search_values.emplace_back(std::vector<std::string>{std::move(text_value)});
+    info.search_values.emplace_back(icm::vector<std::string>{std::move(text_value)});
     return idx_info.run_query(std::move(info));
 }
 
@@ -563,7 +563,7 @@ void collect_index_data(IndexScanDesc scan, ScanKey keys, int32_t nkeys, ScanKey
             }
         }
         if (nkeys > 0) {
-            std::vector<int64_t> row_numbers;
+            icm::vector<int64_t> row_numbers;
             row_numbers.reserve(result.cardinality());
             std::transform(result.begin(), result.end(), std::back_inserter(row_numbers), [](auto v) {
                 return static_cast<int64_t>(v);
diff --git a/cpp/deeplake_pg/nd_utils.hpp b/cpp/deeplake_pg/nd_utils.hpp
@@ -3,6 +3,7 @@
 #include "exceptions.hpp"
 
 #include <icm/shape.hpp>
+#include <icm/vector.hpp>
 #include <nd/adapt.hpp>
 #include <nd/array.hpp>
 #include <nd/none.hpp>
@@ -124,7 +125,7 @@ inline pg::array_type pg_to_nd_typed(ArrayType* array, bool copy_data = true)
         const auto nrows = dims[0];
         const auto ncols = dims[1];
         if (copy_data) {
-            std::vector<nd::array> data_vector;
+            icm::vector<nd::array> data_vector;
             data_vector.reserve(static_cast<size_t>(nrows));
             for (int i = 0; i < nrows; ++i) {
                 data_vector.emplace_back(nd::adapt(std::vector<T>(data + static_cast<size_t>(i) * static_cast<size_t>(ncols),
@@ -564,7 +565,7 @@ inline nd::array datum_to_nd(Datum value, Oid attr_typeid, int32_t typmod)
             return nd::none(nd::dtype::byte, 0);
         } else {
             int nelems = ArrayGetNItems(ARR_NDIM(arr), ARR_DIMS(arr));
-            std::vector<nd::array> elements;
+            icm::vector<nd::array> elements;
             elements.reserve(static_cast<size_t>(nelems));
 
             Datum* datums = nullptr;
@@ -590,7 +591,7 @@ inline nd::array datum_to_nd(Datum value, Oid attr_typeid, int32_t typmod)
             return nd::none(nd::dtype::string, 1);
         } else {
             int nelems = ArrayGetNItems(ARR_NDIM(arr), ARR_DIMS(arr));
-            std::vector<nd::array> elements;
+            icm::vector<nd::array> elements;
             elements.reserve(static_cast<size_t>(nelems));
             Datum* datums = nullptr;
             bool* nulls = nullptr;
@@ -903,7 +904,7 @@ nd::array eval_with_nones(nd::array arr)
         return nd::eval(arr);
     } catch (const nd::invalid_dynamic_eval&) {
     }
-    std::vector<nd::array> result_elements;
+    icm::vector<nd::array> result_elements;
     result_elements.reserve(arr.size());
     for (auto a : arr) {
         if (a.is_none()) {
diff --git a/cpp/deeplake_pg/sync_worker.cpp b/cpp/deeplake_pg/sync_worker.cpp
@@ -67,8 +67,8 @@ void deeplake_sync_worker_sighup(SIGNAL_ARGS)
  */
 void deeplake_sync_tables_from_catalog(const std::string& root_path, icm::string_map<> creds)
 {
-    auto catalog_tables = pg::dl_catalog::load_tables(root_path, creds);
-    auto catalog_columns = pg::dl_catalog::load_columns(root_path, creds);
+    // Load tables and columns in parallel for better performance
+    auto [catalog_tables, catalog_columns] = pg::dl_catalog::load_tables_and_columns(root_path, creds);
 
     for (const auto& meta : catalog_tables) {
         // Skip tables marked as dropping
@@ -168,6 +168,8 @@ PGDLLEXPORT void deeplake_sync_worker_main(Datum main_arg)
     elog(LOG, "pg_deeplake sync worker started");
 
     int64_t last_catalog_version = 0;
+    std::string last_root_path;  // Track root_path to detect changes
+    bool catalog_ensured = false;
 
     while (!got_sigterm) {
         // Handle SIGHUP - reload configuration
@@ -199,10 +201,15 @@ PGDLLEXPORT void deeplake_sync_worker_main(Datum main_arg)
             if (!root_path.empty()) {
                 auto creds = pg::session_credentials::get_credentials();
 
-                // Ensure catalog exists
-                pg::dl_catalog::ensure_catalog(root_path, creds);
+                // Only ensure catalog on first call or when root_path changes
+                if (!catalog_ensured || root_path != last_root_path) {
+                    pg::dl_catalog::ensure_catalog(root_path, creds);
+                    catalog_ensured = true;
+                    last_root_path = root_path;
+                    last_catalog_version = 0;  // Reset version when path changes
+                }
 
-                // Use existing catalog version API to check for changes
+                // Use existing catalog version API to check for changes (now fast with cache)
                 int64_t current_version = pg::dl_catalog::get_catalog_version(root_path, creds);
 
                 if (current_version != last_catalog_version) {
diff --git a/cpp/deeplake_pg/table_data.hpp b/cpp/deeplake_pg/table_data.hpp
@@ -6,6 +6,7 @@
 #include <base/spin_lock.hpp>
 #include <bifrost/column_streamer.hpp>
 #include <deeplake_api/dataset.hpp>
+#include <icm/vector.hpp>
 #include <nd/array.hpp>
 #include <nd/string_array_holder.hpp>
 
@@ -165,16 +166,16 @@ struct table_data
     constexpr static int64_t batch_mask_ = batch_size_ - 1;
 
     streamer_info streamers_;
-    icm::string_map<std::vector<nd::array>> insert_rows_;
+    icm::string_map<icm::vector<nd::array>> insert_rows_;
     std::deque<async::promise<void>> insert_promises_;
-    std::vector<int64_t> delete_rows_;
-    std::vector<std::tuple<int64_t, std::string, nd::array>> update_rows_;
+    icm::vector<int64_t> delete_rows_;
+    icm::vector<std::tuple<int64_t, std::string, nd::array>> update_rows_;
     std::shared_ptr<deeplake_api::dataset> dataset_;
     std::shared_ptr<deeplake_api::dataset> refreshing_dataset_;
     async::promise<void> refresh_promise_;
-    std::vector<bool> requested_columns_;
-    std::vector<Oid> base_typeids_;              // Cached base type OIDs for performance
-    std::vector<int32_t> active_column_indices_; // Maps logical index to TupleDesc index (excludes dropped)
+    icm::vector<bool> requested_columns_;
+    icm::vector<Oid> base_typeids_;              // Cached base type OIDs for performance
+    icm::vector<int32_t> active_column_indices_; // Maps logical index to TupleDesc index (excludes dropped)
     icm::string_map<> creds_;
     TupleDesc tuple_descriptor_;
     http::uri dataset_path_ = http::uri(std::string());
diff --git a/cpp/deeplake_pg/table_data_impl.hpp b/cpp/deeplake_pg/table_data_impl.hpp
@@ -499,7 +499,7 @@ inline bool table_data::flush_updates()
     // Flush the update rows to the dataset
     try {
         streamers_.reset();
-        std::vector<async::promise<void>> update_promises;
+        icm::vector<async::promise<void>> update_promises;
         update_promises.reserve(update_rows_.size());
         for (const auto& [row_number, column_name, new_value] : update_rows_) {
             update_promises.emplace_back(get_dataset()->update_row(row_number, column_name, new_value));
diff --git a/cpp/deeplake_pg/table_storage.cpp b/cpp/deeplake_pg/table_storage.cpp
@@ -291,22 +291,29 @@ void table_storage::load_table_metadata()
 
     // Stateless catalog sync (only when enabled)
     if (pg::stateless_enabled) {
-        pg::dl_catalog::ensure_catalog(root_dir, creds);
-
+        // Fast path: if already loaded, just check version without ensure_catalog
         if (tables_loaded_) {
             const auto current_version = pg::dl_catalog::get_catalog_version(root_dir, creds);
             if (current_version == catalog_version_) {
                 return;
             }
+            // Version changed, need to reload
             tables_.clear();
             views_.clear();
             tables_loaded_ = false;
+            catalog_version_ = current_version; // Reuse the version we just fetched
         }
+
+        // Only ensure catalog exists when we need to load/reload
+        pg::dl_catalog::ensure_catalog(root_dir, creds);
         tables_loaded_ = true;
-        catalog_version_ = pg::dl_catalog::get_catalog_version(root_dir, creds);
+        // Only fetch version if we don't already have it from the check above
+        if (catalog_version_ == 0) {
+            catalog_version_ = pg::dl_catalog::get_catalog_version(root_dir, creds);
+        }
 
-        auto catalog_tables = pg::dl_catalog::load_tables(root_dir, creds);
-        auto catalog_columns = pg::dl_catalog::load_columns(root_dir, creds);
+        // Load tables and columns in parallel for better performance
+        auto [catalog_tables, catalog_columns] = pg::dl_catalog::load_tables_and_columns(root_dir, creds);
         if (!catalog_tables.empty()) {
             for (const auto& meta : catalog_tables) {
                 const std::string qualified_name = meta.schema_name + "." + meta.table_name;
diff --git a/cpp/deeplake_pg/table_storage.hpp b/cpp/deeplake_pg/table_storage.hpp
@@ -183,6 +183,7 @@ class table_storage
     void force_load_table_metadata()
     {
         tables_loaded_ = false;
+        catalog_version_ = 0;  // Reset so version gets re-fetched for new root_path
         load_table_metadata();
     }
     void mark_metadata_stale() noexcept
diff --git a/postgres/tests/py_tests/test_startup_latency.py b/postgres/tests/py_tests/test_startup_latency.py
diff --git a/postgres/tests/py_tests/test_stateless_catalog_resilience.py b/postgres/tests/py_tests/test_stateless_catalog_resilience.py
diff --git a/postgres/tests/py_tests/test_stateless_multi_instance.py b/postgres/tests/py_tests/test_stateless_multi_instance.py
diff --git a/postgres/tests/py_tests/test_stateless_reserved_schema.py b/postgres/tests/py_tests/test_stateless_reserved_schema.py

Original file line number	Diff line number	Diff line change
`@@ -1197,10 +1197,18 @@ static void process_utility(PlannedStmt* pstmt,`
`1197`	`1197`	`}`
`1198`	`1198`	`// When root_path is set, auto-discover tables from the deeplake catalog`
`1199`	`1199`	`if (vstmt->name != nullptr && pg_strcasecmp(vstmt->name, "deeplake.root_path") == 0) {`
`1200`		`- // Reload table metadata from the catalog at the new root_path`
`1201`		`- // This enables stateless multi-instance support where tables are`
`1202`		`- // auto-discovered when pointing to a shared root_path`
`1203`		`- pg::table_storage::instance().force_load_table_metadata();`
	`1200`	`+ // Track the previous root_path to detect actual changes`
	`1201`	`+ static thread_local std::string last_root_path;`
	`1202`	`+ auto current_root_path = pg::session_credentials::get_root_path();`
	`1203`	`+`
	`1204`	`+ if (current_root_path != last_root_path) {`
	`1205`	`+ // Path changed - force full reload`
	`1206`	`+ last_root_path = current_root_path;`
	`1207`	`+ pg::table_storage::instance().force_load_table_metadata();`
	`1208`	`+ } else {`
	`1209`	`+ // Same path - just check for catalog updates (fast path)`
	`1210`	`+ pg::table_storage::instance().load_table_metadata();`
	`1211`	`+ }`
`1204`	`1212`	`}`
`1205`	`1213`	`}`
`1206`	`1214`	`}`
Original file line number	Diff line number	Diff line change
`@@ -108,7 +108,7 @@ struct scan_opaque`
`108`	`108`
`109`	`109`	`query_core::query_result run_index_search(nd::array input_array, std::string func_name, const std::string& column_name, pg::index_info& idx_info)`
`110`	`110`	`{`
`111`		`- std::vector<query_core::expr> args;`
	`111`	`+ icm::vector<query_core::expr> args;`
`112`	`112`	`args.emplace_back(query_core::expr::make_column_ref(column_name, std::string{}));`
`113`	`113`	`args.emplace_back(query_core::expr::make_literal_array(std::move(input_array)));`
`114`	`114`	`const bool is_cosine_similarity = (func_name == "COSINE_SIMILARITY");`
`@@ -194,7 +194,7 @@ icm::roaring run_exact_text_search(std::string text_value, StrategyNumber strate`
`194`	`194`	`return {};`
`195`	`195`	`}`
`196`	`196`	`info.column_name = idx_info.column_name();`
`197`		`- info.search_values.emplace_back(std::vector<std::string>{std::move(text_value)});`
	`197`	`+ info.search_values.emplace_back(icm::vector<std::string>{std::move(text_value)});`
`198`	`198`	`return idx_info.run_query(std::move(info));`
`199`	`199`	`}`
`200`	`200`
`@@ -563,7 +563,7 @@ void collect_index_data(IndexScanDesc scan, ScanKey keys, int32_t nkeys, ScanKey`
`563`	`563`	`}`
`564`	`564`	`}`
`565`	`565`	`if (nkeys > 0) {`
`566`		`- std::vector<int64_t> row_numbers;`
	`566`	`+ icm::vector<int64_t> row_numbers;`
`567`	`567`	`row_numbers.reserve(result.cardinality());`
`568`	`568`	`std::transform(result.begin(), result.end(), std::back_inserter(row_numbers), [](auto v) {`
`569`	`569`	`return static_cast<int64_t>(v);`