Fixed rows count in pg-deeplake.

cpp/deeplake_pg/table_am.cpp

@@ -5,19 +5,24 @@ extern "C" {
 // Must be first to avoid macro conflicts
 #include <postgres.h>
 
+#include <access/multixact.h>
 #include <access/parallel.h>
 #include <access/reloptions.h> // For relation options
 #include <access/tableam.h>
 #include <access/xact.h>
 #include <catalog/namespace.h>
 #include <catalog/storage.h>
+#include <catalog/storage_xlog.h>
+#include <storage/smgr.h>
+#include <storage/bufmgr.h>
 #include <commands/vacuum.h> // For VacuumParams and VACOPT_* flags
 #include <miscadmin.h>
 #include <nodes/bitmapset.h>  // For bitmap operations
 #include <nodes/nodes.h>      // For node types
 #include <nodes/parsenodes.h> // For parse nodes
 #include <pgstat.h>           // For statistics collector integration
 #include <storage/block.h>
+#include <storage/read_stream.h>
 #include <storage/relfilelocator.h>
 #include <utils/builtins.h> // For text conversion functions
 #include <utils/lsyscache.h>
@@ -97,7 +102,7 @@ struct DeeplakeScanData
         , scan_state(std::move(state))
         , memory_context(AllocSetContextCreate(CurrentMemoryContext, "scan_context", ALLOCSET_SMALL_SIZES))
     {
-        const auto num_rows = scan_state.get_table_data().num_rows();
+        const auto num_rows = scan_state.get_table_data().num_total_rows();
         print_progress = pg::print_progress_during_seq_scan && num_rows > 100000;
         if (print_progress) {
             progress_bar = pg::utils::progress_display(num_rows,
@@ -176,15 +181,17 @@ bool deeplake_relation_needs_toast_table(Relation)
     return false;
 }
 
-uint64_t deeplake_relation_size(Relation rel, ForkNumber)
+uint64_t deeplake_relation_size(Relation rel, ForkNumber fork)
 {
-    auto table_id = RelationGetRelid(rel);
-    if (pg::table_storage::instance().table_exists(table_id)) {
-        auto& table_data = pg::table_storage::instance().get_table_data(table_id);
-        auto total_bytes = heimdall::dataset_total_bytes(*table_data.get_read_only_dataset());
-        return total_bytes / BLCKSZ;
-    }
-    return BLCKSZ; // Default block size in bytes
+    // For ANALYZE to work correctly with PostgreSQL's block-based sampling,
+    // we must return a size that matches the actual physical storage file.
+    // PostgreSQL's BlockSampler uses this to decide which blocks to sample,
+    // and read_stream_next_buffer will fail if we report more blocks than exist.
+    //
+    // We always return exactly BLCKSZ (1 block) because that's what we create
+    // in relation_set_new_node. The planner uses relation_estimate_size for
+    // accurate row estimates, which is set correctly via deeplake_estimate_rel_size.
+    return BLCKSZ;
 }
 
 void deeplake_index_validate_scan(Relation heap_rel,
@@ -220,8 +227,12 @@ void deeplake_estimate_rel_size(
     auto table_id = RelationGetRelid(rel);
     if (pg::table_storage::instance().table_exists(table_id)) {
         auto& table_data = pg::table_storage::instance().get_table_data(table_id);
+        // Refresh to ensure we see the latest data (including recently inserted rows)
+        const_cast<pg::table_data&>(table_data).refresh();
+        // Use num_total_rows() which includes uncommitted/staged data
+        auto total_rows = table_data.num_total_rows();
         if (tuples != nullptr) {
-            *tuples = static_cast<double>(table_data.num_rows());
+            *tuples = static_cast<double>(total_rows);
         }
         if (allvisfrac != nullptr) {
             *allvisfrac = 1.0; // Assume all tuples are visible
@@ -239,7 +250,7 @@ void deeplake_estimate_rel_size(
         if (pages != nullptr) {
             constexpr uint32_t min_pages = 1;
             const uint32_t num_blocks =
-                static_cast<uint32_t>(std::ceil(static_cast<double>(table_data.num_rows()) / 65536.0));
+                static_cast<uint32_t>(std::ceil(static_cast<double>(total_rows) / 65536.0));
             *pages = std::max(min_pages, num_blocks);
         }
 
@@ -270,7 +281,6 @@ bool deeplake_scan_analyze_next_tuple(
         return false;
     }
 
-    // Try to fetch next tuple from your columnar storage
     if (!scan_data->scan_state.get_next_tuple(slot)) {
         return false; // no more tuples
     }
@@ -298,14 +308,33 @@ bool deeplake_scan_analyze_next_block(TableScanDesc scan, ReadStream* stream)
         return false;
     }
 
-    // For columnar storage, treat the entire table as one logical "block"
-    // Return true on first call, false on subsequent calls
-    if (!scan_data->analyze_block_returned) {
-        scan_data->analyze_block_returned = true;
-        return true;
+    // PostgreSQL's ANALYZE formula requires:
+    //   totalrows = (liverows / bs.m) * totalblocks
+    //
+    // Where bs.m is tracked by the BlockSampler. To properly set bs.m,
+    // we must consume blocks from the stream via read_stream_next_buffer.
+    // Each call to read_stream_next_buffer triggers the stream's callback
+    // (block_sampling_read_stream_next) which calls BlockSampler_Next,
+    // incrementing bs.m.
+    //
+    // For columnar storage, we have a minimal physical storage file
+    // (created in relation_set_new_node) that allows the stream to work.
+    // We consume ONE block per call to this function, and return true
+    // to indicate there's data to process. scan_analyze_next_tuple will
+    // then return all rows from our columnar storage.
+
+    Buffer buf = read_stream_next_buffer(stream, NULL);
+
+    if (!BufferIsValid(buf)) {
+        return false;  // No more blocks to sample
     }
 
-    return false;
+    // Release the buffer immediately - we don't actually need the physical data
+    // because our columnar storage provides the data via scan_analyze_next_tuple.
+    // But we needed to consume the stream to increment bs.m.
+    ReleaseBuffer(buf);
+
+    return true;  // Indicate we have data to process
 }
 #endif
 
@@ -334,7 +363,7 @@ double deeplake_index_build_range_scan(Relation heap_rel,
     std::vector<Datum> values(nkeys, 0);
     std::vector<uint8_t> nulls(nkeys, 0);
     pg::table_scan tscan(table_id, false, false);
-    const auto num_rows = td.num_rows();
+    const auto num_rows = td.num_total_rows();
     ItemPointerData tid;
     for (auto row = 0; row < num_rows; ++row) {
         auto [block_number, offset_number] = pg::utils::row_number_to_tid(row);
@@ -637,10 +666,14 @@ TableScanDesc deeplake_table_am_routine::scan_begin(Relation relation,
     auto table_id = RelationGetRelid(relation);
     bool is_parallel = (pg::use_parallel_workers && parallel_scan != nullptr);
 
+    // Refresh table data BEFORE creating the scan to ensure num_rows is up to date
+    // (table_scan constructor caches num_rows at construction time)
+    auto& td_for_refresh = pg::table_storage::instance().get_table_data(table_id);
+    const_cast<pg::table_data&>(td_for_refresh).refresh();
+
     // Initialize extended structure with embedded scan data
     new (extended_scan)
         DeeplakeScanData(table_scan(table_id, is_parallel, query_info::current().receiver_registered()));
-    const_cast<pg::table_data&>(extended_scan->scan_state.get_table_data()).refresh();
 
     TableScanDesc scan_desc = &extended_scan->postgres_scan;
     scan_desc->rs_rd = relation;
@@ -704,7 +737,7 @@ void deeplake_table_am_routine::scan_rescan(TableScanDesc scan,
             std::string pre_msg = "Progress of sequential scan for table '" +
                                   scan_data->scan_state.get_table_data().get_table_name() + "'" + " (rescan " +
                                   std::to_string(scan_data->num_rescans) + ")";
-            scan_data->progress_bar.restart(scan_data->scan_state.get_table_data().num_rows(), std::move(pre_msg));
+            scan_data->progress_bar.restart(scan_data->scan_state.get_table_data().num_total_rows(), std::move(pre_msg));
         }
     }
 }
@@ -1081,6 +1114,31 @@ void deeplake_table_am_routine::relation_set_new_node(
                  "Deeplake uses columnar storage which is already compact. Use VACUUM or VACUUM ANALYZE instead.")));
     }
 
+    *freezeXid = RecentXmin;
+    *minmulti = GetOldestMultiXactId();
+
+    // Create physical storage for the relation.
+    // Even though DeepLake uses columnar storage (S3/cloud), PostgreSQL's ANALYZE
+    // requires a physical storage file for block sampling. The read_stream API
+    // expects to read physical blocks from the relation's storage file.
+    // Without this, ANALYZE fails with "could not open file" errors.
+    SMgrRelation srel = RelationCreateStorage(*newrnode, persistence, true);
+
+    // Extend the storage to have at least one block.
+    // This is needed because ANALYZE's BlockSampler will try to read blocks,
+    // and if the file is empty, read_stream_next_buffer will fail.
+    // We create a minimal block so that ANALYZE can proceed.
+    // The actual data comes from our columnar storage via scan_analyze_next_tuple.
+    smgrzeroextend(srel, MAIN_FORKNUM, 0, 1, true);
+
+    if (persistence == RELPERSISTENCE_UNLOGGED) {
+        smgrcreate(srel, INIT_FORKNUM, false);
+        log_smgrcreate(&srel->smgr_rlocator.locator, INIT_FORKNUM);
+        smgrimmedsync(srel, INIT_FORKNUM);
+    }
+
+    smgrclose(srel);
+
     // Get the tuple descriptor
     TupleDesc tupdesc = RelationGetDescr(rel);
     if (tupdesc == nullptr) {

cpp/deeplake_pg/table_data_impl.hpp

@@ -164,6 +164,9 @@ inline void table_data::refresh()
             })
                 .get_future()
                 .get();
+            // After refresh from version change (another backend committed),
+            // use the dataset's actual row count. This correctly handles both
+            // inserts and deletes from other backends.
             num_total_rows_ = dataset_->num_rows();
         }
     } else {
@@ -178,6 +181,9 @@ inline void table_data::refresh()
                 std::swap(dataset_, refreshing_dataset_);
                 dataset_->set_indexing_mode(ds_indexing_mode);
                 refreshing_dataset_->set_indexing_mode(deeplake::indexing_mode::off);
+                // After refresh from version change (another backend committed),
+                // use the dataset's actual row count. This correctly handles both
+                // inserts and deletes from other backends.
                 num_total_rows_ = dataset_->num_rows();
             }
         }
@@ -298,6 +304,12 @@ inline int64_t table_data::num_rows() const noexcept
 
 inline int64_t table_data::num_total_rows() const noexcept
 {
+    // If num_total_rows_ is 0, the dataset may not have been opened yet
+    // (e.g., when table_data was loaded from metadata after RENAME COLUMN).
+    // In this case, open the dataset to get the correct row count.
+    if (num_total_rows_ == 0 && dataset_ == nullptr) {
+        const_cast<table_data*>(this)->open_dataset();
+    }
     return num_total_rows_;
 }
 
@@ -473,6 +485,8 @@ inline bool table_data::flush_deletes()
         return true;
     }
 
+    const auto num_deletes = static_cast<int64_t>(delete_rows_.size());
+
     // Flush the delete rows to the dataset
     try {
         streamers_.reset();
@@ -484,6 +498,10 @@ inline bool table_data::flush_deletes()
     }
 
     delete_rows_.clear();
+
+    // Update the total row count to reflect deleted rows
+    num_total_rows_ -= num_deletes;
+
     return true;
 }
 
@@ -546,7 +564,7 @@ inline Oid table_data::get_table_oid() const noexcept
 inline std::pair<int64_t, int64_t> table_data::get_row_range(int32_t worker_id) const
 {
     ASSERT(worker_id >= 0 && worker_id < max_parallel_workers);
-    const auto total_rows = num_rows();
+    const auto total_rows = num_total_rows();
     const auto total_workers = max_parallel_workers;
     int64_t rows_per_worker = total_rows / total_workers;
     int64_t remaining_rows = total_rows % total_workers;
@@ -570,7 +588,7 @@ inline void table_data::create_streamer(int32_t idx, int32_t worker_id)
         return;
     }
     if (pg::memory_tracker::has_memory_limit()) {
-        const auto column_size = pg::utils::get_column_width(get_base_atttypid(idx), get_atttypmod(idx)) * num_rows();
+        const auto column_size = pg::utils::get_column_width(get_base_atttypid(idx), get_atttypmod(idx)) * num_total_rows();
         pg::memory_tracker::ensure_memory_available(column_size);
     }
     heimdall::column_view_ptr cv = get_column_view(idx);
@@ -579,7 +597,7 @@ inline void table_data::create_streamer(int32_t idx, int32_t worker_id)
         cv = heimdall_common::create_filtered_column(*(cv),
                                                      icm::index_mapping_t<int64_t>::slice({start_row, end_row, 1}));
     }
-    const int64_t row_count = num_rows();
+    const int64_t row_count = num_total_rows();
     const int64_t batch_count = (row_count + batch_size_ - 1) / batch_size_;
     column_batches = std::vector<streamer_info::batch_data>(batch_count);
     for (int64_t i = 0; i < batch_count; ++i) {

cpp/deeplake_pg/table_scan.hpp

@@ -49,6 +49,11 @@ class table_scan
         current_position_ = position;
     }
 
+    inline int64_t get_num_rows() const noexcept
+    {
+        return num_rows_;
+    }
+
     inline const auto& get_table_data() const noexcept
     {
         return table_data_;

cpp/deeplake_pg/table_scan_impl.hpp

@@ -33,7 +33,7 @@ namespace pg {
 inline table_scan::table_scan(Oid table_id, bool is_parallel, bool streamer_only)
     : table_data_(table_storage::instance().get_table_data(table_id))
     , current_position_(0)
-    , num_rows_(table_data_.num_rows())
+    , num_rows_(table_data_.num_total_rows())
     , table_id_(table_id)
     , is_parallel_(is_parallel)
 {
@@ -49,8 +49,6 @@ inline table_scan::table_scan(Oid table_id, bool is_parallel, bool streamer_only
             special_columns_.emplace_back(i);
         }
     }
-    for (int32_t i = 0; i < table_data_.num_columns(); ++i) {
-    }
     if (is_parallel_ || IsParallelWorker()) {
         int32_t worker_number = ParallelWorkerNumber + 1;
         auto [start_row, end_row] = table_data_.get_row_range(worker_number);

cpp/deeplake_pg/timing_guard.cpp

@@ -0,0 +1,61 @@
+#include "timing_guard.hpp"
+
+#include <fstream>
+#include <iomanip>
+#include <sstream>
+#include <thread>
+
+namespace base {
+
+std::ofstream& get_timing_log_file();
+std::chrono::high_resolution_clock::time_point get_session_start_time();
+// needs to be non-const as it's modified in timing_guard constructor/destructor
+// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
+thread_local int timing_indent_level = 0;
+
+std::chrono::high_resolution_clock::time_point get_session_start_time()
+{
+    thread_local auto session_start = std::chrono::high_resolution_clock::now();
+    return session_start;
+}
+
+std::ofstream& get_timing_log_file()
+{
+    thread_local std::ofstream log_file;
+    if (!log_file.is_open()) {
+        // Initialize session start time for this thread
+        get_session_start_time();
+
+        std::ostringstream filename;
+        filename << "/home/ubuntu/deeplake_timing_" << std::this_thread::get_id() << ".log";
+        log_file.open(filename.str());
+    }
+
+    return log_file;
+}
+
+timing_guard::timing_guard(const std::string& operation)
+    : operation_(operation)
+    , start_time_(std::chrono::high_resolution_clock::now())
+{
+    timing_indent_level++;
+}
+
+timing_guard::~timing_guard()
+{
+    auto end_time = std::chrono::high_resolution_clock::now();
+    auto duration_us = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time_).count();
+    auto start_timestamp_us =
+        std::chrono::duration_cast<std::chrono::microseconds>(start_time_ - get_session_start_time()).count();
+
+    constexpr int INDENT_SPACES_PER_LEVEL = 2;
+    constexpr int TIMESTAMP_WIDTH = 8;
+    std::string indent(static_cast<std::string::size_type>(timing_indent_level * INDENT_SPACES_PER_LEVEL), ' ');
+
+    get_timing_log_file() << "[" << std::setw(TIMESTAMP_WIDTH) << start_timestamp_us << "] " << indent << operation_ << " "
+                          << duration_us << '\n';
+
+    timing_indent_level--;
+}
+
+} // namespace base

cpp/deeplake_pg/timing_guard.hpp

@@ -0,0 +1,19 @@
+#pragma once
+
+#include <chrono>
+#include <string>
+
+namespace base {
+
+class timing_guard
+{
+public:
+    timing_guard(const std::string& operation);
+    ~timing_guard();
+
+private:
+    std::string operation_;
+    std::chrono::high_resolution_clock::time_point start_time_;
+};
+
+} // namespace base