infiniflow · JinHai-CN · May 15, 2025 · May 14, 2025 · May 14, 2025 · May 14, 2025
diff --git a/python/test_pysdk/test_update.py b/python/test_pysdk/test_update.py
@@ -266,6 +266,33 @@ def test_update_table_with_one_block(self, suffix):
         res = db_obj.drop_table("test_update_table_with_one_block"+suffix, ConflictType.Error)
         assert res.error_code == ErrorCode.OK
 
+    def test_update_table_with_two_blocks(self, suffix):
+        # connect
+        db_obj = self.infinity_obj.get_database("default_db")
+        db_obj.drop_table("test_update_table_with_two_blocks"+suffix, ConflictType.Ignore)
+        table_obj = db_obj.create_table("test_update_table_with_two_blocks"+suffix,
+                                        {"c1": {"type": "int"}, "c2": {"type": "int"}},
+                                        ConflictType.Error)
+
+        # insert
+        values = [{"c1": 1, "c2": 2} for _ in range(6000)]
+        table_obj.insert(values)
+        insert_res, extra_result = table_obj.output(["*"]).to_df()
+        print(insert_res)
+
+        values = [{"c1": 1, "c2": 2} for _ in range(5000)]
+        table_obj.insert(values)
+        insert_res, extra_result = table_obj.output(["*"]).to_df()
+        print(insert_res)
+
+        # update
+        table_obj.update("c1 = 1", {"c2": 20})
+        delete_res, extra_result = table_obj.output(["*"]).to_df()
+        print(delete_res)
+
+        res = db_obj.drop_table("test_update_table_with_two_blocks"+suffix, ConflictType.Error)
+        assert res.error_code == ErrorCode.OK
+
     def test_update_table_with_one_segment(self, suffix):
         # connect
         db_obj = self.infinity_obj.get_database("default_db")

diff --git a/src/executor/operator/physical_update.cpp b/src/executor/operator/physical_update.cpp
@@ -109,13 +109,7 @@ bool PhysicalUpdate::Execute(QueryContext *query_context, OperatorState *operato
             SharedPtr<DataBlock> output_data_block = DataBlock::Make();
             output_data_block->Init(output_column_vectors);
 
-            Status status = new_txn->Append(*table_info_->db_name_, *table_info_->table_name_, output_data_block);
-            if (!status.ok()) {
-                operator_state->status_ = status;
-                RecoverableError(status);
-                return false;
-            }
-            status = new_txn->Delete(*table_info_->db_name_, *table_info_->table_name_, row_ids);
+            Status status = new_txn->Update(*table_info_->db_name_, *table_info_->table_name_, output_data_block, row_ids);
             if (!status.ok()) {
                 operator_state->status_ = status;
                 RecoverableError(status);

diff --git a/src/storage/new_txn/base_txn_store.cpp b/src/storage/new_txn/base_txn_store.cpp
@@ -66,7 +66,13 @@ String DropIndexTxnStore::ToString() const {
 }
 
 String AppendTxnStore::ToString() const {
-    return fmt::format("{}: database: {}, db_id: {}, table: {}, table_id: {}", TransactionType2Str(type_), db_name_, db_id_, table_name_, table_id_);
+    return fmt::format("{}: database: {}, db_id: {}, table: {}, table_id: {}, appended: {}",
+                       TransactionType2Str(type_),
+                       db_name_,
+                       db_id_,
+                       table_name_,
+                       table_id_,
+                       row_ranges_.size());
 }
 
 String ImportTxnStore::ToString() const {
@@ -127,4 +133,16 @@ String DeleteTxnStore::ToString() const {
                        table_id_,
                        row_ids_.size());
 }
+
+String UpdateTxnStore::ToString() const {
+    return fmt::format("{}: database: {}, db_id: {}, table: {}, table_id: {}, appended: {}, deleted: {}",
+                       TransactionType2Str(type_),
+                       db_name_,
+                       db_id_,
+                       table_name_,
+                       table_id_,
+                       row_ranges_.size(),
+                       row_ids_.size());
+}
+
 } // namespace infinity
diff --git a/src/storage/new_txn/base_txn_store.cppm b/src/storage/new_txn/base_txn_store.cppm
@@ -276,4 +276,30 @@ export struct DeleteTxnStore : public BaseTxnStore {
 
     String ToString() const final;
 };
+
+export struct UpdateTxnStore : public BaseTxnStore {
+    UpdateTxnStore() : BaseTxnStore(TransactionType::kUpdate) {}
+
+    String db_name_{};
+    String db_id_str_{};
+    String table_name_{};
+    String table_id_str_{};
+    u64 db_id_{};
+    u64 table_id_{};
+
+    Vector<SharedPtr<DataBlock>> input_blocks_{};
+    Vector<String> index_ids_{}; // indexes will be appended
+
+    // For data append
+    Vector<Pair<RowID, u64>> row_ranges_{};
+
+    // For mem index
+    Vector<MemIndexRange> mem_indexes_to_append_{};
+    Vector<MemIndexRange> mem_indexes_to_dump_{};
+
+    Vector<RowID> row_ids_{};
+
+    String ToString() const final;
+};
+
 } // namespace infinity
diff --git a/src/storage/new_txn/new_txn.cpp b/src/storage/new_txn/new_txn.cpp
@@ -2047,6 +2047,9 @@ bool NewTxn::CheckConflictTxnStore(NewTxn *previous_txn, String &cause, bool &re
         case TransactionType::kDropColumn: {
             return CheckConflictTxnStore(static_cast<const DropColumnsTxnStore &>(*base_txn_store_), previous_txn, cause, retry_query);
         }
+        case TransactionType::kDropTable: {
+            return CheckConflictTxnStore(static_cast<const DropTableTxnStore &>(*base_txn_store_), previous_txn, cause, retry_query);
+        }
         default: {
             return false;
         }
@@ -3164,6 +3167,29 @@ bool NewTxn::CheckConflictCmd(const WalCmdDropTableV2 &cmd, NewTxn *previous_txn
     return false;
 }
 
+bool NewTxn::CheckConflictTxnStore(const DropTableTxnStore &txn_store, NewTxn *previous_txn, String &cause, bool &retry_query) {
+    const String &db_name = txn_store.db_name_;
+    bool conflict = false;
+    switch (previous_txn->base_txn_store_->type_) {
+        case TransactionType::kDropDB: {
+            DropDBTxnStore *drop_db_txn_store = static_cast<DropDBTxnStore *>(previous_txn->base_txn_store_.get());
+            if (drop_db_txn_store->db_name_ == db_name) {
+                retry_query = false;
+                conflict = true;
+            }
+            break;
+        }
+        default: {
+        }
+    }
+
+    if (conflict) {
+        cause = fmt::format("{} vs. {}", previous_txn->base_txn_store_->ToString(), txn_store.ToString());
+        return true;
+    }
+    return false;
+}
+
 bool NewTxn::CheckConflictTxnStore(const DumpMemIndexTxnStore &txn_store, NewTxn *previous_txn, String &cause, bool &retry_query) {
     const String &db_name = txn_store.db_name_;
     const String &table_name = txn_store.table_name_;

diff --git a/src/storage/new_txn/new_txn.cppm b/src/storage/new_txn/new_txn.cppm
@@ -93,7 +93,6 @@ struct AppendState;
 struct AppendRange;
 enum class DumpIndexCause;
 struct IndexReader;
-struct BaseTxnStore;
 struct TxnCommitterTask;
 
 struct BaseTxnStore;
@@ -107,6 +106,7 @@ struct CompactTxnStore;
 struct CreateIndexTxnStore;
 struct DumpMemIndexTxnStore;
 struct DeleteTxnStore;
+struct DropTableTxnStore;
 
 export struct CheckpointOption {
     TxnTimeStamp checkpoint_ts_ = 0;
@@ -334,11 +334,16 @@ public:
     Status Update(const String &db_name, const String &table_name, const SharedPtr<DataBlock> &input_block, const Vector<RowID> &row_ids);
 
 private:
+    Tuple<Vector<Pair<RowID, u64>>, Status> GetRowRanges(TableMeeta &table_meta, const SharedPtr<DataBlock> &input_block);
+
     Status AppendInner(const String &db_name,
                        const String &table_name,
                        const String &table_key,
                        TableMeeta &table_meta,
-                       const SharedPtr<DataBlock> &input_block);
+                       const SharedPtr<DataBlock> &input_block,
+                       const Vector<Pair<RowID, u64>> &row_ranges);
+
+    Status DeleteInner(const String &db_name, const String &table_name, TableMeeta &table_meta, const Vector<RowID> &row_ids);
 
 public:
     // Status Delete(TableEntry *table_entry, const Vector<RowID> &row_ids, bool check_conflict = true);
@@ -622,6 +627,7 @@ private:
     bool CheckConflictTxnStore(const CreateIndexTxnStore &txn_store, NewTxn *previous_txn, String &cause, bool &retry_query);
     bool CheckConflictTxnStore(const DumpMemIndexTxnStore &txn_store, NewTxn *previous_txn, String &cause, bool &retry_query);
     bool CheckConflictTxnStore(const DeleteTxnStore &txn_store, NewTxn *previous_txn, String &cause, bool &retry_query);
+    bool CheckConflictTxnStore(const DropTableTxnStore &txn_store, NewTxn *previous_txn, String &cause, bool &retry_query);
 
 public:
     static Status Cleanup(TxnTimeStamp ts, KVInstance *kv_instance);