From 3760d44c13f69bf3a3114ffc41abc08ed737bec8 Mon Sep 17 00:00:00 2001
From: Jussi Saurio <jussi.saurio@gmail.com>
Date: Fri, 15 Aug 2025 15:45:33 +0300
Subject: [PATCH] sim: force compound selects to have the same number of result
 columns

---
 simulator/generation/query.rs | 46 +++++++++++++++++++++++++++++++++--
 1 file changed, 44 insertions(+), 2 deletions(-)

diff --git a/simulator/generation/query.rs b/simulator/generation/query.rs
index c0b3ce7cc..7e07c529c 100644
--- a/simulator/generation/query.rs
+++ b/simulator/generation/query.rs
@@ -144,6 +144,43 @@ impl ArbitraryFrom<&SimulatorEnv> for SelectInner {
     }
 }
 
+impl ArbitrarySizedFrom<&SimulatorEnv> for SelectInner {
+    fn arbitrary_sized_from<R: Rng>(
+        rng: &mut R,
+        env: &SimulatorEnv,
+        num_result_columns: usize,
+    ) -> Self {
+        let mut select_inner = SelectInner::arbitrary_from(rng, env);
+        let select_from = &select_inner.from.as_ref().unwrap();
+        let table_names = select_from
+            .joins
+            .iter()
+            .map(|j| j.table.clone())
+            .chain(std::iter::once(select_from.table.clone()))
+            .collect::<Vec<_>>();
+
+        let flat_columns_names = table_names
+            .iter()
+            .flat_map(|t| {
+                env.tables
+                    .iter()
+                    .find(|table| table.name == *t)
+                    .unwrap()
+                    .columns
+                    .iter()
+                    .map(|c| format!("{}.{}", t.clone(), c.name))
+            })
+            .collect::<Vec<_>>();
+        let selected_columns = pick_unique(&flat_columns_names, num_result_columns, rng);
+        let mut columns = Vec::new();
+        for column_name in selected_columns {
+            columns.push(ResultColumn::Column(column_name.clone()));
+        }
+        select_inner.columns = columns;
+        select_inner
+    }
+}
+
 impl Arbitrary for Distinctness {
     fn arbitrary<R: Rng>(rng: &mut R) -> Self {
         match rng.gen_range(0..=5) {
@@ -191,10 +228,15 @@ impl ArbitraryFrom<&SimulatorEnv> for Select {
             0
         };
 
-        let first = SelectInner::arbitrary_from(rng, env);
+        let min_column_count_across_tables =
+            env.tables.iter().map(|t| t.columns.len()).min().unwrap();
+
+        let num_result_columns = rng.gen_range(1..=min_column_count_across_tables);
+
+        let first = SelectInner::arbitrary_sized_from(rng, env, num_result_columns);
 
         let rest: Vec<SelectInner> = (0..num_compound_selects)
-            .map(|_| SelectInner::arbitrary_from(rng, env))
+            .map(|_| SelectInner::arbitrary_sized_from(rng, env, num_result_columns))
             .collect();
 
         Self {