Merge branch 'main' into ci-cd-optimize

Author: khustup2

cpp/deeplake_pg/extension_init.cpp

@@ -46,7 +46,6 @@ namespace pg {
 bool use_parallel_workers = false;
 bool use_deeplake_executor = true;
 bool explain_query_before_execute = false;
-bool ignore_primary_keys = true;
 bool print_runtime_stats = false;
 bool support_json_index = false;
 bool is_filter_pushdown_enabled = true;
@@ -133,17 +132,6 @@ void initialize_guc_parameters()
                              nullptr // check_hook, assign_hook, show_hook
     );
 
-    DefineCustomBoolVariable("pg_deeplake.ignore_primary_keys",
-                             "If set to true, PRIMARY KEY constraints will be ignored during table creation.",
-                             nullptr,                  // optional long description
-                             &pg::ignore_primary_keys, // linked C variable
-                             true,                     // default value
-                             PGC_USERSET,              // context (USERSET, SUSET, etc.)
-                             0,                        // flags
-                             nullptr,
-                             nullptr,
-                             nullptr // check_hook, assign_hook, show_hook
-    );
 
     DefineCustomBoolVariable("pg_deeplake.print_runtime_stats",
                              "Enable runtime statistics printing for pg_deeplake operations.",
@@ -625,70 +613,6 @@ static void process_utility(PlannedStmt* pstmt,
             list_free_deep(stmt->options);
             stmt->options = NIL;
         }
-        // Remove PRIMARY KEY constraints
-        if (pg::ignore_primary_keys && deeplake_table && stmt->tableElts != nullptr) {
-            List* new_table_elts = NIL;
-            ListCell* lc = nullptr;
-            bool has_primary_key = false;
-            // Get table name from the CreateStmt
-            std::string table_name = stmt->relation ? stmt->relation->relname : "";
-            std::map<std::string, std::set<std::string>> primary_keys;
-            foreach (lc, stmt->tableElts) {
-                Node* element = (Node*)lfirst(lc);
-                // Handle table-level PRIMARY KEY constraints
-                if (IsA(element, Constraint)) {
-                    Constraint* constraint = (Constraint*)element;
-                    if (constraint->contype == CONSTR_PRIMARY) {
-                        has_primary_key = true;
-                        // Extract primary key column names
-                        if (constraint->keys != nullptr) {
-                            ListCell* key_lc = nullptr;
-                            foreach (key_lc, constraint->keys) {
-                                std::string col_name = strVal(lfirst(key_lc));
-                                elog(DEBUG1,
-                                     "Removing table-level PRIMARY KEY constraint on table '%s' for column: %s",
-                                     table_name.c_str(),
-                                     col_name.c_str());
-                                primary_keys[table_name].insert(col_name);
-                            }
-                        }
-                        continue; // Skip adding this constraint
-                    }
-                }
-
-                // Handle column-level PRIMARY KEY constraints (e.g., "col INT PRIMARY KEY")
-                if (IsA(element, ColumnDef)) {
-                    ColumnDef* coldef = (ColumnDef*)element;
-                    if (coldef->constraints != nullptr) {
-                        List* new_constraints = NIL;
-                        ListCell* const_lc = nullptr;
-
-                        foreach (const_lc, coldef->constraints) {
-                            Constraint* constraint = (Constraint*)lfirst(const_lc);
-                            if (constraint->contype == CONSTR_PRIMARY) {
-                                has_primary_key = true;
-                                elog(DEBUG1,
-                                     "Removing column-level PRIMARY KEY constraint on table '%s' for column: %s",
-                                     table_name.c_str(),
-                                     coldef->colname);
-                                primary_keys[table_name].insert(coldef->colname);
-                                continue;
-                            }
-                            new_constraints = lappend(new_constraints, constraint);
-                        }
-
-                        // Update column's constraints list
-                        coldef->constraints = new_constraints;
-                    }
-                }
-
-                new_table_elts = lappend(new_table_elts, element);
-            }
-            if (has_primary_key) {
-                stmt->tableElts = new_table_elts;
-                pg::table_storage::instance().set_primary_keys(std::move(primary_keys));
-            }
-        }
     }
 
     std::optional<pg::utils::parallel_workers_switcher> switcher;

cpp/deeplake_pg/table_am.cpp

@@ -28,6 +28,7 @@ extern "C" {
 #include <utils/lsyscache.h>
 #include <utils/rel.h>
 #include <utils/relcache.h>
+#include <utils/snapshot.h> // For SNAPSHOT_DIRTY and snapshot types
 #include <utils/varlena.h> // For text functions
 
 #ifdef __cplusplus
@@ -966,13 +967,27 @@ bool deeplake_table_am_routine::index_fetch_tuple(struct IndexFetchTableData* sc
 
     pg::utils::memory_context_switcher context_switcher(idx_scan->memory_context);
     idx_scan->scan_state.set_current_position(utils::tid_to_row_number(tid));
+
     if (!idx_scan->scan_state.get_next_tuple(slot)) {
-        *all_dead = true;
+        if (all_dead != nullptr) {
+            *all_dead = true;
+        }
         return false;
     }
 
+    // For SNAPSHOT_DIRTY (used by btree unique checking),
+    // we need to indicate that no in-progress transaction is affecting this tuple.
+    // This prevents PostgreSQL from trying to look up transaction status in pg_subtrans.
+    // Deeplake tuples are always immediately visible (no MVCC).
+    if (snapshot != nullptr && snapshot->snapshot_type == SNAPSHOT_DIRTY) {
+        snapshot->xmin = InvalidTransactionId;
+        snapshot->xmax = InvalidTransactionId;
+    }
+
     *call_again = false;
-    *all_dead = false;
+    if (all_dead != nullptr) {
+        *all_dead = false;
+    }
     return true;
 }
 

cpp/deeplake_pg/table_data.hpp

@@ -77,7 +77,6 @@ struct table_data
     inline void clear_delete_rows() noexcept;
     inline void add_update_row(int64_t row_id, icm::string_map<nd::array> update_row);
     inline void clear_update_rows() noexcept;
-    inline void set_primary_keys(const std::set<std::string>& primary_keys);
     inline Oid get_table_oid() const noexcept;
     inline bool flush();
 

cpp/deeplake_pg/table_data_impl.hpp

@@ -530,32 +530,6 @@ inline bool table_data::flush_updates()
     return true;
 }
 
-inline void table_data::set_primary_keys(const std::set<std::string>& primary_keys)
-{
-    get_dataset()->set_indexing_mode(deeplake::indexing_mode::always);
-    bool index_created = false;
-    for (const auto& column_name : primary_keys) {
-        auto& column = get_dataset()->get_column(column_name);
-        auto index_holder = column.index_holder();
-        if (index_holder != nullptr) {
-            continue;
-        }
-        auto column_type_kind = column.type().kind();
-        if (column_type_kind == deeplake_core::type_kind::generic && !column.type().data_type().is_array() &&
-            nd::dtype_is_numeric(column.type().data_type().get_dtype())) {
-            column.create_index(deeplake_core::index_type(
-                deeplake_core::numeric_index_type(deeplake_core::deeplake_index_type::type::inverted_index)));
-            index_created = true;
-            elog(DEBUG1, "Created numeric index on table '%s' column '%s'", table_name_.c_str(), column_name.c_str());
-        } else {
-            elog(DEBUG1, "Column %s is not supported yet for indexing", column_name.c_str());
-        }
-    }
-    if (index_created) {
-        commit();
-    }
-}
-
 inline Oid table_data::get_table_oid() const noexcept
 {
     return table_oid_;

cpp/deeplake_pg/table_storage.cpp

@@ -677,14 +677,6 @@ void table_storage::create_table(const std::string& table_name, Oid table_id, Tu
         ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR), errmsg("%s", e.what())));
     }
 
-    if (auto it = primary_keys_.find(simple_table_name); it != primary_keys_.end() && !it->second.empty()) {
-        try {
-            td.set_primary_keys(it->second);
-        } catch (const base::exception& e) {
-            elog(WARNING, "Failed to set primary keys for table %s: %s", table_name.c_str(), e.what());
-        }
-    }
-
     tables_.emplace(table_id, std::move(td));
     up_to_date_ = false;
 }
@@ -792,7 +784,10 @@ bool table_storage::fetch_tuple(Oid table_id, ItemPointer tid, TupleTableSlot* s
         ExecClearTuple(slot);
 
         const auto row_number = utils::tid_to_row_number(tid);
-        if (row_number >= table_data.num_rows()) {
+        // Use num_total_rows() to include uncommitted rows in the current transaction.
+        // This is necessary for AFTER triggers (like FK checks) that need to see
+        // rows inserted earlier in the same transaction.
+        if (row_number >= table_data.num_total_rows()) {
             return false;
         }
         Datum* values = slot->tts_values;

cpp/deeplake_pg/table_storage.hpp

@@ -223,11 +223,6 @@ class table_storage
         up_to_date_ = up_to_date;
     }
 
-    inline void set_primary_keys(std::map<std::string, std::set<std::string>>&& primary_keys) noexcept
-    {
-        primary_keys_ = std::move(primary_keys);
-    }
-
 private:
     table_storage() = default;
 
@@ -236,7 +231,6 @@ class table_storage
 
     std::unordered_map<Oid, table_data> tables_;
     std::unordered_map<Oid, std::pair<std::string, std::string>> views_;
-    std::map<std::string, std::set<std::string>> primary_keys_;
     std::string schema_name_ = "public";
     bool tables_loaded_ = false;
     bool up_to_date_ = true;

cpp/deeplake_pg/utils.hpp

@@ -46,7 +46,6 @@ inline static constexpr const char* index_type_option_name = "index_type";
 extern bool use_parallel_workers;
 extern bool use_deeplake_executor;
 extern bool explain_query_before_execute;
-extern bool ignore_primary_keys;
 extern bool print_runtime_stats;
 extern bool is_filter_pushdown_enabled;
 extern int32_t max_streamable_column_width;

postgres/tests/py_tests/test_constraint_enforcement.py

@@ -0,0 +1,69 @@
+"""
+Tests for constraint enforcement in pg_deeplake.
+
+These tests verify that PRIMARY KEY, UNIQUE, and FOREIGN KEY constraints
+work correctly with deeplake tables.
+"""
+import pytest
+import asyncpg
+
+
+@pytest.mark.asyncio
+async def test_primary_key_rejects_duplicates(db_conn: asyncpg.Connection):
+    """PRIMARY KEY should reject duplicate values."""
+    try:
+        await db_conn.execute("DROP TABLE IF EXISTS pk_test CASCADE")
+        await db_conn.execute("""
+            CREATE TABLE pk_test (id INT PRIMARY KEY, name TEXT) USING deeplake
+        """)
+
+        await db_conn.execute("INSERT INTO pk_test VALUES (1, 'alice')")
+
+        with pytest.raises(asyncpg.UniqueViolationError):
+            await db_conn.execute("INSERT INTO pk_test VALUES (1, 'bob')")
+
+    finally:
+        await db_conn.execute("DROP TABLE IF EXISTS pk_test CASCADE")
+
+
+@pytest.mark.asyncio
+async def test_unique_constraint_rejects_duplicates(db_conn: asyncpg.Connection):
+    """UNIQUE constraint should reject duplicate values."""
+    try:
+        await db_conn.execute("DROP TABLE IF EXISTS unique_test CASCADE")
+        await db_conn.execute("""
+            CREATE TABLE unique_test (id INT PRIMARY KEY, email TEXT UNIQUE) USING deeplake
+        """)
+
+        await db_conn.execute("INSERT INTO unique_test VALUES (1, 'alice@test.com')")
+
+        with pytest.raises(asyncpg.UniqueViolationError):
+            await db_conn.execute("INSERT INTO unique_test VALUES (2, 'alice@test.com')")
+
+    finally:
+        await db_conn.execute("DROP TABLE IF EXISTS unique_test CASCADE")
+
+
+@pytest.mark.asyncio
+async def test_foreign_key_insert(db_conn: asyncpg.Connection):
+    """INSERT into child table with FK should trigger parent lookup."""
+    try:
+        await db_conn.execute("DROP TABLE IF EXISTS fk_child CASCADE")
+        await db_conn.execute("DROP TABLE IF EXISTS fk_parent CASCADE")
+
+        await db_conn.execute("""
+            CREATE TABLE fk_parent (id INT PRIMARY KEY) USING deeplake
+        """)
+        await db_conn.execute("""
+            CREATE TABLE fk_child (
+                id INT PRIMARY KEY,
+                parent_id INT REFERENCES fk_parent(id)
+            ) USING deeplake
+        """)
+
+        await db_conn.execute("INSERT INTO fk_parent VALUES (1)")
+        await db_conn.execute("INSERT INTO fk_child VALUES (1, 1)")
+
+    finally:
+        await db_conn.execute("DROP TABLE IF EXISTS fk_child CASCADE")
+        await db_conn.execute("DROP TABLE IF EXISTS fk_parent CASCADE")