Feat/yugabyte

.github/workflows/IntegrationTests.yml

@@ -38,6 +38,119 @@ jobs:
           clang-format --dump-config
           make format-check
 
+  linux-yugabyte:
+    name: YugabyteDB Tests
+    needs: format-check
+    runs-on: ubuntu-latest
+
+    env:
+      GEN: ninja
+      CC: 'ccache gcc'
+      CXX: 'ccache g++'
+      CCACHE_DIR: ${{ github.workspace }}/ccache
+      VCPKG_TARGET_TRIPLET: x64-linux-release
+      VCPKG_HOST_TRIPLET: x64-linux-release
+      VCPKG_TOOLCHAIN_PATH: ${{ github.workspace }}/vcpkg/scripts/buildsystems/vcpkg.cmake
+
+    steps:
+    - name: Checkout
+      uses: actions/checkout@v6
+      with:
+        fetch-depth: 0
+        submodules: 'true'
+
+    - name: Install Dependencies
+      run: |
+        sudo apt-get update -y -q -o=Dpkg::Use-Pty=0
+        sudo apt-get install -y -q -o=Dpkg::Use-Pty=0 \
+          build-essential \
+          ccache \
+          cmake \
+          ninja-build \
+          postgresql-client
+
+    - name: Start YugabyteDB
+      run: |
+        docker run -d --name yugabyte \
+          -p 5433:5433 -p 7000:7000 -p 9000:9000 \
+          yugabytedb/yugabyte:2025.1.4.0-b103 \
+          bin/yugabyted start --daemon=false
+
+    - name: Cache Key
+      id: cache_key
+      run: |
+        DUCKDB_VERSION=$(cd duckdb && git rev-parse --short HEAD)
+        EXT_VERSION=$(git rev-parse --short HEAD)
+        KEY="${{ runner.os }}-${{ runner.arch }}-${DUCKDB_VERSION}-${EXT_VERSION}-yugabyte"
+        echo "value=${KEY}" >> "${GITHUB_OUTPUT}"
+
+    - name: Restore Cache
+      uses: actions/cache/restore@v5
+      with:
+        path: ${{ github.workspace }}/ccache
+        key: ${{ steps.cache_key.outputs.value }}
+
+    - name: Setup vcpkg
+      uses: lukka/run-vcpkg@v11.1
+      with:
+        vcpkgGitCommitId: 84bab45d415d22042bd0b9081aea57f362da3f35
+
+    - name: Build extension
+      run: |
+        make release
+
+    - name: Save Cache
+      uses: actions/cache/save@v5
+      with:
+        path: ${{ github.workspace }}/ccache
+        key: ${{ steps.cache_key.outputs.value }}
+
+    - name: Wait for YugabyteDB readiness
+      env:
+        PGHOST: localhost
+        PGPORT: '5433'
+        PGUSER: yugabyte
+        PGPASSWORD: yugabyte
+      run: |
+        echo "Waiting for YugabyteDB to accept connections..."
+        for i in $(seq 1 40); do
+          if psql -d yugabyte -c "SELECT count(*) FROM yb_servers()" 2>/dev/null; then
+            echo "YugabyteDB cluster is ready (attempt $i)"
+            psql -d yugabyte -c "SELECT host, port, node_type FROM yb_servers()"
+            exit 0
+          fi
+          echo "Attempt $i/40 — waiting 10s..."
+          sleep 10
+        done
+        echo "ERROR: YugabyteDB did not become ready after 400 seconds"
+        docker logs yugabyte 2>&1 | tail -100
+        exit 1
+
+    - name: Setup YugabyteDB test data
+      env:
+        PGHOST: localhost
+        PGPORT: '5433'
+        PGUSER: yugabyte
+        PGPASSWORD: yugabyte
+      run: |
+        source ./create-yugabyte-tables.sh
+
+    - name: Run YugabyteDB tests
+      env:
+        PGHOST: localhost
+        PGPORT: '5433'
+        PGUSER: yugabyte
+        PGPASSWORD: yugabyte
+        YUGABYTE_TEST_DATABASE_AVAILABLE: 1
+        LOCAL_EXTENSION_REPO: 'build/release/repository'
+      run: |
+        make test
+
+    - name: YugabyteDB logs on failure
+      if: failure()
+      run: |
+        docker logs yugabyte 2>&1 | tail -200
+
   linux-tests:
     name: Linux Tests
     needs: format-check
@@ -262,7 +375,7 @@ jobs:
         threadsan: [0, 1]
 
     env:
-      CMAKE_BUILD_PARALLEL_LEVEL: 2
+      CMAKE_BUILD_PARALLEL_LEVEL: 1
       CC: 'ccache gcc'
       CXX: 'ccache g++'
       CCACHE_DIR: ${{ github.workspace }}/ccache
@@ -367,11 +480,8 @@ jobs:
     - name: Dependencies
       shell: bash
       run: |
-        choco install \
-          ccache \
-          make \
-          ninja \
-          -y --force --no-progress
+        choco install ccache make -y --force --no-progress
+        choco install ninja -y --force --no-progress || pip install ninja
 
     - name: Build Environment
       shell: bash

.gitignore

@@ -5,3 +5,5 @@ cmake-build-debug
 .vscode
 .cache
 duckdb_unittest_tempdir
+# CocoIndex Code (ccc)
+/.cocoindex_code/

create-yugabyte-tables.sh

@@ -0,0 +1,198 @@
+#!/bin/bash
+set -e
+set -x
+
+# Create test tables on YugabyteDB for integration tests.
+# Expects YSQL connection via PGHOST/PGPORT/PGUSER/PGPASSWORD env vars.
+#
+# YugabyteDB can be slow to accept connections after startup.
+# Retry with backoff before giving up.
+
+MAX_RETRIES=30
+RETRY_DELAY=5
+
+echo "Waiting for YugabyteDB to accept connections..."
+for i in $(seq 1 $MAX_RETRIES); do
+  if psql -d yugabyte -c "SELECT 1" >/dev/null 2>&1; then
+    echo "YugabyteDB is ready (attempt $i)"
+    break
+  fi
+  if [ "$i" -eq "$MAX_RETRIES" ]; then
+    echo "ERROR: YugabyteDB did not become ready after $((MAX_RETRIES * RETRY_DELAY)) seconds"
+    exit 1
+  fi
+  echo "Attempt $i/$MAX_RETRIES failed, retrying in ${RETRY_DELAY}s..."
+  sleep $RETRY_DELAY
+done
+
+dropdb --if-exists postgresscanner || true
+createdb postgresscanner
+
+# Hash-partitioned table (default for YugabyteDB) — 100k rows to exercise parallel scan
+psql -d postgresscanner -c "
+CREATE TABLE hash_test (
+  id INTEGER PRIMARY KEY,
+  name TEXT,
+  value INTEGER
+);
+INSERT INTO hash_test SELECT g, 'row_' || g, g * 10 FROM generate_series(1, 100000) g;
+ANALYZE hash_test;
+"
+
+# Wide table with various types to stress the COPY path
+psql -d postgresscanner -c "
+CREATE TABLE wide_test (
+  id INTEGER PRIMARY KEY,
+  col_text TEXT,
+  col_int BIGINT,
+  col_float DOUBLE PRECISION,
+  col_bool BOOLEAN,
+  col_ts TIMESTAMP,
+  col_date DATE
+);
+INSERT INTO wide_test
+SELECT g,
+       'text_' || g,
+       g * 100000::BIGINT,
+       g * 3.14159,
+       (g % 2 = 0),
+       '2024-01-01'::TIMESTAMP + (g || ' seconds')::INTERVAL,
+       '2024-01-01'::DATE + g
+FROM generate_series(1, 50000) g;
+ANALYZE wide_test;
+"
+
+# Simple test table for attach/detach cycles
+psql -d postgresscanner -c "
+CREATE TABLE test (i INTEGER);
+INSERT INTO test VALUES (1), (2), (3), (NULL);
+"
+
+# Null test table
+psql -d postgresscanner -c "
+CREATE TABLE nulltest (
+  c1 INTEGER, c2 INTEGER, c3 INTEGER, c4 INTEGER, c5 INTEGER,
+  c6 INTEGER, c7 INTEGER, c8 INTEGER, c9 INTEGER, c10 INTEGER
+);
+INSERT INTO nulltest VALUES (1, 2, 3, 4, 5, 6, 7, 8, 9, 10);
+INSERT INTO nulltest VALUES (1, NULL, 3, 4, NULL, 6, 7, 8, NULL, 10);
+INSERT INTO nulltest VALUES (NULL, NULL, 3, 4, 5, 6, 7, NULL, NULL, NULL);
+INSERT INTO nulltest VALUES (NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL);
+"
+
+# Multi-column hash key table — tests compound partition key parallel scan
+psql -d postgresscanner -c "
+CREATE TABLE multi_hash (
+  region TEXT,
+  id INTEGER,
+  data TEXT,
+  PRIMARY KEY (region, id)
+);
+INSERT INTO multi_hash SELECT 'region_' || (g % 5), g, 'data_' || g FROM generate_series(1, 10000) g;
+ANALYZE multi_hash;
+"
+
+# Range-partitioned tables (no hash key — yb_num_hash_key_columns = 0)
+# These must fall back to single-threaded scan, NOT yb_hash_code() ranges.
+
+# Single-column range key
+psql -d postgresscanner -c "
+CREATE TABLE range_single (
+  id INTEGER,
+  name TEXT,
+  value INTEGER,
+  PRIMARY KEY (id ASC)
+);
+INSERT INTO range_single SELECT g, 'range_' || g, g * 10 FROM generate_series(1, 10000) g;
+ANALYZE range_single;
+"
+
+# Compound range key (timeseries pattern)
+psql -d postgresscanner -c "
+CREATE TABLE range_ts (
+  ts TIMESTAMP,
+  sensor_id INTEGER,
+  reading DOUBLE PRECISION,
+  PRIMARY KEY (ts ASC, sensor_id ASC)
+);
+INSERT INTO range_ts
+SELECT '2024-01-01'::TIMESTAMP + (g || ' seconds')::INTERVAL,
+       g % 100,
+       random() * 1000
+FROM generate_series(1, 20000) g;
+ANALYZE range_ts;
+"
+
+# Range key with DESC ordering
+psql -d postgresscanner -c "
+CREATE TABLE range_desc (
+  created_at TIMESTAMP,
+  id INTEGER,
+  payload TEXT,
+  PRIMARY KEY (created_at DESC, id DESC)
+);
+INSERT INTO range_desc
+SELECT '2024-06-01'::TIMESTAMP - (g || ' seconds')::INTERVAL,
+       g,
+       'payload_' || g
+FROM generate_series(1, 15000) g;
+ANALYZE range_desc;
+"
+
+# Colocated database — all tables share a single tablet (no hash partitioning)
+# This exercises the non-parallel scan fallback and tests that yb_table_properties
+# returns 0 tablets for colocated tables.
+# YugabyteDB colocated databases must be created with colocation=true at CREATE DATABASE time.
+dropdb --if-exists postgresscanner_colocated || true
+psql -d yugabyte -c "CREATE DATABASE postgresscanner_colocated WITH colocation = true"
+psql -d postgresscanner_colocated -c "
+CREATE TABLE coloc_test (
+  id INTEGER PRIMARY KEY,
+  name TEXT,
+  value INTEGER
+);
+INSERT INTO coloc_test SELECT g, 'coloc_' || g, g * 10 FROM generate_series(1, 10000) g;
+ANALYZE coloc_test;
+"
+
+psql -d postgresscanner_colocated -c "
+CREATE TABLE coloc_wide (
+  id INTEGER PRIMARY KEY,
+  col_text TEXT,
+  col_int BIGINT,
+  col_float DOUBLE PRECISION,
+  col_bool BOOLEAN
+);
+INSERT INTO coloc_wide
+SELECT g, 'text_' || g, g * 100000::BIGINT, g * 3.14, (g % 2 = 0)
+FROM generate_series(1, 5000) g;
+ANALYZE coloc_wide;
+"
+
+# Non-colocated table in the same database for contrast
+psql -d postgresscanner_colocated -c "
+CREATE TABLE non_coloc_test (
+  id INTEGER PRIMARY KEY,
+  name TEXT,
+  value INTEGER
+) WITH (colocation = false);
+INSERT INTO non_coloc_test SELECT g, 'nocoloc_' || g, g * 10 FROM generate_series(1, 10000) g;
+ANALYZE non_coloc_test;
+"
+
+echo "YugabyteDB test tables created successfully"
+echo ""
+echo "  Database: postgresscanner (non-colocated, default)"
+echo "    hash_test:     100,000 rows (hash-partitioned, single key)"
+echo "    wide_test:      50,000 rows (hash-partitioned, multiple types)"
+echo "    multi_hash:     10,000 rows (hash-partitioned, compound key)"
+echo "    range_single:   10,000 rows (range ASC, single key)"
+echo "    range_ts:       20,000 rows (range ASC compound key, timeseries)"
+echo "    range_desc:     15,000 rows (range DESC compound key)"
+echo "    test:                4 rows (simple)"
+echo "    nulltest:            4 rows (null patterns)"
+echo ""
+echo "  Database: postgresscanner_colocated"
+echo "    coloc_test:     10,000 rows (colocated, single tablet)"
+echo "    coloc_wide:      5,000 rows (colocated, multiple types)"
+echo "    non_coloc_test: 10,000 rows (non-colocated in colocated db)"

src/include/postgres_connection.hpp

@@ -31,6 +31,7 @@ struct OwnedPostgresConnection {
 
 	PGconn *connection;
 	mutex connection_lock;
+	PostgresInstanceType instance_type = PostgresInstanceType::POSTGRES;
 };
 
 class PostgresConnection {
@@ -63,6 +64,8 @@ class PostgresConnection {
 
 	void BeginCopyTo(ClientContext &context, PostgresCopyState &state, PostgresCopyFormat format,
 	                 const string &schema_name, const string &table_name, const vector<string> &column_names);
+	void CommitAndRestartCopy(ClientContext &context, PostgresCopyState &state, PostgresCopyFormat format,
+	                          const string &schema_name, const string &table_name, const vector<string> &column_names);
 	void CopyData(data_ptr_t buffer, idx_t size);
 	void CopyData(PostgresBinaryWriter &writer);
 	void CopyData(PostgresTextWriter &writer);

src/include/postgres_scanner.hpp

@@ -50,6 +50,10 @@ struct PostgresBindData : public FunctionData {
 	bool use_text_protocol = false;
 	idx_t max_threads = 1;
 
+	idx_t yb_num_tablets = 0;
+	idx_t yb_num_hash_key_columns = 0;
+	vector<string> yb_hash_partition_columns;
+
 public:
 	void SetTablePages(idx_t approx_num_pages);
 

src/include/postgres_version.hpp

@@ -12,7 +12,7 @@
 
 namespace duckdb {
 
-enum class PostgresInstanceType { UNKNOWN, POSTGRES, AURORA, REDSHIFT };
+enum class PostgresInstanceType { UNKNOWN, POSTGRES, AURORA, REDSHIFT, YUGABYTE };
 
 struct PostgresVersion {
 	PostgresVersion() {