Function id_array_visitor

datafusion/common/src/cse.rs:762–911

()

    #[test]
    fn id_array_visitor() -> Result<()> {
        let alias_generator = AliasGenerator::new();
        let eliminator = CSE::new(TestTreeNodeCSEController::new(
            &alias_generator,
            TestTreeNodeMask::Normal,
        ));

        let a_plus_1 = TestTreeNode::new(
            vec![
                TestTreeNode::new_leaf("a".to_string()),
                TestTreeNode::new_leaf("1".to_string()),
            ],
            "+".to_string(),
        );
        let avg_c = TestTreeNode::new(
            vec![TestTreeNode::new_leaf("c".to_string())],
            "avg".to_string(),
        );
        let sum_a_plus_1 = TestTreeNode::new(vec![a_plus_1], "sum".to_string());
        let sum_a_plus_1_minus_avg_c =
            TestTreeNode::new(vec![sum_a_plus_1, avg_c], "-".to_string());
        let root = TestTreeNode::new(
            vec![
                sum_a_plus_1_minus_avg_c,
                TestTreeNode::new_leaf("2".to_string()),
            ],
            "*".to_string(),
        );

        let [sum_a_plus_1_minus_avg_c, _] = root.children.as_slice() else {
            panic!("Cannot extract subtree references")
        };
        let [sum_a_plus_1, avg_c] = sum_a_plus_1_minus_avg_c.children.as_slice() else {
            panic!("Cannot extract subtree references")
        };
        let [a_plus_1] = sum_a_plus_1.children.as_slice() else {
            panic!("Cannot extract subtree references")
        };

        // skip aggregates
        let mut id_array = vec![];
        eliminator.node_to_id_array(&root, &mut NodeStats::new(), &mut id_array)?;

        // Collect distinct hashes and set them to 0 in `id_array`
        fn collect_hashes(
            id_array: &mut IdArray<'_, TestTreeNode<String>>,
        ) -> HashSet<u64> {
            id_array
                .iter_mut()
                .flat_map(|(_, id_option)| {
                    id_option.as_mut().map(|node_id| {
                        let hash = node_id.hash;
                        node_id.hash = 0;
                        hash
                    })
                })
                .collect::<HashSet<_>>()
        }