refactor: account for count alias in roundtrip (#1096)

Please be sure to look over the pull request guidelines here:
https://github.com/spaceandtimelabs/sxt-proof-of-sql/blob/main/CONTRIBUTING.md#submit-pr.

# Please go through the following checklist
- [ ] The PR title and commit messages adhere to guidelines here:
https://github.com/spaceandtimelabs/sxt-proof-of-sql/blob/main/CONTRIBUTING.md.
In particular `!` is used if and only if at least one breaking change
has been introduced.
- [ ] I have run the ci check script with `source
scripts/run_ci_checks.sh`.
- [ ] I have run the clean commit check script with `source
scripts/check_commits.sh`, and the commit history is certified to follow
clean commit guidelines as described here:
https://github.com/spaceandtimelabs/sxt-proof-of-sql/blob/main/COMMIT_GUIDELINES.md
- [ ] The latest changes from `main` have been incorporated to this PR
by simple rebase if possible, if not, then conflicts are resolved
appropriately.

# Rationale for this change

We need to ensure, for now, that the count alias is correct, so that
roundtrip serialization is the identity function.

# What changes are included in this PR?

Enforcing the count alias to be correct on the deserialized plan.

# Are these changes tested?
Yes.

Author: stuarttimwhite

crates/proof-of-sql/src/sql/evm_proof_plan/plans.rs

@@ -17,6 +17,7 @@ use alloc::{
     string::{String, ToString},
     vec::Vec,
 };
+use core::iter;
 use serde::{Deserialize, Serialize};
 use sqlparser::ast::Ident;
 
@@ -215,6 +216,28 @@ fn try_unwrap_output_column_names(
     Ok(output_column_names)
 }
 
+fn try_unwrap_output_column_names_with_count_alias(
+    output_column_names: Option<&IndexSet<String>>,
+    length: usize,
+    count_alias: &String,
+) -> EVMProofPlanResult<IndexSet<String>> {
+    let output_column_names = match output_column_names {
+        Some(output_column_names) => {
+            if length > output_column_names.len() {
+                return Err(EVMProofPlanError::InvalidOutputColumnName);
+            }
+            output_column_names.clone()
+        }
+        None => (0..length)
+            .map(|i| i.to_string())
+            .filter(|name| name != count_alias)
+            .take(length - 1)
+            .chain(iter::once(count_alias.clone()))
+            .collect::<IndexSet<_>>(),
+    };
+    Ok(output_column_names)
+}
+
 /// Represents a filter execution plan in EVM.
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 pub(crate) struct EVMLegacyFilterExec {
@@ -618,8 +641,11 @@ impl EVMAggregateExec {
         output_column_names: Option<&IndexSet<String>>,
     ) -> EVMProofPlanResult<AggregateExec> {
         let required_alias_count = self.group_by_exprs.len() + self.sum_expr.len() + 1;
-        let output_column_names =
-            try_unwrap_output_column_names(output_column_names, required_alias_count)?;
+        let output_column_names = try_unwrap_output_column_names_with_count_alias(
+            output_column_names,
+            required_alias_count,
+            &self.count_alias_name,
+        )?;
         let input = self
             .input_plan
             .try_into_proof_plan(table_refs, column_refs, None)?;
@@ -805,6 +831,7 @@ mod tests {
             proof_plans::{DynProofPlan, SortMergeJoinExec},
         },
     };
+    use indexmap::IndexSet;
 
     #[test]
     fn we_can_put_projection_exec_in_evm() {
@@ -2455,4 +2482,43 @@ mod tests {
             Err(EVMProofPlanError::InvalidOutputColumnName)
         ));
     }
+
+    #[test]
+    fn we_can_unwrap_correct_output_column_names_when_none() {
+        let output_column_names =
+            try_unwrap_output_column_names_with_count_alias(None, 2, &"0".to_string()).unwrap();
+        let expected_output_column_names: IndexSet<
+            String,
+            core::hash::BuildHasherDefault<ahash::AHasher>,
+        > = vec!["1".to_string(), "0".to_string()].into_iter().collect();
+        assert_eq!(output_column_names, expected_output_column_names);
+    }
+
+    #[test]
+    fn we_can_unwrap_correct_output_column_names_when_some() {
+        let expected_output_column_names: IndexSet<String, _> =
+            vec!["a".to_string(), "b".to_string()].into_iter().collect();
+        let output_column_names = try_unwrap_output_column_names_with_count_alias(
+            Some(&expected_output_column_names),
+            2,
+            &"b".to_string(),
+        )
+        .unwrap();
+
+        assert_eq!(output_column_names, expected_output_column_names);
+    }
+
+    #[test]
+    fn we_can_unwrap_err_when_mismatching_count_alias() {
+        let expected_output_column_names: IndexSet<String, _> =
+            vec!["a".to_string(), "b".to_string()].into_iter().collect();
+        let err = try_unwrap_output_column_names_with_count_alias(
+            Some(&expected_output_column_names),
+            3,
+            &"b".to_string(),
+        )
+        .unwrap_err();
+
+        assert!(matches!(err, EVMProofPlanError::InvalidOutputColumnName));
+    }
 }