Related to #35941 Previous PR: #36034 This patch makes the switch branching logic correct and make the unit test work for cases which does not select the whole dataset. Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>