@Test
  public void topKOperatorSortTest() throws Exception {
    AtomicInteger count1 = new AtomicInteger(0);
    AtomicInteger count2 = new AtomicInteger(0);
    long[] time1 = new long[] {1, 2, 3, 4, 5, 6};
    long[] time2 = new long[] {3, 3, 6, 7, 8, 9};
    long[] ans = new long[] {1, 2, 3, 3, 3, 4, 5, 6, 6, 7, 8, 9};

    TsBlock[] tsBlocks1 = constructOutput(time1);
    TsBlock[] tsBlocks2 = constructOutput(time2);

    ListenableFuture<Boolean> mockFuture = Mockito.mock(ListenableFuture.class);
    doReturn(mockFuture).when(childOperator1).isBlocked();
    doReturn(mockFuture).when(childOperator2).isBlocked();
    when(mockFuture.isDone()).thenReturn(true);
    when(childOperator1.nextWithTimer())
        .thenAnswer(
            (Answer<TsBlock>)
                invocationOnMock -> {
                  int count = count1.getAndIncrement();
                  if (count == 0) {
                    return tsBlocks1[0];
                  } else if (count == 1) {
                    return tsBlocks1[1];
                  } else {
                    return null;
                  }
                });
    when(childOperator2.nextWithTimer())
        .thenAnswer(
            (Answer<TsBlock>)
                invocationOnMock -> {
                  int count = count2.getAndIncrement();
                  if (count == 0) {
                    return tsBlocks2[0];
                  } else if (count == 1) {
                    return tsBlocks2[1];
                  } else {
                    return null;
                  }
                });

    when(childOperator1.hasNextWithTimer())
        .thenAnswer(
            (Answer<Boolean>)
                invocationOnMock -> {
                  int count = count1.get();
                  return count < 2;
                });

    when(childOperator2.hasNextWithTimer())
        .thenAnswer(
            (Answer<Boolean>)
                invocationOnMock -> {
                  int count = count2.get();
                  return count < 2;
                });

    QueryId queryId = new QueryId("stub_query");
    ExecutorService instanceNotificationExecutor =
        IoTDBThreadPoolFactory.newFixedThreadPool(1, "test-instance-notification");
    FragmentInstanceId instanceId =
        new FragmentInstanceId(new PlanFragmentId(queryId, 0), "stub-instance");
    FragmentInstanceStateMachine stateMachine =
        new FragmentInstanceStateMachine(instanceId, instanceNotificationExecutor);
    FragmentInstanceContext fragmentInstanceContext =
        createFragmentInstanceContext(instanceId, stateMachine);
    DriverContext driverContext = new DriverContext(fragmentInstanceContext, 0);
    PlanNodeId planNodeId1 = new PlanNodeId("1");
    driverContext.addOperatorContext(1, planNodeId1, TreeMergeSortOperator.class.getSimpleName());

    TreeTopKOperator topKOperator =
        new TreeTopKOperator(
            driverContext.getOperatorContexts().get(0),
            Arrays.asList(childOperator1, childOperator2),
            Collections.singletonList(TSDataType.INT64),
            MergeSortComparator.getComparator(
                Collections.singletonList(new SortItem(OrderByKey.TIME, Ordering.ASC)),
                Collections.singletonList(-1),
                Collections.singletonList(TSDataType.INT64)),
            20,
            true);
    OperatorContext.setMaxRunTime(new Duration(500, TimeUnit.MILLISECONDS));

    int index = 0;
    while (topKOperator.isBlocked().isDone() && topKOperator.hasNext()) {
      TsBlock result = topKOperator.next();
      if (result == null) {
        continue;
      }
      for (int i = 0; i < result.getPositionCount(); i++) {
        long time = result.getTimeByIndex(i);
        assertEquals(time, ans[index++]);
      }
    }
    assertEquals(index, ans.length);
  }
The original data and structure of this test is same as {@link MergeTreeSortOperatorTest}.
iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/TreeTopKOperatorTest.java