fix(clickhouse): implement re_scan for nested loop joins (supabase#546)

* fix(clickhouse): implement re_scan for nested loop joins

Fixes supabase#532

PostgreSQL uses Nested Loop joins when joining foreign tables.
During a Nested Loop join, the inner (right) table is scanned
multiple times - once for each row of the outer table. The FDW
framework calls re_scan() to restart the inner scan for each
outer row.

The ClickHouse FDW was using the default no-op re_scan()
implementation, which meant that after the first scan completed,
subsequent rescans returned no data because:
- is_scan_complete was still true
- row_receiver channel was exhausted
- No new streaming task was spawned

This fix implements re_scan() to restart the async streaming:
- Abort existing streaming task if running
- Reset scan state flags
- Reinitialize the bounded channel
- Spawn new streaming task with the same SQL query
- Fetch the first row to initialize the rescan

Also adds integration tests for inner, left, and right joins.

* code refactor and upgrade version

---------

Co-authored-by: Bo Lu <[email protected]>

Author: JohnCari

wrappers/src/fdw/clickhouse_fdw/README.md

@@ -11,6 +11,7 @@ This is a foreign data wrapper for [ClickHouse](https://clickhouse.com/). It is
 
 | Version | Date       | Notes                                                |
 | ------- | ---------- | ---------------------------------------------------- |
+| 0.1.10  | 2026-02-04 | Implement re_scan() for nested loop joins            |
 | 0.1.9   | 2025-11-08 | Added stream_buffer_size foreign table option        |
 | 0.1.8   | 2025-10-27 | Refactor to read rows with async streaming           |
 | 0.1.7   | 2025-05-22 | Added more data types support                        |

wrappers/src/fdw/clickhouse_fdw/clickhouse_fdw.rs

@@ -364,7 +364,7 @@ where
 }
 
 #[wrappers_fdw(
-    version = "0.1.9",
+    version = "0.1.10",
     author = "Supabase",
     website = "https://github.com/supabase/wrappers/tree/main/wrappers/src/fdw/clickhouse_fdw",
     error_type = "ClickHouseFdwError"
@@ -386,6 +386,7 @@ pub(crate) struct ClickHouseFdw {
     tgt_cols: Vec<Column>,
     sql_query: String,
     params: Vec<Qual>,
+    stream_buffer_size: usize,
 }
 
 impl ClickHouseFdw {
@@ -547,6 +548,33 @@ impl ClickHouseFdw {
             Ok(())
         }
     }
+
+    // Helper method to set up the streaming mechanism
+    fn setup_streaming(&mut self) -> ClickHouseFdwResult<()> {
+        // Create bounded channel
+        let (tx, rx) =
+            channel::bounded::<ClickHouseFdwResult<Option<ConvertedRow>>>(self.stream_buffer_size);
+        self.row_receiver = Some(rx);
+
+        // Clone data needed by the async task
+        let conn_str = self.conn_str.clone();
+        let sql = self.sql_query.clone();
+        let tgt_cols = self.tgt_cols.clone();
+        let params = self.params.clone();
+        let tx_clone = tx.clone();
+
+        // Spawn the async streaming task
+        let streaming_task = self.rt.spawn(async move {
+            stream_data_to_channel(conn_str, sql, tgt_cols, params, tx_clone).await;
+        });
+
+        self.streaming_task = Some(streaming_task);
+
+        // Fetch the first row to initialize the scan
+        self.fetch_next_row()?;
+
+        Ok(())
+    }
 }
 
 impl ForeignDataWrapper<ClickHouseFdwError> for ClickHouseFdw {
@@ -578,6 +606,7 @@ impl ForeignDataWrapper<ClickHouseFdwError> for ClickHouseFdw {
             tgt_cols: Vec::new(),
             sql_query: String::new(),
             params: Vec::new(),
+            stream_buffer_size: 1024,
         })
     }
 
@@ -597,34 +626,14 @@ impl ForeignDataWrapper<ClickHouseFdwError> for ClickHouseFdw {
         self.current_row_data = None;
 
         // get stream buffer size from options, with validation
-        let stream_buffer_size = options
+        self.stream_buffer_size = options
             .get("stream_buffer_size")
             .map(|s| s.parse::<usize>().map(|size| size.clamp(1, 100_000)))
             .transpose()
             .map_err(ClickHouseFdwError::ParseIntError)?
             .unwrap_or(1024);
 
-        // create bounded channel
-        let (tx, rx) =
-            channel::bounded::<ClickHouseFdwResult<Option<ConvertedRow>>>(stream_buffer_size);
-        self.row_receiver = Some(rx);
-
-        // clone data needed by the async task
-        let conn_str = self.conn_str.clone();
-        let sql = self.sql_query.clone();
-        let tgt_cols = self.tgt_cols.clone();
-        let params = self.params.clone();
-        let tx_clone = tx.clone();
-
-        // spawn the async streaming task
-        let streaming_task = self.rt.spawn(async move {
-            stream_data_to_channel(conn_str, sql, tgt_cols, params, tx_clone).await;
-        });
-
-        self.streaming_task = Some(streaming_task);
-
-        // fetch the first row to initialize the scan
-        self.fetch_next_row()?;
+        self.setup_streaming()?;
 
         Ok(())
     }
@@ -669,6 +678,24 @@ impl ForeignDataWrapper<ClickHouseFdwError> for ClickHouseFdw {
         Ok(())
     }
 
+    fn re_scan(&mut self) -> ClickHouseFdwResult<()> {
+        // Abort existing streaming task if running
+        if let Some(task) = self.streaming_task.take() {
+            task.abort();
+        }
+
+        // Drop existing channel receiver
+        self.row_receiver = None;
+
+        // Reset state
+        self.is_scan_complete = false;
+        self.current_row_data = None;
+
+        self.setup_streaming()?;
+
+        Ok(())
+    }
+
     fn begin_modify(&mut self, options: &HashMap<String, String>) -> ClickHouseFdwResult<()> {
         self.table = require_option("table", options)?.to_string();
         self.rowid_col = require_option("rowid_column", options)?.to_string();

wrappers/src/fdw/clickhouse_fdw/tests.rs

@@ -485,4 +485,157 @@ mod tests {
                 .is_empty());
         });
     }
+
+    #[pg_test]
+    fn clickhouse_join_test() {
+        Spi::connect_mut(|c| {
+            let clickhouse_pool = ch::Pool::new("tcp://default:default@localhost:9000/default");
+
+            let rt = create_async_runtime().expect("failed to create runtime");
+            let mut handle = rt
+                .block_on(async { clickhouse_pool.get_handle().await })
+                .expect("handle");
+
+            // Create two tables in ClickHouse for join testing
+            rt.block_on(async {
+                handle.execute("DROP TABLE IF EXISTS join_t1").await?;
+                handle.execute("DROP TABLE IF EXISTS join_t2").await?;
+                handle
+                    .execute("CREATE TABLE join_t1 (k Int16) engine = Memory")
+                    .await?;
+                handle
+                    .execute("CREATE TABLE join_t2 (k Int16) engine = Memory")
+                    .await?;
+                // Insert values (1, 2, 3) into each table
+                handle
+                    .execute("INSERT INTO join_t1 VALUES (1), (2), (3)")
+                    .await?;
+                handle
+                    .execute("INSERT INTO join_t2 VALUES (1), (2), (3)")
+                    .await
+            })
+            .expect("join tables in ClickHouse");
+
+            // Set up FDW and server
+            c.update(
+                r#"CREATE FOREIGN DATA WRAPPER clickhouse_wrapper
+                         HANDLER click_house_fdw_handler VALIDATOR click_house_fdw_validator"#,
+                None,
+                &[],
+            )
+            .unwrap();
+            c.update(
+                r#"CREATE SERVER my_clickhouse_server
+                         FOREIGN DATA WRAPPER clickhouse_wrapper
+                         OPTIONS (
+                           conn_string 'tcp://default:default@localhost:9000/default'
+                         )"#,
+                None,
+                &[],
+            )
+            .unwrap();
+
+            // Create foreign tables for join testing
+            c.update(
+                r#"
+                  CREATE FOREIGN TABLE join_t1 (
+                    k smallint
+                  )
+                  SERVER my_clickhouse_server
+                  OPTIONS (table 'join_t1')
+                "#,
+                None,
+                &[],
+            )
+            .unwrap();
+            c.update(
+                r#"
+                  CREATE FOREIGN TABLE join_t2 (
+                    k smallint
+                  )
+                  SERVER my_clickhouse_server
+                  OPTIONS (table 'join_t2')
+                "#,
+                None,
+                &[],
+            )
+            .unwrap();
+
+            // Test inner join - should return 3 rows, not 1
+            let inner_join_count: i64 = c
+                .select(
+                    "SELECT COUNT(*) FROM join_t1 JOIN join_t2 ON join_t1.k = join_t2.k",
+                    None,
+                    &[],
+                )
+                .unwrap()
+                .first()
+                .get_one::<i64>()
+                .unwrap()
+                .unwrap();
+            assert_eq!(
+                inner_join_count, 3,
+                "Inner join should return 3 rows, got {inner_join_count}",
+            );
+
+            // Verify inner join values are correct
+            let result = c
+                .select(
+                    "SELECT join_t1.k as k1, join_t2.k as k2 FROM join_t1 JOIN join_t2 ON join_t1.k = join_t2.k ORDER BY join_t1.k",
+                    None,
+                    &[],
+                )
+                .unwrap();
+            assert_eq!(result.len(), 3, "Inner join should return 3 rows");
+            let mut values: Vec<(i16, i16)> = Vec::new();
+            for row in result {
+                let k1: i16 = row.get_by_name("k1").unwrap().unwrap();
+                let k2: i16 = row.get_by_name("k2").unwrap().unwrap();
+                values.push((k1, k2));
+            }
+            assert_eq!(
+                values,
+                vec![(1, 1), (2, 2), (3, 3)],
+                "Inner join values incorrect"
+            );
+
+            // Test left join - should return 3 rows with correct values
+            let left_join_count: i64 = c
+                .select(
+                    "SELECT COUNT(*) FROM join_t1 LEFT JOIN join_t2 ON join_t1.k = join_t2.k",
+                    None,
+                    &[],
+                )
+                .unwrap()
+                .first()
+                .get_one::<i64>()
+                .unwrap()
+                .unwrap();
+            assert_eq!(
+                left_join_count, 3,
+                "Left join should return 3 rows, got {left_join_count}",
+            );
+
+            // Test right join - should return 3 rows with correct values
+            let right_join_count: i64 = c
+                .select(
+                    "SELECT COUNT(*) FROM join_t1 RIGHT JOIN join_t2 ON join_t1.k = join_t2.k",
+                    None,
+                    &[],
+                )
+                .unwrap()
+                .first()
+                .get_one::<i64>()
+                .unwrap()
+                .unwrap();
+            assert_eq!(
+                right_join_count, 3,
+                "Right join should return 3 rows, got {right_join_count}",
+            );
+
+            // Cleanup
+            c.update("DROP FOREIGN TABLE join_t1", None, &[]).unwrap();
+            c.update("DROP FOREIGN TABLE join_t2", None, &[]).unwrap();
+        });
+    }
 }