Function test_constant_split_multi_exprs

core/iox_query/src/exec/split.rs:757–853

()

    #[tokio::test]
    async fn test_constant_split_multi_exprs() {
        // test that it works with a constant expression
        test_helpers::maybe_start_logging();
        let batch0 = RecordBatch::try_from_iter(vec![(
            "int_col",
            Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef,
        )])
        .unwrap();

        // Test 1: 3 streams but all data is sent to the second one
        let input = make_input(vec![vec![batch0.clone()]]);
        // use `false` & `true` to send all outputs to second stream
        let split_expr1 = df_physical_expr(&input, lit(false));
        let split_expr2 = df_physical_expr(&input, lit(true));
        let split_exec: Arc<dyn ExecutionPlan> =
            Arc::new(StreamSplitExec::new(input, vec![split_expr1, split_expr2]));

        let output0 = test_collect_partition(Arc::clone(&split_exec), 0).await;
        let expected = vec!["+---------+", "| int_col |", "+---------+", "+---------+"];
        assert_batches_sorted_eq!(&expected, &output0);

        let output1 = test_collect_partition(Arc::clone(&split_exec), 1).await;
        let expected = vec![
            "+---------+",
            "| int_col |",
            "+---------+",
            "| 1       |",
            "| 2       |",
            "| 3       |",
            "+---------+",
        ];
        assert_batches_sorted_eq!(&expected, &output1);

        let output2 = test_collect_partition(split_exec, 2).await;
        let expected = vec!["+---------+", "| int_col |", "+---------+", "+---------+"];
        assert_batches_sorted_eq!(&expected, &output2);

        // -----------------------
        // Test 2: 3 streams but all data is sent to the last one
        let input = make_input(vec![vec![batch0.clone()]]);

        // use `false` & `false` to send all outputs to third stream
        let split_expr1 = df_physical_expr(&input, lit(false));
        let split_expr2 = df_physical_expr(&input, lit(false));
        let split_exec: Arc<dyn ExecutionPlan> =
            Arc::new(StreamSplitExec::new(input, vec![split_expr1, split_expr2]));

        let output0 = test_collect_partition(Arc::clone(&split_exec), 0).await;
        let expected = vec!["+---------+", "| int_col |", "+---------+", "+---------+"];
        assert_batches_sorted_eq!(&expected, &output0);

        let output1 = test_collect_partition(Arc::clone(&split_exec), 1).await;
        let expected = vec!["+---------+", "| int_col |", "+---------+", "+---------+"];
        assert_batches_sorted_eq!(&expected, &output1);

        let output2 = test_collect_partition(Arc::clone(&split_exec), 2).await;
        let expected = vec![
            "+---------+",