Function test_basic_split_multi_exprs

core/iox_query/src/exec/split.rs:653–721

()

    #[tokio::test]
    async fn test_basic_split_multi_exprs() {
        test_helpers::maybe_start_logging();
        let batch0 = RecordBatch::try_from_iter(vec![
            (
                "int_col",
                Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef,
            ),
            (
                "str_col",
                Arc::new(StringArray::from(vec!["one", "two", "three"])) as ArrayRef,
            ),
        ])
        .unwrap();

        let batch1 = RecordBatch::try_from_iter(vec![
            (
                "int_col",
                Arc::new(Int64Array::from(vec![4, -2])) as ArrayRef,
            ),
            (
                "str_col",
                Arc::new(StringArray::from(vec!["four", "negative 2"])) as ArrayRef,
            ),
        ])
        .unwrap();

        let input = make_input(vec![vec![batch0, batch1]]);
        // int_col < 2
        let split_expr1 = df_physical_expr(&input, col("int_col").lt(lit::<i16>(2)));
        // 2 <= int_col < 3
        let expr = col("int_col")
            .gt_eq(lit::<i16>(2))
            .and(col("int_col").lt(lit::<i16>(3)));
        let split_expr2 = df_physical_expr(&input, expr);
        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 | str_col    |",
            "+---------+------------+",
            "| -2      | negative 2 |",
            "| 1       | one        |",
            "+---------+------------+",
        ];
        assert_batches_sorted_eq!(&expected, &output0);

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