CI.yml

name: CI

permissions:
  id-token: write

on:
  push:
    branches:
      - main
    tags:
      - "*"
  pull_request:
  workflow_dispatch:
    inputs:
      # Latest commit to include with the release. If omitted, use the latest commit on the main branch.
      sha:
        description: Commit SHA
        type: string

defaults:
  run:
    shell: bash

env:
  PYTHON_VERSION: '3.9'

jobs:
  create-sdist:
    runs-on: ubuntu-latest
    strategy:
      fail-fast: false
      matrix:
        package: [polars_ds]

    steps:
      - uses: actions/checkout@v4
        with:
          ref: ${{ inputs.sha }}

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: ${{ env.PYTHON_VERSION }}

      - name: Create source distribution
        uses: PyO3/maturin-action@v1
        with:
          command: sdist
          args: >
            --manifest-path Cargo.toml
            --out dist

      - name: Test sdist
        run: |
          pip install --force-reinstall --verbose dist/*.tar.gz
          pip install typing_extensions
          pip install -r requirements.txt
          python -c 'import polars_ds as pds'
          python -c 'from polars_ds import linear_models'
          python -c 'from polars_ds.ts_features import *'
          python -c 'from polars_ds.spatial import *'
          python -c 'from polars_ds.sample_and_split import *'

      - name: Upload sdist
        uses: actions/upload-artifact@v4
        with:
          name: sdist-${{ matrix.package }}
          path: dist/*.tar.gz

      - name: Test Correctness
        run: | 
          python -m pip install --upgrade pip
          pip install pytest
          pip install .
          pip install -r tests/requirements-test.txt
          pytest tests/test_*

      - name: Test Notebooks
        run: | 
          python -m pip install --upgrade pip
          pip install jupyter ipython ipykernel nbconvert
          pip install -r tests/requirements-test.txt
          jupyter execute examples/basics.ipynb
          jupyter execute examples/diagnosis.ipynb
          jupyter execute examples/sample_and_split.ipynb


  build-wheels:
    runs-on: ${{ matrix.os }}
    needs: [create-sdist]
    strategy:
      fail-fast: false
      matrix:
        package: [polars_ds]
        os: [ubuntu-latest, macos-13, windows-latest]
        architecture: [x86-64, aarch64]
        exclude:
          - os: windows-latest
            architecture: aarch64

    steps:
      - uses: actions/checkout@v4
        with:
          ref: ${{ inputs.sha }}

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: ${{ env.PYTHON_VERSION }}

      - name: Determine CPU features for x86-64
        id: features
        if: matrix.architecture == 'x86-64'
        env:
          IS_MACOS: ${{ matrix.os == 'macos-13' }}
        run: |
          if [[ "$IS_MACOS" = true ]]; then
            FEATURES=+sse3,+ssse3,+sse4.1,+sse4.2,+popcnt,+avx,+fma,+pclmulqdq
          else
            FEATURES=+sse3,+ssse3,+sse4.1,+sse4.2,+popcnt,+avx,+avx2,+fma,+bmi1,+bmi2,+lzcnt,+pclmulqdq
          fi
          echo "features=$FEATURES" >> $GITHUB_OUTPUT

      - name: Set RUSTFLAGS for x86-64
        if: matrix.architecture == 'x86-64'
        env:
          FEATURES: ${{ steps.features.outputs.features }}
        run: echo "RUSTFLAGS=-C target-feature=${{ steps.features.outputs.features }}" >> $GITHUB_ENV

      - name: Set Rust target for aarch64
        if: matrix.architecture == 'aarch64'
        id: target
        run: |
          TARGET=${{ matrix.os == 'macos-13' && 'aarch64-apple-darwin' || 'aarch64-unknown-linux-gnu'}}
          echo "target=$TARGET" >> $GITHUB_OUTPUT

      - name: Set jemalloc for aarch64 Linux
        if: matrix.architecture == 'aarch64' && matrix.os == 'ubuntu-latest'
        run: |
          echo "JEMALLOC_SYS_WITH_LG_PAGE=16" >> $GITHUB_ENV

      - name: Build wheel
        uses: PyO3/maturin-action@v1
        with:
          command: build
          target: ${{ steps.target.outputs.target }}
          args: >
            --release
            --manifest-path Cargo.toml
            --out dist
          manylinux: ${{ matrix.architecture == 'aarch64' && '2_24' || 'auto' }}

      - name: Test wheel
        # Only test on x86-64 for now as this matches the runner architecture
        if: matrix.architecture == 'x86-64'
        run: |
          pip install --force-reinstall --verbose dist/*.whl
          pip install typing_extensions
          python -c 'import polars_ds'

      - name: Upload wheel
        uses: actions/upload-artifact@v4
        with:
          name: wheel-${{ matrix.package }}-${{ matrix.os }}-${{ matrix.architecture }}
          path: dist/*.whl

  release:
    name: Release
    runs-on: ubuntu-latest
    if: "startsWith(github.ref, 'refs/tags/')"
    needs: [build-wheels, create-sdist]
    permissions:
        id-token: write
    steps:
      - uses: actions/download-artifact@v4
        with:
          pattern: wheel-*
          merge-multiple: true

      - uses: actions/download-artifact@v4
        with:
          pattern: sdist-*
          merge-multiple: true
          
      - name: Publish to PyPI
        uses: PyO3/maturin-action@v1
        with:
          command: upload
          args: --non-interactive --skip-existing *