build_python_connect.yml

#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#

name: Python Connect Parity (Spark 4.3.0)

# Cancel an in-progress run when a newer commit supersedes it (per workflow + ref).
concurrency:
  group: ${{ github.workflow }}-${{ github.ref }}
  cancel-in-progress: true

# Mirrors Apache Spark's build_python_connect.yml approach, adapted for the
# Rust-backed client: download the Spark 4.3.0 Connect server, start it, pull the
# official v4.3.0 connect test files, and run them through our client - the gate is
# "every test passes except the checked-in known-environmental-failures manifest".
# A separate job enforces Rust test coverage across all crates.

on:
  push:
    branches: [master, "branch-*"]
  pull_request:
    branches: [master, "branch-*"]
  workflow_dispatch:

env:
  SPARK_VERSION: "4.3.0"
  RUST_PARITY_LOG_DIR: /tmp/parity-test-logs
  SPARK_CONNECT_TESTING_REMOTE: sc://localhost:15002
  SPARK_TESTING: "1"
  SPARK_SKIP_CONNECT_COMPAT_TESTS: "1"   # force remote-only (pure connect, no local JVM)

jobs:
  # ---------------------------------------------------------------------------
  # Parity gate: the official connect suite through our client, against a real
  # Spark 4.3.0 Connect server (see scripts/run_official_tests.py).
  # ---------------------------------------------------------------------------
  connect-parity:
    name: Connect parity (official suite via our client, Python ${{ matrix.python-version }})
    runs-on: ubuntu-latest
    # The official suite runs serially (--jobs 1) and takes ~45-60 min; bound the job so
    # a genuinely hung test fails it instead of running to GitHub's 6-hour default.
    timeout-minutes: 90
    strategy:
      fail-fast: false
      matrix:
        python-version: ["3.11"]
    steps:
      - uses: actions/checkout@v4

      - name: Set up Python ${{ matrix.python-version }}
        uses: actions/setup-python@v5
        with:
          python-version: ${{ matrix.python-version }}
          cache: pip

      - name: Set up JDK 17
        uses: actions/setup-java@v4
        with:
          distribution: temurin
          java-version: "17"

      - name: Install Rust toolchain
        # rustup is preinstalled; avoid third-party toolchain actions (ASF Actions
        # policy allowlists actions/* only).
        run: |
          rustup toolchain install stable --profile minimal
          rustup default stable

      - name: Cache cargo & target
        uses: actions/cache@v4
        with:
          path: |
            ~/.cargo/registry
            ~/.cargo/git
            target
          key: ${{ runner.os }}-cargo-parity-${{ hashFiles('**/Cargo.lock') }}
          restore-keys: ${{ runner.os }}-cargo-parity-

      - name: Install system dependencies
        run: |
          sudo apt-get update
          sudo apt-get install -y protobuf-compiler

      - name: Install Python dependencies
        run: |
          python -m pip install --upgrade pip
          # Match the supported pandas 2.x environment used by pandas-parity.yml.
          # Spark 4.2's tests and Python workers are not compatible with pandas 3.x.
          python -m pip install pytest 'pandas>=2.2,<3.0' pyarrow numpy grpcio grpcio-status googleapis-common-protos cloudpickle zstandard py4j

      - name: Build the Rust extension into the skin
        # The parity diff drives both clients via PYTHONPATH (our ./python for ours,
        # the spark checkout for reference), so we do NOT install a wheel (which would
        # collide in site-packages with the reference pyspark). We just build the
        # extension and drop it into python/pyspark so ./python is a working package.
        run: |
          cargo build -p pyspark-rs --release
          cp target/release/lib_pyspark.so python/pyspark/_pyspark.so

      - name: Cache Spark ${{ env.SPARK_VERSION }} distribution
        uses: actions/cache@v4
        with:
          path: ~/spark-dist
          key: spark-dist-${{ env.SPARK_VERSION }}-hadoop3

      - name: Download Apache Spark ${{ env.SPARK_VERSION }} (Connect server)
        id: spark-download
        run: |
          PKG="spark-${SPARK_VERSION}-bin-hadoop3-connect"
          if [ ! -d "${HOME}/spark-dist/spark-4.3.0-bin-hadoop3"* ]; then
            URL="https://archive.apache.org/dist/spark/spark-${SPARK_VERSION}/${PKG}.tgz"
            mkdir -p "${HOME}/spark-dist"
            echo "Downloading ${URL}"
            wget -q -O "${HOME}/spark-dist/${PKG}.tgz" "${URL}"
            tar -xzf "${HOME}/spark-dist/${PKG}.tgz" -C "${HOME}/spark-dist"
            rm -f "${HOME}/spark-dist/${PKG}.tgz"
          else
            echo "Using cached Spark distribution."
          fi
          SPARK_HOME=$(find "${HOME}/spark-dist" -maxdepth 1 -type d -name 'spark-4.3.0-bin-hadoop3*' | head -1)
          echo "SPARK_HOME=${SPARK_HOME}" >> "$GITHUB_OUTPUT"

      - name: Cache Spark ${{ env.SPARK_VERSION }} source (official test files)
        uses: actions/cache@v4
        with:
          path: ~/spark-source
          key: spark-source-${{ env.SPARK_VERSION }}

      - name: Clone Apache Spark v${{ env.SPARK_VERSION }} (official test files + reference client)
        id: spark-source
        run: |
          if [ ! -d "${HOME}/spark-source/python" ]; then
            git clone --depth 1 --branch "v${SPARK_VERSION}" \
              https://github.com/apache/spark.git "${HOME}/spark-source"
          else
            echo "Using cached Spark source."
          fi
          echo "SPARK_SOURCE=${HOME}/spark-source" >> "$GITHUB_OUTPUT"

      - name: Start Spark Connect server
        env:
          # The server spawns Python workers for UDF/arrow/mapInPandas execution; they
          # need pandas + pyarrow (installed above) plus pyspark (from SPARK_HOME/python).
          PYSPARK_PYTHON: python
          PYSPARK_DRIVER_PYTHON: python
        run: |
          SPARK_HOME="${{ steps.spark-download.outputs.SPARK_HOME }}"
          echo "Starting Connect server on ${SPARK_CONNECT_TESTING_REMOTE}"
          # Give the driver extra heap: the parity gate runs all 84 files serially against
          # this one server, and a 1g default degrades late in the run (GC pressure),
          # which is what made the streaming listener/observation files flake.
          "${SPARK_HOME}/sbin/start-connect-server.sh" \
            --driver-memory 3g \
            --conf spark.sql.session.localRelationSizeLimit=1073741824 \
            > /tmp/connect-server.log 2>&1 &
          for i in $(seq 1 90); do
            if (echo > /dev/tcp/localhost/15002) 2>/dev/null; then
              echo "Connect server is up (after ${i}s)."; break
            fi
            if [ "$i" -eq 90 ]; then echo "Server failed to start:"; cat /tmp/connect-server.log; exit 1; fi
            sleep 1
          done

      # NOTE (SPARK-59032): the behavioral e2e tests used to run here, sharing this
      # server with the parity gate. They now run in their own job (connect-e2e,
      # below) so the parity run starts against a server nothing else has touched.
      # No test is skipped.

      - name: Parity - official connect suite through OUR Rust client
        id: parity
        env:
          PYTHONUNBUFFERED: "1"   # stream per-file progress into the live log
          RUST_PYSPARK_SO: ${{ github.workspace }}/python/pyspark/_pyspark.so
        run: |
          # pipefail so the script's exit code (nonzero on any unexpected failure) gates
          # the job - otherwise `| tee` would mask it and the gate would be falsely green.
          set -o pipefail
          # Run the official connect suite ONLY through the Rust transport plugin (which
          # genuinely exercises our client), requiring every test to pass except the
          # checked-in manifest of known *environmental* failures (scripts/
          # parity_known_failures.txt) - tests the reference client also fails in this
          # single-node pure-Connect environment. Now that the client is at full parity we
          # no longer run the reference client on every file just to derive that baseline,
          # which roughly halves this job. Regenerate the manifest with
          # scripts/gen_parity_skiplist.py when the pinned Spark version bumps.
          # Run serially (--jobs 1), matching Apache's own connect test runner
          # (./python/run-tests --parallelism=1): the streaming-query-listener tests wait
          # on real server-pushed progress events and are timeout-sensitive, so parallel
          # runs starve them under server contention. Serial removes that artifact.
          # The resource/artifact tests need a real multi-executor cluster, so they run
          # in the separate local-cluster phase below (mirroring Apache's
          # build_python_connect.yml, which runs exactly these under local-cluster) and
          # are excluded here where the server is plain local[*].
          python scripts/run_official_tests.py \
            --spark "${{ steps.spark-source.outputs.SPARK_SOURCE }}" \
            --remote "${SPARK_CONNECT_TESTING_REMOTE}" \
            --jobs 1 \
            --exclude "test_artifact.py,test_artifact_localcluster.py,test_parity_resources.py" \
            2>&1 | tee /tmp/parity_ours.log

      - name: Drop-in live-server tests (native Arrow handoff and API behaviors)
        # The parity gate above runs the reference pyspark over our transport, so the
        # drop-in's own Rust-backed classes are only exercised against a server here.
        working-directory: python
        env:
          SPARK_REMOTE: sc://localhost:15002
          PYTHONPATH: ${{ github.workspace }}/python
        run: python -m pytest tests/test_arrow_e2e.py tests/test_dropin_e2e.py tests/test_ml_e2e.py -q

      - name: Restart Connect server in local-cluster mode
        # Resource profiles and artifact distribution across executors can only be
        # exercised against a real multi-executor cluster, so - mirroring Apache's
        # build_python_connect.yml - bounce the server into local-cluster mode for those
        # tests (the parity run above uses the default local[*] and excludes these files).
        env:
          PYSPARK_PYTHON: python
          PYSPARK_DRIVER_PYTHON: python
        run: |
          SPARK_HOME="${{ steps.spark-download.outputs.SPARK_HOME }}"
          echo "Stopping the local[*] Connect server..."
          "${SPARK_HOME}/sbin/stop-connect-server.sh" || true
          pkill -f 'org.apache.spark.sql.connect.service' 2>/dev/null || true
          pkill -f 'spark-connect' 2>/dev/null || true
          for i in $(seq 1 30); do
            if ! (echo > /dev/tcp/localhost/15002) 2>/dev/null; then break; fi
            sleep 1
          done
          echo "Starting Connect server in local-cluster mode..."
          # The -connect distribution bundles the Connect jar, so no --packages is needed.
          "${SPARK_HOME}/sbin/start-connect-server.sh" \
            --master "local-cluster[2, 4, 1024]" \
            --driver-memory 3g \
            --conf spark.sql.session.localRelationSizeLimit=1073741824 \
            > /tmp/connect-server-localcluster.log 2>&1 &
          for i in $(seq 1 120); do
            if (echo > /dev/tcp/localhost/15002) 2>/dev/null; then
              echo "Connect server (local-cluster) is up (after ${i}s)."; break
            fi
            if [ "$i" -eq 120 ]; then
              echo "Server failed to start:"; cat /tmp/connect-server-localcluster.log; exit 1
            fi
            sleep 1
          done

      - name: Parity - resource & artifact tests (local-cluster)
        id: parity-localcluster
        env:
          PYTHONUNBUFFERED: "1"
          RUST_PYSPARK_SO: ${{ github.workspace }}/python/pyspark/_pyspark.so
        run: |
          set -o pipefail
          # Mirror Apache's local-cluster testnames: resource profiles + artifact
          # distribution across executors, run through our Rust client. --no-drift-check:
          # none of these are skiplisted.
          python scripts/run_official_tests.py \
            --spark "${{ steps.spark-source.outputs.SPARK_SOURCE }}" \
            --remote "${SPARK_CONNECT_TESTING_REMOTE}" \
            --jobs 1 \
            --no-drift-check \
            --only "pyspark.resource.tests.test_connect_resources,pyspark.sql.tests.connect.client.test_artifact,pyspark.sql.tests.connect.client.test_artifact_localcluster,pyspark.sql.tests.connect.test_parity_resources" \
            2>&1 | tee /tmp/parity_localcluster.log

      - name: Upload parity log + server log
        if: always()
        uses: actions/upload-artifact@v4
        with:
          name: parity-ours-py${{ matrix.python-version }}
          path: |
            /tmp/parity-test-logs/
            /tmp/parity_ours.log
            /tmp/parity_localcluster.log
            /tmp/connect-server.log
            /tmp/connect-server-localcluster.log
          if-no-files-found: warn

  # ---------------------------------------------------------------------------
  # Behavioral e2e tests (Rust drop-in against a live 4.3.0 server), in their OWN job
  # with a separate Connect server so they never share server state with the
  # parity gate. The parity gate then starts against a server nothing else has
  # touched, which avoids the SPARK-59032 arrow-UDF hang (a UDF that crashes a
  # Python worker on the shared server wedged worker reuse for later tests).
  # ---------------------------------------------------------------------------
  connect-e2e:
    name: Rust behavioral e2e
    runs-on: ubuntu-latest
    timeout-minutes: 45
    steps:
      - uses: actions/checkout@v4

      - name: Set up Python 3.11
        uses: actions/setup-python@v5
        with:
          python-version: "3.11"
          cache: pip

      - name: Set up JDK 17
        uses: actions/setup-java@v4
        with:
          distribution: temurin
          java-version: "17"

      - name: Install Rust toolchain (+ wasm32 target)
        run: |
          rustup toolchain install stable --profile minimal
          rustup default stable
          rustup target add wasm32-unknown-unknown

      - name: Cache cargo & target
        uses: actions/cache@v4
        with:
          path: |
            ~/.cargo/registry
            ~/.cargo/git
            target
          key: ${{ runner.os }}-cargo-e2e-${{ hashFiles('**/Cargo.lock') }}
          restore-keys: ${{ runner.os }}-cargo-e2e-

      - name: Install system dependencies
        run: |
          sudo apt-get update
          sudo apt-get install -y protobuf-compiler

      - name: Install Python dependencies (server-side UDF workers, WASM UDF packer)
        run: |
          python -m pip install --upgrade pip
          python -m pip install 'pandas>=2.2,<3.0' pyarrow numpy wasmtime "pyspark==${SPARK_VERSION}"

      - name: Cache Spark ${{ env.SPARK_VERSION }} distribution
        uses: actions/cache@v4
        with:
          path: ~/spark-dist
          key: spark-dist-${{ env.SPARK_VERSION }}-hadoop3

      - name: Download Apache Spark ${{ env.SPARK_VERSION }} (Connect server)
        id: spark-download
        run: |
          PKG="spark-${SPARK_VERSION}-bin-hadoop3-connect"
          if [ ! -d "${HOME}/spark-dist/spark-4.3.0-bin-hadoop3"* ]; then
            URL="https://archive.apache.org/dist/spark/spark-${SPARK_VERSION}/${PKG}.tgz"
            mkdir -p "${HOME}/spark-dist"
            echo "Downloading ${URL}"
            wget -q -O "${HOME}/spark-dist/${PKG}.tgz" "${URL}"
            tar -xzf "${HOME}/spark-dist/${PKG}.tgz" -C "${HOME}/spark-dist"
            rm -f "${HOME}/spark-dist/${PKG}.tgz"
          else
            echo "Using cached Spark distribution."
          fi
          SPARK_HOME=$(find "${HOME}/spark-dist" -maxdepth 1 -type d -name 'spark-4.3.0-bin-hadoop3*' | head -1)
          echo "SPARK_HOME=${SPARK_HOME}" >> "$GITHUB_OUTPUT"

      - name: Start Spark Connect server
        env:
          PYSPARK_PYTHON: python
          PYSPARK_DRIVER_PYTHON: python
        run: |
          SPARK_HOME="${{ steps.spark-download.outputs.SPARK_HOME }}"
          # A small senderMaxStreamSize makes the server end the response stream of every
          # result over 1 MiB before it is complete, so that the client's reattach path is
          # exercised (by default it only is for results over 1 GiB or 2 minutes).
          "${SPARK_HOME}/sbin/start-connect-server.sh" \
            --driver-memory 3g \
            --conf spark.sql.session.localRelationSizeLimit=1073741824 \
            --conf spark.connect.execute.reattachable.senderMaxStreamSize=1048576 \
            > /tmp/connect-server.log 2>&1 &
          for i in $(seq 1 90); do
            if (echo > /dev/tcp/localhost/15002) 2>/dev/null; then
              echo "Connect server is up (after ${i}s)."; break
            fi
            if [ "$i" -eq 90 ]; then echo "Server failed to start:"; cat /tmp/connect-server.log; exit 1; fi
            sleep 1
          done

      - name: Rust behavioral e2e tests (against the live server)
        env:
          SPARK_REMOTE: sc://localhost:15002
        run: |
          cargo test -p apache-spark-connect \
            --test e2e_functions --test e2e_columns --test e2e_dataframe \
            --test e2e_coverage --test e2e_linear_svc --test e2e_integration \
            -- --test-threads=1

      - name: Rust WASM UDF e2e test (against the live server)
        env:
          SPARK_REMOTE: sc://localhost:15002
        run: cargo test --manifest-path tests/wasm-udf/Cargo.toml --test e2e_wasm_udf

  # ---------------------------------------------------------------------------
  # E2E performance benchmark: our client vs reference, same 4.3.0 server.
  # ---------------------------------------------------------------------------
  connect-benchmark:
    name: E2E benchmark vs reference
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with:
          python-version: "3.11"
          cache: pip
      - uses: actions/setup-java@v4
        with:
          distribution: temurin
          java-version: "17"
      - name: Install Rust toolchain
        run: |
          rustup toolchain install stable --profile minimal
          rustup default stable
      - name: Cache cargo & target
        uses: actions/cache@v4
        with:
          path: |
            ~/.cargo/registry
            ~/.cargo/git
            target
          key: ${{ runner.os }}-cargo-bench-${{ hashFiles('**/Cargo.lock') }}
          restore-keys: ${{ runner.os }}-cargo-bench-
      - name: Install deps
        run: |
          sudo apt-get update && sudo apt-get install -y protobuf-compiler
          python -m pip install --upgrade pip
          python -m pip install 'pandas>=2.2,<3.0' pyarrow numpy grpcio grpcio-status googleapis-common-protos cloudpickle zstandard py4j "pyspark==${SPARK_VERSION}"
      - name: Build the Rust extension into the skin
        # Build the extension into ./python/pyspark so our client is usable via
        # PYTHONPATH. We do NOT install a wheel: our package and the reference
        # pyspark both own the `pyspark` namespace and would collide in site-packages.
        run: |
          cargo build -p pyspark-rs --release
          cp target/release/lib_pyspark.so python/pyspark/_pyspark.so
      - name: Cache Spark ${{ env.SPARK_VERSION }} distribution
        uses: actions/cache@v4
        with:
          path: ~/spark-dist
          key: spark-dist-${{ env.SPARK_VERSION }}-hadoop3
      - name: Cache Spark ${{ env.SPARK_VERSION }} source (official test files)
        uses: actions/cache@v4
        with:
          path: ~/spark-source
          key: spark-source-${{ env.SPARK_VERSION }}
      - name: Download & start Spark ${{ env.SPARK_VERSION }} server
        run: |
          PKG="spark-${SPARK_VERSION}-bin-hadoop3-connect"
          if [ ! -d "${HOME}/spark-dist/spark-4.3.0-bin-hadoop3"* ]; then
            mkdir -p "${HOME}/spark-dist"
            wget -q -O "/tmp/${PKG}.tgz" "https://archive.apache.org/dist/spark/spark-${SPARK_VERSION}/${PKG}.tgz"
            tar -xzf "/tmp/${PKG}.tgz" -C "${HOME}/spark-dist"
            rm -f "/tmp/${PKG}.tgz"
          fi
          if [ ! -d "${HOME}/spark-source/python" ]; then
            git clone --depth 1 --branch "v${SPARK_VERSION}" https://github.com/apache/spark.git "${HOME}/spark-source"
          fi
          SPARK_HOME=$(find "${HOME}/spark-dist" -maxdepth 1 -type d -name 'spark-4.3.0-bin-hadoop3*' | head -1)
          "${SPARK_HOME}/sbin/start-connect-server.sh" > /tmp/server.log 2>&1 &
          for i in $(seq 1 90); do (echo > /dev/tcp/localhost/15002) 2>/dev/null && break; sleep 1; done
      - name: Benchmark reference client
        run: |
          # Reference uses the pip-installed pyspark 4.3.0 (no ./python on the path).
          python scripts/benchmark_e2e.py --isolate --label reference --out /tmp/bench_reference.json
      - name: Benchmark our client
        run: |
          # ./python shadows the installed pyspark, so this exercises our client.
          PYTHONPATH="$(pwd)/python" \
            python scripts/benchmark_e2e.py --isolate --label ours --out /tmp/bench_ours.json
      - name: Compare
        run: python scripts/benchmark_e2e.py --compare /tmp/bench_reference.json /tmp/bench_ours.json | tee /tmp/bench_compare.txt
      - name: Upload benchmark results
        if: always()
        uses: actions/upload-artifact@v4
        with:
          name: benchmark-results
          path: |
            /tmp/bench_reference.json
            /tmp/bench_ours.json
            /tmp/bench_compare.txt
          if-no-files-found: warn

  # ---------------------------------------------------------------------------
  # Rust checks - must be fully green.
  # ---------------------------------------------------------------------------
  rust-checks:
    name: Rust checks
    runs-on: ubuntu-latest
    # Hard-fail on any compiler warning in our workspace crates (cargo caps dependency
    # lints to `allow`, so third-party deps are unaffected). Job-scoped on purpose: the
    # parity job above builds the pyspark-rs extension, which this must not gate.
    env:
      RUSTFLAGS: "-D warnings"
    steps:
      - uses: actions/checkout@v4
      - name: Install Rust toolchain
        run: |
          rustup toolchain install stable --profile minimal
          rustup default stable
      - name: Cache cargo & target
        uses: actions/cache@v4
        with:
          path: |
            ~/.cargo/registry
            ~/.cargo/git
            target
          key: ${{ runner.os }}-cargo-checks-${{ hashFiles('**/Cargo.lock') }}
          restore-keys: ${{ runner.os }}-cargo-checks-
      - name: Install system dependencies
        run: |
          sudo apt-get update && sudo apt-get install -y protobuf-compiler
      - name: Cargo build
        run: cargo build --release
      - name: Cargo test
        run: cargo test -p apache-spark-connect -p apache-spark-connect-core --release
      - name: No-stub audit
        run: bash scripts/audit_no_stubs.sh