fix: update pk_cache in compat reader (#3576)

* fix: update pk_cache in compat reader Signed-off-by: Ruihang Xia <[email protected]> * add sqlness case Signed-off-by: Ruihang Xia <[email protected]> * update document Signed-off-by: Ruihang Xia <[email protected]> * add more sqlness case Signed-off-by: Ruihang Xia <[email protected]> * avoid mysterious bug Signed-off-by: Ruihang Xia <[email protected]> --------- Signed-off-by: Ruihang Xia <[email protected]>
GreptimeTeam · Mar 26, 2024 · 7456515 · 7456515
1 parent 83c1b48
commit 7456515
Show file tree

Hide file tree

Showing 4 changed files with 54 additions and 9 deletions.
diff --git a/src/mito2/src/read.rs b/src/mito2/src/read.rs
@@ -191,6 +191,9 @@ impl Batch {
     }
 
     /// Replaces the primary key of the batch.
+    ///
+    /// Notice that this [Batch] also contains a maybe-exist `pk_values`.
+    /// Be sure to update that field as well.
     pub fn set_primary_key(&mut self, primary_key: Vec<u8>) {
         self.primary_key = primary_key;
     }

diff --git a/src/mito2/src/read/compat.rs b/src/mito2/src/read/compat.rs
@@ -119,6 +119,14 @@ impl CompatPrimaryKey {
         )?;
 
         batch.set_primary_key(buffer);
+
+        // update cache
+        if let Some(pk_values) = &mut batch.pk_values {
+            for value in &self.values {
+                pk_values.push(value.clone());
+            }
+        }
+
         Ok(batch)
     }
 }

diff --git a/tests/cases/standalone/common/alter/alter_table.result b/tests/cases/standalone/common/alter/alter_table.result
@@ -1,4 +1,4 @@
-CREATE TABLE test_alt_table(i INTEGER, j TIMESTAMP TIME INDEX);
+CREATE TABLE test_alt_table(h INTEGER, i INTEGER, j TIMESTAMP TIME INDEX, PRIMARY KEY (h, i));
 
 Affected Rows: 0
 
@@ -7,11 +7,18 @@ DESC TABLE test_alt_table;
 +--------+----------------------+-----+------+---------+---------------+
 | Column | Type                 | Key | Null | Default | Semantic Type |
 +--------+----------------------+-----+------+---------+---------------+
-| i      | Int32                |     | YES  |         | FIELD         |
+| h      | Int32                | PRI | YES  |         | TAG           |
+| i      | Int32                | PRI | YES  |         | TAG           |
 | j      | TimestampMillisecond | PRI | NO   |         | TIMESTAMP     |
 +--------+----------------------+-----+------+---------+---------------+
 
-ALTER TABLE test_alt_table ADD COLUMN k INTEGER;
+INSERT INTO test_alt_table VALUES (1, 1, 0), (2, 2, 1);
+
+Affected Rows: 2
+
+-- TODO: It may result in an error if `k` is with type INTEGER.
+-- Error: 3001(EngineExecuteQuery), Invalid argument error: column types must match schema types, expected Int32 but found Utf8 at column index 3
+ALTER TABLE test_alt_table ADD COLUMN k STRING PRIMARY KEY;
 
 Affected Rows: 0
 
@@ -20,11 +27,29 @@ DESC TABLE test_alt_table;
 +--------+----------------------+-----+------+---------+---------------+
 | Column | Type                 | Key | Null | Default | Semantic Type |
 +--------+----------------------+-----+------+---------+---------------+
-| i      | Int32                |     | YES  |         | FIELD         |
+| h      | Int32                | PRI | YES  |         | TAG           |
+| i      | Int32                | PRI | YES  |         | TAG           |
 | j      | TimestampMillisecond | PRI | NO   |         | TIMESTAMP     |
-| k      | Int32                |     | YES  |         | FIELD         |
+| k      | String               | PRI | YES  |         | TAG           |
 +--------+----------------------+-----+------+---------+---------------+
 
+SELECT * FROM test_alt_table;
+
++---+---+-------------------------+---+
+| h | i | j                       | k |
++---+---+-------------------------+---+
+| 1 | 1 | 1970-01-01T00:00:00     |   |
+| 2 | 2 | 1970-01-01T00:00:00.001 |   |
++---+---+-------------------------+---+
+
+SELECT * FROM test_alt_table WHERE i = 1;
+
++---+---+---------------------+---+
+| h | i | j                   | k |
++---+---+---------------------+---+
+| 1 | 1 | 1970-01-01T00:00:00 |   |
++---+---+---------------------+---+
+
 -- SQLNESS ARG restart=true
 ALTER TABLE test_alt_table ADD COLUMN m INTEGER;
 
@@ -35,9 +60,10 @@ DESC TABLE test_alt_table;
 +--------+----------------------+-----+------+---------+---------------+
 | Column | Type                 | Key | Null | Default | Semantic Type |
 +--------+----------------------+-----+------+---------+---------------+
-| i      | Int32                |     | YES  |         | FIELD         |
+| h      | Int32                | PRI | YES  |         | TAG           |
+| i      | Int32                | PRI | YES  |         | TAG           |
 | j      | TimestampMillisecond | PRI | NO   |         | TIMESTAMP     |
-| k      | Int32                |     | YES  |         | FIELD         |
+| k      | String               | PRI | YES  |         | TAG           |
 | m      | Int32                |     | YES  |         | FIELD         |
 +--------+----------------------+-----+------+---------+---------------+
 

diff --git a/tests/cases/standalone/common/alter/alter_table.sql b/tests/cases/standalone/common/alter/alter_table.sql
@@ -1,11 +1,19 @@
-CREATE TABLE test_alt_table(i INTEGER, j TIMESTAMP TIME INDEX);
+CREATE TABLE test_alt_table(h INTEGER, i INTEGER, j TIMESTAMP TIME INDEX, PRIMARY KEY (h, i));
 
 DESC TABLE test_alt_table;
 
-ALTER TABLE test_alt_table ADD COLUMN k INTEGER;
+INSERT INTO test_alt_table VALUES (1, 1, 0), (2, 2, 1);
+
+-- TODO: It may result in an error if `k` is with type INTEGER.
+-- Error: 3001(EngineExecuteQuery), Invalid argument error: column types must match schema types, expected Int32 but found Utf8 at column index 3
+ALTER TABLE test_alt_table ADD COLUMN k STRING PRIMARY KEY;
 
 DESC TABLE test_alt_table;
 
+SELECT * FROM test_alt_table;
+
+SELECT * FROM test_alt_table WHERE i = 1;
+
 -- SQLNESS ARG restart=true
 ALTER TABLE test_alt_table ADD COLUMN m INTEGER;