Merge pull request #206 from mkhairi/fix/join-limit-and-qualified-projection

fix: join LIMIT applied before WHERE, and same-named output columns collapsing

Author: farhan-syah

nodedb/src/control/server/http/routes/query_stream.rs

@@ -122,6 +122,9 @@ pub(super) fn ndjson_body_stream(
                     return;
                 }
             };
+            // Row maps are keyed by `ShapedRows::cell_keys`, so each NDJSON
+            // line serializes as-is; two output columns sharing a name emit
+            // `{"id": …, "id_1": …}` rather than collapsing to one cell.
             let shaped = shape_decoded_rows(&value, projection.as_ref());
             for row in shaped.rows {
                 if emitted >= limit {

nodedb/src/control/server/http/routes/result_shape.rs

@@ -47,6 +47,11 @@ pub(super) fn shape_http_payload(
         database_id,
         tenant_id,
     )? {
+        // Each row map is already keyed by `ShapedRows::cell_keys`, so it
+        // serializes to JSON as-is. When two output columns share a name the
+        // later one carries a `_<n>` suffix (`SELECT w.id, b.id` →
+        // `{"id": …, "id_1": …}`) — a JSON object cannot repeat a key, and
+        // dropping the duplicate would silently lose a projected column.
         ShapeOutcome::Rows(shaped) => Ok(HttpShaped::Rows(
             shaped
                 .rows
@@ -111,6 +116,8 @@ pub(super) fn ddl_results_to_json(
                     "tag": command,
                 }));
             }
+            // Keyed by `ShapedRows::cell_keys`, same JSON contract as
+            // `shape_http_payload` above (duplicate names take a `_<n>` key).
             DdlResult::Rows(shaped) => {
                 for row in shaped.rows {
                     rows.push(serde_json::Value::Object(row));

nodedb/src/control/server/native/dispatch/conversion.rs

@@ -146,15 +146,18 @@ pub(crate) fn calvin_native_response(
 /// `json_to_value_display`; a column absent from a given row's map becomes
 /// `Value::Null`.
 pub(crate) fn to_native_columns_rows(shaped: &ShapedRows) -> (Vec<String>, Vec<Vec<Value>>) {
+    // Cells live in the row maps under per-column keys (display names may
+    // repeat across columns, e.g. `SELECT w.id, b.id`), so read through the
+    // shared accessor rather than by display name.
+    let cell_keys = shaped.cell_keys();
     let rows = shaped
         .rows
         .iter()
         .map(|row| {
-            shaped
-                .columns
+            cell_keys
                 .iter()
-                .map(|col| {
-                    row.get(col.as_str())
+                .map(|key| {
+                    row.get(key.as_str())
                         .map(json_to_value_display)
                         .unwrap_or(Value::Null)
                 })

nodedb/src/control/server/pgwire/handler/shape_encode.rs

@@ -28,22 +28,27 @@ use crate::control::server::response_shape::types::{DdlColType, ShapedRows};
 
 use super::super::ddl_encode::col_type_to_field_with_format;
 
-/// Encode one flat row object into a pgwire `DataRow`, using `columns` (in
-/// order) to look up cells in `row` and `column_types` (parallel to `columns`)
-/// to pick each cell's text rendering.
+/// Encode one flat row object into a pgwire `DataRow`, using `cell_keys` (in
+/// order) to look up cells in `row` and `column_types` (parallel to
+/// `cell_keys`) to pick each cell's text rendering.
+///
+/// `cell_keys` are the per-column unique row-map keys derived from the
+/// display names via `response_shape::project::cell_keys` — identical to the
+/// display names except where those repeat (`SELECT w.id, b.id`), in which
+/// case later duplicates carry a `_n` suffix so both cells survive the map.
 ///
 /// Missing keys and explicit JSON `null` both encode as SQL NULL. Every other
 /// cell renders per its column type via [`encode_typed_cell`]; a
 /// missing/short `column_types` entry defaults to `Text`.
 pub(in crate::control::server::pgwire) fn encode_shaped_row(
     schema: &Arc<Vec<FieldInfo>>,
-    columns: &[String],
+    cell_keys: &[String],
     column_types: &[DdlColType],
     formats: &[FieldFormat],
     row: &serde_json::Map<String, serde_json::Value>,
 ) -> PgWireResult<DataRow> {
     let mut encoder = DataRowEncoder::new(schema.clone());
-    for (idx, name) in columns.iter().enumerate() {
+    for (idx, name) in cell_keys.iter().enumerate() {
         let ct = column_types.get(idx).copied().unwrap_or(DdlColType::Text);
         let format = formats.get(idx).copied().unwrap_or(FieldFormat::Text);
         match row.get(name) {
@@ -141,6 +146,10 @@ pub(in crate::control::server::pgwire) fn shaped_query_response(
     shaped: ShapedRows,
     formats: &[FieldFormat],
 ) -> (Response, Option<String>) {
+    // Cells live in the row maps under per-column keys that differ from the
+    // display names only when two columns share a name; derived here before
+    // the struct is destructured.
+    let keys = shaped.cell_keys();
     let ShapedRows {
         columns,
         column_types,
@@ -161,7 +170,7 @@ pub(in crate::control::server::pgwire) fn shaped_query_response(
 
     let encoded_rows: Vec<PgWireResult<DataRow>> = rows
         .iter()
-        .map(|row| encode_shaped_row(&schema, &columns, &column_types, formats, row))
+        .map(|row| encode_shaped_row(&schema, &keys, &column_types, formats, row))
         .collect();
 
     let response = Response::Query(QueryResponse::new(

nodedb/src/control/server/pgwire/handler/stream_response.rs

@@ -107,6 +107,10 @@ pub(crate) fn streaming_shaped_response(
         .iter()
         .map(|c| c.display_name.clone())
         .collect();
+    // Cells live in the shaped row maps under unique per-column keys
+    // (display names may repeat across columns, e.g. `SELECT w.id, b.id`);
+    // derive the same keys the shaper used so every column reads its own cell.
+    let cell_keys = crate::control::server::response_shape::project::cell_keys(&display_columns);
     // Advertise each projected column's real catalog type so the streaming
     // path's RowDescription OIDs match the non-streaming `shaped_query_response`
     // (and the extended-query Describe path); `column_types` also drives the
@@ -160,7 +164,7 @@ pub(crate) fn streaming_shaped_response(
                     break;
                 }
                 let encoded =
-                    encode_shaped_row(&row_schema, &display_columns, &column_types, &row_formats, row)?;
+                    encode_shaped_row(&row_schema, &cell_keys, &column_types, &row_formats, row)?;
                 emitted += 1;
                 yield encoded;
             }

nodedb/src/control/server/response_shape/compose.rs

@@ -244,9 +244,15 @@ pub fn shape_decoded_rows(decoded: &JsonValue, projection: Option<&OutputSchema>
             let lookup_keys: Vec<String> = s.columns.iter().map(|c| c.lookup_key.clone()).collect();
             let display_names: Vec<String> =
                 s.columns.iter().map(|c| c.display_name.clone()).collect();
+            // Row cells are stored under per-column unique keys, not display
+            // names: duplicate display names (`SELECT w.id, b.id` → `id`,
+            // `id`) would collide in the row map and collapse both wire
+            // columns to the last value. Encoders re-derive the same keys via
+            // `cell_keys` when reading cells.
+            let keys = super::project::cell_keys(&display_names);
             let projected_rows = rows
                 .iter()
-                .map(|row| project_row(row, &lookup_keys, &display_names))
+                .map(|row| project_row(row, &lookup_keys, &display_names, &keys))
                 .collect();
             // Carry each projected column's real catalog type, aligned in
             // order with `display_names`. Only the pgwire encoder consumes
@@ -280,10 +286,15 @@ pub fn shape_decoded_rows(decoded: &JsonValue, projection: Option<&OutputSchema>
 /// Select and rename one flat row's fields per the projection lists, trying
 /// each candidate key in order: the full lookup key, then the bare
 /// (post-dot) column name, then the SELECT alias.
+///
+/// Cells are inserted under `cell_keys` (unique per column, see
+/// [`super::project::cell_keys`]) rather than the display names, which may
+/// repeat across columns and would otherwise collapse in the output map.
 fn project_row(
     row: &Map<String, JsonValue>,
     lookup_keys: &[String],
     display_names: &[String],
+    cell_keys: &[String],
 ) -> Map<String, JsonValue> {
     let mut out = Map::new();
     for (i, lookup_key) in lookup_keys.iter().enumerate() {
@@ -313,7 +324,8 @@ fn project_row(
             })
             .cloned()
             .unwrap_or(JsonValue::Null);
-        out.insert(display_name.to_string(), value);
+        let cell_key = cell_keys.get(i).map(String::as_str).unwrap_or(display_name);
+        out.insert(cell_key.to_string(), value);
     }
     out
 }
@@ -372,3 +384,44 @@ fn single_result_row(text: String) -> ShapedRows {
         notice: None,
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Two projected columns sharing the display name `id` (`SELECT w.id,
+    /// b.id`) must keep both values in the shaped row instead of collapsing
+    /// to the last table's value.
+    #[test]
+    fn project_row_keeps_both_columns_with_duplicate_display_names() {
+        let mut row = Map::new();
+        row.insert("w.id".to_string(), JsonValue::String("w1".to_string()));
+        row.insert("b.id".to_string(), JsonValue::String("b1".to_string()));
+
+        let lookup_keys = vec!["w.id".to_string(), "b.id".to_string()];
+        let display_names = vec!["id".to_string(), "id".to_string()];
+        let keys = super::super::project::cell_keys(&display_names);
+
+        let out = project_row(&row, &lookup_keys, &display_names, &keys);
+        assert_eq!(out.len(), 2, "both cells must survive the projection");
+        assert_eq!(out.get("id"), Some(&JsonValue::String("w1".to_string())));
+        assert_eq!(out.get("id_1"), Some(&JsonValue::String("b1".to_string())));
+    }
+
+    /// Duplicate-free projections still store cells under the display name
+    /// (`cell_keys` is the identity), so existing readers are unaffected.
+    #[test]
+    fn project_row_uses_display_names_when_unique() {
+        let mut row = Map::new();
+        row.insert("w.id".to_string(), JsonValue::String("w1".to_string()));
+        row.insert("b.title".to_string(), JsonValue::String("t".to_string()));
+
+        let lookup_keys = vec!["w.id".to_string(), "b.title".to_string()];
+        let display_names = vec!["id".to_string(), "title".to_string()];
+        let keys = super::super::project::cell_keys(&display_names);
+
+        let out = project_row(&row, &lookup_keys, &display_names, &keys);
+        assert_eq!(out.get("id"), Some(&JsonValue::String("w1".to_string())));
+        assert_eq!(out.get("title"), Some(&JsonValue::String("t".to_string())));
+    }
+}

nodedb/src/control/server/response_shape/project.rs

@@ -55,3 +55,69 @@ pub fn is_scan_wrapper(map: &serde_json::Map<String, serde_json::Value>) -> bool
         && matches!(map.get("id"), Some(serde_json::Value::String(_)))
         && matches!(map.get("data"), Some(serde_json::Value::Object(_)))
 }
+
+/// Unique per-column cell keys for a shaped column list.
+///
+/// SQL output column names may legally repeat — `SELECT w.id, b.id` displays
+/// both columns as `id` — but a shaped row is a JSON map and cannot hold two
+/// cells under one key: inserting the second cell would overwrite the first,
+/// making both wire columns render the last value. The shaper therefore
+/// stores the first occurrence of a name under the name itself and each later
+/// duplicate under `<name>_<n>` (n = 1, 2, …), skipping any candidate that
+/// collides with another column's name. Row writers and every protocol
+/// encoder derive keys through this one function so they always agree; for a
+/// duplicate-free column list this is the identity mapping.
+pub fn cell_keys(columns: &[String]) -> Vec<String> {
+    use std::collections::HashSet;
+
+    let names: HashSet<&str> = columns.iter().map(String::as_str).collect();
+    let mut used: HashSet<String> = HashSet::with_capacity(columns.len());
+    columns
+        .iter()
+        .map(|name| {
+            if used.insert(name.clone()) {
+                return name.clone();
+            }
+            let mut n = 1usize;
+            loop {
+                let candidate = format!("{name}_{n}");
+                if !names.contains(candidate.as_str()) && used.insert(candidate.clone()) {
+                    return candidate;
+                }
+                n += 1;
+            }
+        })
+        .collect()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::cell_keys;
+
+    fn keys(cols: &[&str]) -> Vec<String> {
+        cell_keys(&cols.iter().map(|s| s.to_string()).collect::<Vec<_>>())
+    }
+
+    /// Duplicate-free column lists map to themselves.
+    #[test]
+    fn cell_keys_is_identity_without_duplicates() {
+        assert_eq!(keys(&["id", "name", "score"]), ["id", "name", "score"]);
+        assert_eq!(keys(&[]), Vec::<String>::new());
+    }
+
+    /// Later duplicates get a `_n` suffix; the first occurrence keeps the
+    /// bare name (`SELECT w.id, b.id` → `id`, `id_1`).
+    #[test]
+    fn cell_keys_suffixes_later_duplicates() {
+        assert_eq!(keys(&["id", "id"]), ["id", "id_1"]);
+        assert_eq!(keys(&["id", "id", "id"]), ["id", "id_1", "id_2"]);
+    }
+
+    /// A suffix candidate never collides with a real column of that name:
+    /// with columns `id, id, id_1` the second `id` skips `id_1` (taken by a
+    /// real column) and becomes `id_2`.
+    #[test]
+    fn cell_keys_skips_candidates_that_shadow_real_columns() {
+        assert_eq!(keys(&["id", "id", "id_1"]), ["id", "id_2", "id_1"]);
+    }
+}

nodedb/src/control/server/response_shape/types.rs

@@ -228,6 +228,10 @@ pub struct ShapedRows {
     /// type OIDs; the native and http entrypoints ignore it. `Text` is used
     /// wherever the source type is unknown.
     pub column_types: Vec<DdlColType>,
+    /// One map per row. Cells are keyed by [`ShapedRows::cell_keys`], NOT by
+    /// `columns` directly — SQL output names may repeat (`SELECT w.id, b.id`
+    /// displays both as `id`) and a map cannot hold two cells under one key.
+    /// Read cells through `cell_keys` so each column reads its own value.
     pub rows: Vec<serde_json::Map<String, serde_json::Value>>,
     pub notice: Option<String>,
 }
@@ -238,4 +242,24 @@ impl ShapedRows {
     pub fn text_types(n: usize) -> Vec<DdlColType> {
         vec![DdlColType::Text; n]
     }
+
+    /// Per-column keys for reading cells out of [`ShapedRows::rows`], parallel
+    /// to `columns`.
+    ///
+    /// This is the single source of truth every consumer shares — the pgwire
+    /// encoders, the native converter, and the HTTP JSON serializers all key
+    /// rows through it, so the row-map layout can never drift from what a
+    /// consumer expects.
+    ///
+    /// Identical to `columns` unless two output columns share a name, in which
+    /// case later duplicates take a `_<n>` suffix (see
+    /// [`super::project::cell_keys`]). Because the HTTP transports serialize a
+    /// row map directly to JSON, that suffix is user-visible there: a
+    /// duplicate-name `SELECT w.id, b.id` emits `{"id": …, "id_1": …}`, since
+    /// a JSON object likewise cannot carry the same key twice. pgwire and
+    /// native are positional on the wire and still report both columns as
+    /// `id`, matching PostgreSQL.
+    pub fn cell_keys(&self) -> Vec<String> {
+        super::project::cell_keys(&self.columns)
+    }
 }

nodedb/src/data/executor/handlers/join/hash_handlers.rs

@@ -333,7 +333,17 @@ impl CoreLoop {
         // budget) and, if the probe fills it, surface a deterministic
         // `ResourcesExhausted` rather than dropping the excess rows. A budget
         // of 0 means "unlimited" → truly unbounded output.
-        let (probe_limit, enforce_output_budget) = if join.limit != usize::MAX {
+        //
+        // A user LIMIT may only cap the probe when there are no post-join
+        // WHERE filters: `filter_and_project` runs AFTER the probe, so
+        // truncating to `n` first would emit the first `n` ON-matched rows and
+        // then discard those failing the WHERE clause — under-filling (or
+        // emptying) the result even though later probe rows match. With
+        // post-filters present the probe runs under the budget ceiling and the
+        // user LIMIT is applied after filtering, below.
+        let has_post_filters = !join.post_filter_bytes.is_empty();
+        let (probe_limit, enforce_output_budget) = if join.limit != usize::MAX && !has_post_filters
+        {
             (join.limit, false)
         } else if budget == 0 {
             (usize::MAX, false)
@@ -374,6 +384,13 @@ impl CoreLoop {
             );
         }
 
+        // Deferred user LIMIT: when post-join WHERE filters exist the probe
+        // ran unbounded (see above) so the LIMIT must be applied here, after
+        // the filters have retained the matching rows.
+        if has_post_filters && join.limit != usize::MAX {
+            results.truncate(join.limit);
+        }
+
         let payload = super::super::super::response_codec::encode_binary_rows(&results);
         self.response_with_payload(join.task, payload)
     }