Repository: tikv/tikv
File: components/tidb_query_executors/src/index_scan_executor.rs (Rust)

Code before:
        // 1. Partition ID in key doesn't break handle decoding
        // 2. Partition ID from value is still extracted correctly
        // 3. All columns are output correctly

        const TABLE_ID: i64 = 100;
        const INDEX_ID: i64 = 5;
        const PARTITION_ID: i64 = 42;

        // Column definitions: indexed column, handle, partition ID
        let columns_info = vec![
            {
                let mut ci = ColumnInfo::default();
                ci.set_column_id(1); // Indexed column
                ci.as_mut_accessor().set_tp(FieldTypeTp::LongLong);
                ci
            },
            {
                let mut ci = ColumnInfo::default();
                ci.set_column_id(-1); // Handle column
                ci.set_pk_handle(true);
                ci.as_mut_accessor().set_tp(FieldTypeTp::LongLong);
                ci
            },
            {
                let mut ci = ColumnInfo::default();
                ci.set_column_id(table::EXTRA_PARTITION_ID_COL_ID); // -2
                ci.as_mut_accessor().set_tp(FieldTypeTp::LongLong);
                ci
            },
        ];

        let schema: Vec<FieldType> = vec![
            FieldTypeTp::LongLong.into(), // indexed column
            FieldTypeTp::LongLong.into(), // handle
            FieldTypeTp::LongLong.into(), // partition id
        ];

        // Build index key with V1 format:
        // [table_prefix][index_id][indexed_col][PARTITION_ID_FLAG][partition_id][handle]
        let indexed_value = 999i64;
        let handle_value = 123i64;

        // Encode indexed column
        let mut index_key_data = datum::encode_key(&mut EvalContext::default(), &[Datum::I64(indexed_value)]).unwrap();
        // Add partition ID prefix (V1 format - new location)
        index_key_data.push(table::INDEX_VALUE_PARTITION_ID_FLAG);
        index_key_data.write_i64(PARTITION_ID).unwrap();
        // Add handle
        let handle_data = datum::encode_key(&mut EvalContext::default(), &[Datum::I64(handle_value)]).unwrap();
        index_key_data.extend(handle_data);


Review comment (suggestion):
Consider adding a test case for a key with a valid handle flag (INT_FLAG or UINT_FLAG) but insufficient data for the handle value itself. For example, test a key like `[INT_FLAG, 0x01]` which has only 1 byte of handle data instead of the required 8 bytes. This would verify that the error handling in the `read_i64()`/`read_u64()` calls works correctly.

Code after:
        // Test error: empty key
        let key = vec![];
        idx_exe.decode_int_handle_from_key(&key).unwrap_err();

        // Test error: INT_FLAG with insufficient data (only 1 byte instead of 8)
        let key = vec![datum::INT_FLAG, 0x01];
        idx_exe.decode_int_handle_from_key(&key).unwrap_err();

        // Test error: UINT_FLAG with insufficient data (only 1 byte instead of 8)
        let key = vec![datum::UINT_FLAG, 0x01];
        idx_exe.decode_int_handle_from_key(&key).unwrap_err();

        // Test error: partition flag with INT_FLAG but insufficient handle data
        let mut key = vec![table::INDEX_VALUE_PARTITION_ID_FLAG];
        key.write_i64(500).unwrap(); // partition id
        key.push(datum::INT_FLAG);
        key.push(0x01); // only 1 byte instead of 8
        idx_exe.decode_int_handle_from_key(&key).unwrap_err();

        // Test error: partition flag with UINT_FLAG but insufficient handle data
        let mut key = vec![table::INDEX_VALUE_PARTITION_ID_FLAG];
        key.write_i64(600).unwrap(); // partition id
        key.push(datum::UINT_FLAG);
        key.push(0x01); // only 1 byte instead of 8
        idx_exe.decode_int_handle_from_key(&key).unwrap_err();
    }

    #[test]
    fn test_non_unique_global_index_v1_partition_in_key_and_value() {
        // V1 format: partition ID in both KEY and VALUE
        // This test validates that:
        // 1. Partition ID in key doesn't break handle decoding
        // 2. Partition ID from value is still extracted correctly
        // 3. All columns are output correctly
        // 4. Edge cases: zero and negative partition IDs work correctly

        const TABLE_ID: i64 = 100;
        const INDEX_ID: i64 = 5;

        // Column definitions: indexed column, handle, partition ID
        let columns_info = vec![
            {
                let mut ci = ColumnInfo::default();
                ci.set_column_id(1); // Indexed column
                ci.as_mut_accessor().set_tp(FieldTypeTp::LongLong);
                ci
            },
            {
                let mut ci = ColumnInfo::default();
                ci.set_column_id(-1); // Handle column
                ci.set_pk_handle(true);