Merge pull request #3134 from activeloopai/drop-column-fixes

Fixed out of bounds fix after drop column.

Author: khustup2

cpp/deeplake_pg/column_statistics.cpp

@@ -218,20 +218,10 @@ bool inject_column_statistics(Relation rel, int16_t attnum)
         return false;
     }
 
-    // Get DeepLake column view - attnum is 1-based, column index is 0-based
-    int32_t col_idx = attnum - 1;
-
-    // Skip dropped columns by finding the actual column index
-    int32_t logical_idx = 0;
-    for (int32_t i = 0; i < tupdesc->natts && logical_idx <= col_idx; ++i) {
-        Form_pg_attribute a = TupleDescAttr(tupdesc, i);
-        if (!a->attisdropped) {
-            if (logical_idx == col_idx) {
-                col_idx = i;
-                break;
-            }
-            logical_idx++;
-        }
+    // Map PG attnum to logical column index (handles dropped columns correctly)
+    const auto col_idx = table_data.logical_index_for_attnum(attnum);
+    if (col_idx < 0) {
+        return false;
     }
 
     heimdall::column_view_ptr column_view;

cpp/deeplake_pg/deeplake_executor.cpp

@@ -77,7 +77,10 @@ void analyze_plan(PlannedStmt* plan)
             if (attnum <= 0) { // Only positive attribute numbers are real columns
                 continue;
             }
-            auto col_idx = static_cast<int32_t>(attnum - 1);
+            auto col_idx = table_data->logical_index_for_attnum(attnum);
+            if (col_idx < 0) {
+                continue; // Dropped column or out of range
+            }
             if (!table_data->is_column_requested(col_idx)) {
                 table_data->set_column_requested(col_idx, true);
                 if (!table_data->column_has_streamer(col_idx) && table_data->can_stream_column(col_idx)) {

cpp/deeplake_pg/table_am.cpp

@@ -367,10 +367,13 @@ double deeplake_index_build_range_scan(Relation heap_rel,
     AttrNumber* indexkeys = index_info->ii_IndexAttrNumbers;
     const auto table_id = RelationGetRelid(heap_rel);
     auto& td = pg::table_storage::instance().get_table_data(table_id);
+    // Map index key attnums to logical column indices
+    std::vector<int32_t> key_logical_indices(nkeys, -1);
     for (int32_t i = 0; i < nkeys; ++i) {
-        int32_t attnum = indexkeys[i] - 1;
-        if (attnum >= 0 && !td.column_has_streamer(attnum) && td.can_stream_column(attnum)) {
-            td.create_streamer(attnum, -1);
+        auto logical = td.logical_index_for_attnum(indexkeys[i]);
+        key_logical_indices[i] = logical;
+        if (logical >= 0 && !td.column_has_streamer(logical) && td.can_stream_column(logical)) {
+            td.create_streamer(logical, -1);
         }
     }
     std::vector<Datum> values(nkeys, 0);
@@ -382,12 +385,12 @@ double deeplake_index_build_range_scan(Relation heap_rel,
         auto [block_number, offset_number] = pg::utils::row_number_to_tid(row);
         ItemPointerSet(&tid, block_number, offset_number);
         for (int32_t i = 0; i < nkeys; ++i) {
-            int32_t attnum = indexkeys[i] - 1;
-            if (attnum < 0) [[unlikely]] {
+            auto logical = key_logical_indices[i];
+            if (logical < 0) [[unlikely]] {
                 nulls[i] = true;
                 values[i] = 0;
             } else [[likely]] {
-                auto [value, null] = tscan.get_datum(attnum, row);
+                auto [value, null] = tscan.get_datum(logical, row);
                 values[i] = value;
                 nulls[i] = null;
             }

cpp/deeplake_pg/table_data.hpp

@@ -67,6 +67,7 @@ struct table_data
     inline std::string get_atttypename(AttrNumber attr_num) const noexcept;
     inline bool is_column_dropped(AttrNumber attr_num) const noexcept;
     inline int32_t get_tupdesc_index(AttrNumber attr_num) const noexcept;
+    inline int32_t logical_index_for_attnum(int32_t attnum) const noexcept;
     inline bool is_column_nullable(AttrNumber attr_num) const noexcept;
     inline bool is_column_indexed(AttrNumber attr_num) const noexcept;
     inline int32_t num_columns() const noexcept;
@@ -176,6 +177,7 @@ struct table_data
     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::vector<int32_t> tupdesc_to_logical_;    // Maps TupleDesc index to logical index (-1 if dropped)
     icm::string_map<> creds_;
     TupleDesc tuple_descriptor_;
     http::uri dataset_path_ = http::uri(std::string());

cpp/deeplake_pg/table_data_impl.hpp

@@ -64,6 +64,12 @@ inline table_data::table_data(
         }
     }
 
+    // Build reverse mapping: TupleDesc index → logical index (-1 for dropped)
+    tupdesc_to_logical_.resize(tuple_descriptor_->natts, -1);
+    for (int32_t logical = 0; logical < static_cast<int32_t>(active_column_indices_.size()); ++logical) {
+        tupdesc_to_logical_[active_column_indices_[logical]] = logical;
+    }
+
     const auto num_active = active_column_indices_.size();
     requested_columns_.resize(num_active, false);
     base_typeids_.resize(num_active);
@@ -272,6 +278,15 @@ inline int32_t table_data::get_tupdesc_index(AttrNumber attr_num) const noexcept
     return active_column_indices_[attr_num];
 }
 
+inline int32_t table_data::logical_index_for_attnum(int32_t attnum) const noexcept
+{
+    const auto tupdesc_idx = attnum - 1;
+    if (tupdesc_idx < 0 || tupdesc_idx >= static_cast<int32_t>(tupdesc_to_logical_.size())) {
+        return -1;
+    }
+    return tupdesc_to_logical_[tupdesc_idx];
+}
+
 inline bool table_data::is_column_indexed(AttrNumber attr_num) const noexcept
 {
     return pg::pg_index::get_oid(table_name_, get_atttypename(attr_num)) != InvalidOid;
@@ -318,13 +333,14 @@ inline void table_data::add_insert_slots(int32_t nslots, TupleTableSlot** slots)
     for (int32_t i = 0; i < num_columns(); ++i) {
         auto& column_values = insert_rows_[get_atttypename(i)];
         const auto dt = get_column_view(i)->dtype();
+        const auto slot_pos = get_tupdesc_index(i);
         for (int32_t k = 0; k < nslots; ++k) {
             auto slot = slots[k];
             nd::array val;
-            if (slot->tts_isnull[i]) {
+            if (slot->tts_isnull[slot_pos]) {
                 val = (nd::dtype_is_numeric(dt) ? nd::adapt(0) : nd::none(dt, 0));
             } else {
-                val = pg::utils::datum_to_nd(slot->tts_values[i], get_base_atttypid(i), get_atttypmod(i));
+                val = pg::utils::datum_to_nd(slot->tts_values[slot_pos], get_base_atttypid(i), get_atttypmod(i));
             }
             column_values.push_back(std::move(val));
         }

cpp/deeplake_pg/table_scan_impl.hpp

@@ -158,19 +158,23 @@ inline std::pair<Datum, bool> table_scan::get_datum(int32_t column_number, int64
 inline void table_scan::convert_nd_to_pg(int64_t row_number, Datum* values, bool* nulls) const noexcept
 {
     for (auto col : null_columns_) {
-        nulls[col] = true;
+        const auto slot_pos = table_data_.get_tupdesc_index(col);
+        nulls[slot_pos] = true;
     }
     for (auto col : scored_columns_) {
-        nulls[col] = false;
+        const auto slot_pos = table_data_.get_tupdesc_index(col);
+        nulls[slot_pos] = false;
     }
     for (auto col : special_columns_) {
-        values[col] = pg::utils::make_special_datum(table_id_, row_number, col, table_data_.get_base_atttypid(col));
-        nulls[col] = false;
+        const auto slot_pos = table_data_.get_tupdesc_index(col);
+        values[slot_pos] = pg::utils::make_special_datum(table_id_, row_number, col, table_data_.get_base_atttypid(col));
+        nulls[slot_pos] = false;
     }
     for (auto col : process_columns_) {
+        const auto slot_pos = table_data_.get_tupdesc_index(col);
         auto [datum, is_null] = get_datum(col, row_number);
-        values[col] = datum;
-        nulls[col] = is_null;
+        values[slot_pos] = datum;
+        nulls[slot_pos] = is_null;
     }
 }
 

cpp/deeplake_pg/table_storage.cpp

@@ -951,10 +951,15 @@ void table_storage::create_table(const std::string& table_name, Oid table_id, Tu
 
 void table_storage::drop_table(const std::string& table_name)
 {
-    pg::table_ddl_lock_guard ddl_lock;
+    // Load metadata BEFORE acquiring the DDL lock.
+    // force_load_table_metadata() may trigger CREATE TABLE (via SPI) for tables
+    // in the S3 catalog that don't exist in pg_class yet.  CREATE TABLE goes
+    // through the table AM which also acquires the DDL lock — doing that while
+    // we already hold it would self-deadlock (LWLocks are not recursive).
     if (!table_exists(table_name)) {
         force_load_table_metadata();
     }
+    pg::table_ddl_lock_guard ddl_lock;
     if (table_exists(table_name)) {
         auto& table_data = get_table_data(table_name);
         auto creds = session_credentials::get_credentials();