Function test_first_group_acc

datafusion/functions-aggregate/src/first_last.rs:1531–1622

()

    #[test]
    fn test_first_group_acc() -> Result<()> {
        let schema = Arc::new(Schema::new(vec![
            Field::new("a", DataType::Int64, true),
            Field::new("b", DataType::Int64, true),
            Field::new("c", DataType::Int64, true),
            Field::new("d", DataType::Int32, true),
            Field::new("e", DataType::Boolean, true),
        ]));

        let sort_keys = [PhysicalSortExpr {
            expr: col("c", &schema).unwrap(),
            options: SortOptions::default(),
        }];

        let mut group_acc = FirstLastGroupsAccumulator::try_new(
            PrimitiveValueState::<Int64Type>::new(DataType::Int64),
            sort_keys.into(),
            true,
            &[DataType::Int64],
            true,
        )?;

        let mut val_with_orderings = {
            let mut val_with_orderings = Vec::<ArrayRef>::new();

            let vals = Arc::new(Int64Array::from(vec![Some(1), None, Some(3), Some(-6)]));
            let orderings = Arc::new(Int64Array::from(vec![1, -9, 3, -6]));

            val_with_orderings.push(vals);
            val_with_orderings.push(orderings);

            val_with_orderings
        };

        group_acc.update_batch(
            &val_with_orderings,
            &[0, 1, 2, 1],
            Some(&BooleanArray::from(vec![true, true, false, true])),
            3,
        )?;
        assert_eq!(
            group_acc.size_of_orderings,
            group_acc.compute_size_of_orderings()
        );

        let state = group_acc.state(EmitTo::All)?;

        let expected_state: Vec<Arc<dyn Array>> = vec![
            Arc::new(Int64Array::from(vec![Some(1), Some(-6), None])),
            Arc::new(Int64Array::from(vec![Some(1), Some(-6), None])),
            Arc::new(BooleanArray::from(vec![true, true, false])),
        ];
        assert_eq!(state, expected_state);

        assert_eq!(
            group_acc.size_of_orderings,
            group_acc.compute_size_of_orderings()
        );