Fix PARTITION BY bug for spatial projections

src/execute/layer.rs

@@ -530,19 +530,7 @@ where
     // Note: Facet aesthetics are already in partition_by via add_discrete_columns_to_partition_by,
     // so we don't add facet.get_variables() here (which would add original column names
     // instead of aesthetic column names, breaking pre-stat transforms like domain censoring).
-    let mut group_by: Vec<String> = Vec::new();
-    for col in &layer.partition_by {
-        group_by.push(col.clone());
-    }
-
-    // Add literal aesthetic columns to group_by so they survive stat transforms.
-    // Since literal columns contain constant values (same for every row), adding them
-    // to GROUP BY doesn't affect aggregation results - they're simply preserved.
-    for col in &literal_columns {
-        if !group_by.contains(col) {
-            group_by.push(col.clone());
-        }
-    }
+    let group_by = crate::util::set_union(layer.partition_by.clone(), &literal_columns);
 
     // Apply statistical transformation (uses aesthetic names)
     let stat_result = layer.geom.apply_stat_transform(

src/execute/mod.rs

@@ -3312,6 +3312,56 @@ mod tests {
         );
     }
 
+    #[cfg(all(feature = "duckdb", feature = "spatial"))]
+    #[test]
+    fn test_partition_by_preserved_through_map_projection() {
+        let reader = DuckDBReader::from_connection_string("duckdb://memory").unwrap();
+        reader
+            .connection()
+            .execute_batch("INSTALL spatial; LOAD spatial")
+            .unwrap();
+        reader
+            .connection()
+            .execute(
+                "CREATE TABLE routes AS \
+                 SELECT 10.0 AS lon, 50.0 AS lat, 'A' AS direction, 1 AS grp \
+                 UNION ALL SELECT 20.0, 55.0, 'A', 1 \
+                 UNION ALL SELECT 30.0, 52.0, 'R' , 2",
+                duckdb::params![],
+            )
+            .unwrap();
+
+        let query = r#"
+            SELECT lon, lat, direction, grp FROM routes
+            VISUALISE lon AS x, lat AS y
+            DRAW path PARTITION BY direction, grp
+            PROJECT x, y TO orthographic
+        "#;
+
+        let result = prepare_data_with_reader(query, &reader);
+        assert!(
+            result.is_ok(),
+            "PARTITION BY columns should survive map projection: {:?}",
+            result.err()
+        );
+
+        let prepared = result.unwrap();
+        let layer = &prepared.specs[0].layers[0];
+        let data_key = layer.data_key.as_ref().unwrap();
+        let df = prepared.data.get(data_key).unwrap();
+        let col_names = df.get_column_names();
+        assert!(
+            col_names.iter().any(|c| c == "direction"),
+            "partition column 'direction' missing from output data; columns: {:?}",
+            col_names
+        );
+        assert!(
+            col_names.iter().any(|c| c == "grp"),
+            "partition column 'grp' missing from output data; columns: {:?}",
+            col_names
+        );
+    }
+
     #[cfg(feature = "duckdb")]
     #[test]
     fn test_case_insensitive_facet() {

src/plot/layer/geom/line.rs

@@ -93,7 +93,7 @@ impl GeomTrait for Line {
         if !needs_projection(projection) {
             return Ok(query.to_string());
         }
-        let columns = mappings.column_names();
+        let columns = crate::util::set_union(mappings.column_names(), partition_by);
         let pos1_col = naming::aesthetic_column("pos1");
         let densified = densify_edges(
             query,

src/plot/layer/geom/path.rs

@@ -53,7 +53,7 @@ impl GeomTrait for Path {
         if !needs_projection(projection) {
             return Ok(query.to_string());
         }
-        let columns = mappings.column_names();
+        let columns = crate::util::set_union(mappings.column_names(), partition_by);
         let densified = densify_edges(
             query,
             dialect,

src/plot/layer/geom/point.rs

@@ -61,10 +61,10 @@ impl GeomTrait for Point {
         projection: &Projection,
         dialect: &dyn SqlDialect,
         mappings: &mut Mappings,
-        _partition_by: &mut Vec<String>,
+        partition_by: &mut Vec<String>,
         _parameters: &mut std::collections::HashMap<String, crate::plot::types::ParameterValue>,
     ) -> Result<String> {
-        let columns = mappings.column_names();
+        let columns = crate::util::set_union(mappings.column_names(), partition_by);
         project_position_columns(query, projection, dialect, &columns)
     }
 }

src/plot/layer/geom/polygon.rs

@@ -54,7 +54,7 @@ impl GeomTrait for Polygon {
         if !needs_projection(projection) {
             return Ok(query.to_string());
         }
-        let columns = mappings.column_names();
+        let columns = crate::util::set_union(mappings.column_names(), partition_by);
         let densified = densify_edges(query, dialect, &columns, partition_by, None, true, 1.0, 360);
         project_position_columns(&densified, projection, dialect, &columns)
     }

src/plot/layer/geom/ribbon.rs

@@ -71,6 +71,8 @@ impl GeomTrait for Ribbon {
         partition_by.push(naming::DENSIFY_ID_COLUMN.to_string());
         parameters.insert("densified".to_string(), ParameterValue::Boolean(true));
 
+        let expanded_columns = crate::util::set_union(expanded_columns, partition_by);
+
         let densified = densify_edges(
             &expanded,
             dialect,

src/plot/layer/geom/rule.rs

@@ -58,7 +58,7 @@ impl GeomTrait for Rule {
         // The rule input always has one position column named __ggsql_aes_pos1__
         // regardless of whether the aesthetic key is "pos1" or "pos2" — the executor
         // normalizes it to the pos1 slot.
-        let columns = vec![naming::aesthetic_column("pos1")];
+        let columns = crate::util::set_union(vec![naming::aesthetic_column("pos1")], partition_by);
 
         let has_pos1 = mappings.contains_key("pos1");
         let bbox_expr = match projection.coord.coord_kind() {
@@ -80,6 +80,8 @@ impl GeomTrait for Rule {
         partition_by.push(naming::DENSIFY_ID_COLUMN.to_string());
         parameters.insert("densified".to_string(), ParameterValue::Boolean(true));
 
+        let expanded_columns = crate::util::set_union(expanded_columns, partition_by);
+
         let densified = densify_edges(
             &expanded,
             dialect,

src/plot/layer/geom/segment.rs

@@ -69,6 +69,8 @@ impl GeomTrait for Segment {
         partition_by.push(naming::DENSIFY_ID_COLUMN.to_string());
         parameters.insert("densified".to_string(), ParameterValue::Boolean(true));
 
+        let expanded_columns = crate::util::set_union(expanded_columns, partition_by);
+
         let densified = densify_edges(
             &expanded,
             dialect,

src/plot/layer/geom/text.rs

@@ -76,10 +76,10 @@ impl GeomTrait for Text {
         projection: &Projection,
         dialect: &dyn SqlDialect,
         mappings: &mut Mappings,
-        _partition_by: &mut Vec<String>,
+        partition_by: &mut Vec<String>,
         _parameters: &mut std::collections::HashMap<String, crate::plot::types::ParameterValue>,
     ) -> Result<String> {
-        let columns = mappings.column_names();
+        let columns = crate::util::set_union(mappings.column_names(), partition_by);
         project_position_columns(query, projection, dialect, &columns)
     }
 

src/plot/layer/geom/tile.rs

@@ -222,7 +222,7 @@ impl GeomTrait for Tile {
             return Ok(query.to_string());
         }
 
-        let columns = mappings.column_names();
+        let columns = crate::util::set_union(mappings.column_names(), partition_by);
 
         // Only densify continuous tiles (those parameterized by pos1min/pos1max/pos2min/pos2max).
         // Discrete tiles use categorical positions and don't appear on maps.
@@ -240,6 +240,8 @@ impl GeomTrait for Tile {
         partition_by.push(naming::DENSIFY_ID_COLUMN.to_string());
         parameters.insert("densified".to_string(), ParameterValue::Boolean(true));
 
+        let expanded_columns = crate::util::set_union(expanded_columns, partition_by);
+
         let densified = densify_edges(
             &expanded,
             dialect,

src/plot/projection/coord/map.rs

@@ -98,7 +98,8 @@ pub(crate) fn apply_map_transforms(
         }
 
         layer_queries[idx] = if is_spatial {
-            let columns = layer.mappings.column_names();
+            let columns =
+                crate::util::set_union(layer.mappings.column_names(), &layer.partition_by);
             let geom_col_quoted = naming::quote_ident(&naming::aesthetic_column("geometry"));
             let wkb_expr = dialect.sql_geometry_to_wkb(&geom_col_quoted);
             dialect.sql_select_replace(

src/util.rs

@@ -77,6 +77,16 @@ fn format_quoted_list<T: Display>(items: &[T], quote: char, conjunction: &str) -
     format_list(&quoted, conjunction)
 }
 
+/// Return the set union of `old` and `new`, preserving order and deduplicating.
+pub fn set_union(mut old: Vec<String>, new: &[String]) -> Vec<String> {
+    for item in new {
+        if !old.contains(item) {
+            old.push(item.clone());
+        }
+    }
+    old
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;