Repository: tikv/tikv
File: tests/integrations/coprocessor/test_select.rs (Rust)

Code before:
            vec![
                Datum::Bytes(b"name:0".to_vec()),
                Datum::I64(2),
                Datum::I64(1)
            ],
            vec![
                Datum::Bytes(b"name:1".to_vec()),
                Datum::I64(4),
                Datum::I64(15)
            ],
        ],
        intermediate_rows
    );
}

// fix issue https://github.com/tikv/tikv/issues/19299 to test the case
// when _tidb_commit_ts column is placed before PK handle
#[test]
fn test_select_with_commit_ts_column_before_pk_handle() {
    use tidb_query_datatype::{FieldTypeAccessor, FieldTypeTp};

    let data = vec![
        (1, Some("name:0"), 2),
        (2, Some("name:4"), 3),
        (4, Some("name:3"), 1),
        (5, Some("name:1"), 4),
    ];

    let product = ProductTable::new();
    let (_, endpoint, _) = init_with_data_ext(&product, &data);

    let mut commit_ts_col = tipb::ColumnInfo::default();
    commit_ts_col.set_column_id(tidb_query_datatype::codec::table::EXTRA_COMMIT_TS_COL_ID);
    commit_ts_col
        .as_mut_accessor()
        .set_tp(FieldTypeTp::LongLong);

    let mut select = DagSelect::from(&product);
    select
        .execs
        .get_mut(0)
        .unwrap()
        .mut_tbl_scan()
        .mut_columns()
        .insert(0, commit_ts_col);
    select.output_offsets = Some(vec![1, 2, 3]);

    let req = select.build();
    let mut resp = handle_select(&endpoint, req);
    let spliter = DagChunkSpliter::new(resp.take_chunks().into(), 3);
    for (row, (id, name, cnt)) in spliter.zip(data) {

Review comment (performance):
The comment above this test is misleading. This test does not test the case when _tidb_commit_ts column is placed before PK handle - it only tests batch row limit functionality. The test that actually tests the commit_ts column scenario is test_select_with_commit_ts_column_before_pk_handle above. Consider either removing this comment or clarifying that this is an existing test that was moved to the end of the file (with a reference to why it was moved if relevant to the issue).

Code after:
            vec![
                Datum::Bytes(b"name:0".to_vec()),
                Datum::I64(2),
                Datum::I64(1)
            ],
            vec![
                Datum::Bytes(b"name:1".to_vec()),
                Datum::I64(4),
                Datum::I64(15)
            ],
        ],
        intermediate_rows
    );
}

#[test]
fn test_select_with_commit_ts() {
    use futures::executor::block_on;
    use tidb_query_datatype::{FieldTypeAccessor, FieldTypeTp};

    let data = vec![
        (1, Some("name:0"), 2),
        (2, Some("name:4"), 3),
        (4, Some("name:3"), 1),
        (5, Some("name:1"), 4),
    ];

    let product = ProductTable::new();
    let (store, endpoint, _) = init_with_data_ext(&product, &data);
    let storage = store.get_storage().get_storage();

    for commit_ts_at_first in [true, false] {
        let mut commit_ts_col = tipb::ColumnInfo::default();
        commit_ts_col.set_column_id(tidb_query_datatype::codec::table::EXTRA_COMMIT_TS_COL_ID);
        commit_ts_col
            .as_mut_accessor()
            .set_tp(FieldTypeTp::LongLong);

        let mut select = DagSelect::from(&product);
        let columns = select
            .execs
            .get_mut(0)
            .unwrap()
            .mut_tbl_scan()
            .mut_columns();
        if commit_ts_at_first {
            // To test the issue https://github.com/tikv/tikv/issues/19299
            // when _tidb_commit_ts column is placed before PK handle should not panic.
            columns.insert(0, commit_ts_col);
        } else {
            columns.push(commit_ts_col);