Datafusion aggregate

.github/docker-compose.yaml

@@ -11,5 +11,7 @@ services:
         container_name: dask-worker
         image: daskdev/dask:latest
         command: dask-worker dask-scheduler:8786
+        environment:
+            EXTRA_CONDA_PACKAGES: "pyarrow>=4.0.0"  # required for parquet IO
         volumes:
             - /tmp:/tmp

.github/workflows/release.yml

@@ -16,11 +16,6 @@ jobs:
     if: github.repository == 'dask-contrib/dask-sql'
     steps:
       - uses: actions/checkout@v2
-      - name: Cache local Maven repository
-        uses: actions/cache@v2
-        with:
-          path: ~/.m2/repository
-          key: ${{ runner.os }}-maven-v1-jdk11-${{ hashFiles('**/pom.xml') }}
       - name: Set up Python
         uses: conda-incubator/setup-miniconda@v2
         with:
@@ -29,7 +24,7 @@ jobs:
           python-version: "3.8"
           channel-priority: strict
           activate-environment: dask-sql
-          environment-file: continuous_integration/environment-3.8-jdk11-dev.yaml
+          environment-file: continuous_integration/environment-3.8-dev.yaml
       - name: Install dependencies
         run: |
           pip install setuptools wheel twine

.github/workflows/test-upstream.yml

@@ -5,64 +5,23 @@ on:
   workflow_dispatch: # allows you to trigger the workflow run manually
 
 jobs:
-  build:
-    # This build step should be similar to the deploy build, to make sure we actually test
-    # the future deployable
-    name: Build the jar on ubuntu
-    runs-on: ubuntu-latest
-    if: github.repository == 'dask-contrib/dask-sql'
-    defaults:
-      run:
-        shell: bash -l {0}
-    steps:
-      - uses: actions/checkout@v2
-      - name: Cache local Maven repository
-        uses: actions/cache@v2
-        with:
-          path: ~/.m2/repository
-          key: ${{ runner.os }}-maven-v1-jdk11-${{ hashFiles('**/pom.xml') }}
-      - name: Set up Python
-        uses: conda-incubator/setup-miniconda@v2
-        with:
-          miniforge-variant: Mambaforge
-          use-mamba: true
-          python-version: "3.8"
-          channel-priority: strict
-          activate-environment: dask-sql
-          environment-file: continuous_integration/environment-3.8-jdk11-dev.yaml
-      - name: Install dependencies and build the jar
-        run: |
-          python setup.py build_ext
-      - name: Upload the jar
-        uses: actions/upload-artifact@v1
-        with:
-          name: jar
-          path: dask_sql/jar/DaskSQL.jar
-
   test-dev:
-    name: "Test upstream dev (${{ matrix.os }}, java: ${{ matrix.java }}, python: ${{ matrix.python }})"
-    needs: build
+    name: "Test upstream dev (${{ matrix.os }}, python: ${{ matrix.python }})"
     runs-on: ${{ matrix.os }}
     env:
-      CONDA_FILE: continuous_integration/environment-${{ matrix.python }}-jdk${{ matrix.java }}-dev.yaml
+      CONDA_FILE: continuous_integration/environment-${{ matrix.python }}-dev.yaml
     defaults:
       run:
         shell: bash -l {0}
     strategy:
       fail-fast: false
       matrix:
-        java: [8, 11]
         os: [ubuntu-latest, windows-latest]
         python: ["3.8", "3.9", "3.10"]
     steps:
       - uses: actions/checkout@v2
         with:
           fetch-depth: 0 # Fetch all history for all branches and tags.
-      - name: Cache local Maven repository
-        uses: actions/cache@v2
-        with:
-          path: ~/.m2/repository
-          key: ${{ runner.os }}-maven-v1-jdk${{ matrix.java }}-${{ hashFiles('**/pom.xml') }}
       - name: Set up Python
         uses: conda-incubator/setup-miniconda@v2
         with:
@@ -72,21 +31,12 @@ jobs:
           channel-priority: strict
           activate-environment: dask-sql
           environment-file: ${{ env.CONDA_FILE }}
-      - name: Download the pre-build jar
-        uses: actions/download-artifact@v1
-        with:
-          name: jar
-          path: dask_sql/jar/
       - name: Install hive testing dependencies for Linux
         if: matrix.os == 'ubuntu-latest'
         run: |
           mamba install -c conda-forge sasl>=0.3.1
           docker pull bde2020/hive:2.3.2-postgresql-metastore
           docker pull bde2020/hive-metastore-postgresql:2.3.0
-      - name: Set proper JAVA_HOME for Windows
-        if: matrix.os == 'windows-latest'
-        run: |
-          echo "JAVA_HOME=${{ env.CONDA }}\envs\dask-sql\Library" >> $GITHUB_ENV
       - name: Install upstream dev Dask / dask-ml
         run: |
           python -m pip install --no-deps git+https://github.com/dask/dask

.github/workflows/test.yml

@@ -1,7 +1,7 @@
 ---
 # Test the main branch and every pull request by
-# 1. building the jar on ubuntu
-# 2. testing code (using the build jar) on ubuntu and windows, with different java versions
+# 1. build dask_planner (Arrow DataFusion Rust bindings) on ubuntu
+# 2. testing code (using the build DataFusion bindings) on ubuntu and windows
 name: Test Python package
 on:
   push:
@@ -36,55 +36,20 @@ jobs:
         with:
           keyword: "[test-upstream]"
 
-  build:
-    # This build step should be similar to the deploy build, to make sure we actually test
-    # the future deployable
-    name: Build the jar on ubuntu
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v2
-      - name: Cache local Maven repository
-        uses: actions/cache@v2
-        with:
-          path: ~/.m2/repository
-          key: ${{ runner.os }}-maven-v1-jdk11-${{ hashFiles('**/pom.xml') }}
-      - name: Set up Python
-        uses: conda-incubator/setup-miniconda@v2
-        with:
-          miniforge-variant: Mambaforge
-          use-mamba: true
-          python-version: "3.8"
-          channel-priority: strict
-          activate-environment: dask-sql
-          environment-file: continuous_integration/environment-3.8-jdk11-dev.yaml
-      - name: Build the jar
-        run: |
-          python setup.py build_ext
-      - name: Upload the jar
-        uses: actions/upload-artifact@v1
-        with:
-          name: jar
-          path: dask_sql/jar/DaskSQL.jar
-
   test:
-    name: "Test (${{ matrix.os }}, java: ${{ matrix.java }}, python: ${{ matrix.python }})"
-    needs: [detect-ci-trigger, build]
+    name: "Build & Test (${{ matrix.os }}, python: ${{ matrix.python }}, Rust: ${{ matrix.toolchain }})"
+    needs: [detect-ci-trigger]
     runs-on: ${{ matrix.os }}
     env:
-      CONDA_FILE: continuous_integration/environment-${{ matrix.python }}-jdk${{ matrix.java }}-dev.yaml
+      CONDA_FILE: continuous_integration/environment-${{ matrix.python }}-dev.yaml
     strategy:
       fail-fast: false
       matrix:
-        java: [8, 11]
         os: [ubuntu-latest, windows-latest]
         python: ["3.8", "3.9", "3.10"]
+        toolchain: [stable]
     steps:
       - uses: actions/checkout@v2
-      - name: Cache local Maven repository
-        uses: actions/cache@v2
-        with:
-          path: ~/.m2/repository
-          key: ${{ runner.os }}-maven-v1-jdk${{ matrix.java }}-${{ hashFiles('**/pom.xml') }}
       - name: Set up Python
         uses: conda-incubator/setup-miniconda@v2
         with:
@@ -94,21 +59,21 @@ jobs:
           channel-priority: strict
           activate-environment: dask-sql
           environment-file: ${{ env.CONDA_FILE }}
-      - name: Download the pre-build jar
-        uses: actions/download-artifact@v1
+      - name: Setup Rust Toolchain
+        uses: actions-rs/toolchain@v1
+        id: rust-toolchain
         with:
-          name: jar
-          path: dask_sql/jar/
+          toolchain: stable
+          override: true
+      - name: Build the Rust DataFusion bindings
+        run: |
+          python setup.py build install
       - name: Install hive testing dependencies for Linux
         if: matrix.os == 'ubuntu-latest'
         run: |
           mamba install -c conda-forge sasl>=0.3.1
           docker pull bde2020/hive:2.3.2-postgresql-metastore
           docker pull bde2020/hive-metastore-postgresql:2.3.0
-      - name: Set proper JAVA_HOME for Windows
-        if: matrix.os == 'windows-latest'
-        run: |
-          echo "JAVA_HOME=${{ env.CONDA }}\envs\dask-sql\Library" >> $GITHUB_ENV
       - name: Optionally install upstream dev Dask / dask-ml
         if: needs.detect-ci-trigger.outputs.triggered == 'true'
         run: |
@@ -133,15 +98,10 @@ jobs:
 
   cluster:
     name: "Test in a dask cluster"
-    needs: [detect-ci-trigger, build]
+    needs: [detect-ci-trigger]
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
-      - name: Cache local Maven repository
-        uses: actions/cache@v2
-        with:
-          path: ~/.m2/repository
-          key: ${{ runner.os }}-maven-v1-jdk11-${{ hashFiles('**/pom.xml') }}
       - name: Set up Python
         uses: conda-incubator/setup-miniconda@v2
         with:
@@ -150,12 +110,16 @@ jobs:
           python-version: "3.8"
           channel-priority: strict
           activate-environment: dask-sql
-          environment-file: continuous_integration/environment-3.8-jdk11-dev.yaml
-      - name: Download the pre-build jar
-        uses: actions/download-artifact@v1
-        with:
-          name: jar
-          path: dask_sql/jar/
+          environment-file: continuous_integration/environment-3.8-dev.yaml
+      - name: Setup Rust Toolchain
+        uses: actions-rs/toolchain@v1
+        id: rust-toolchain
+        with:
+          toolchain: stable
+          override: true
+      - name: Build the Rust DataFusion bindings
+        run: |
+          python setup.py build install
       - name: Install dependencies
         run: |
           mamba install python-blosc lz4 -c conda-forge
@@ -184,29 +148,21 @@ jobs:
 
   import:
     name: "Test importing with bare requirements"
-    needs: [detect-ci-trigger, build]
+    needs: [detect-ci-trigger]
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
-      - name: Cache local Maven repository
-        uses: actions/cache@v2
-        with:
-          path: ~/.m2/repository
-          key: ${{ runner.os }}-maven-v1-jdk11-${{ hashFiles('**/pom.xml') }}
       - name: Set up Python
         uses: conda-incubator/setup-miniconda@v2
         with:
           python-version: "3.8"
           mamba-version: "*"
           channels: conda-forge,defaults
           channel-priority: strict
-      - name: Download the pre-build jar
-        uses: actions/download-artifact@v1
-        with:
-          name: jar
-          path: dask_sql/jar/
       - name: Install dependencies and nothing else
         run: |
+          conda install setuptools-rust
+          conda install pyarrow>=4.0.0
           pip install -e .
 
           which python

.github/workflows/update-gpuci.yml

@@ -13,39 +13,70 @@ jobs:
     steps:
       - uses: actions/checkout@v2
 
+      - name: Parse current axis YAML
+        uses: the-coding-turtle/[email protected]
+        with:
+          file: continuous_integration/gpuci/axis.yaml
+
       - name: Get latest cuDF nightly version
-        id: latest_version
+        id: cudf_latest
         uses: jacobtomlinson/[email protected]
         with:
           org: "rapidsai-nightly"
           package: "cudf"
           version_system: "CalVer"
 
-      - name: Strip git tags from versions
+      - name: Get latest cuML nightly version
+        id: cuml_latest
+        uses: jacobtomlinson/[email protected]
+        with:
+          org: "rapidsai-nightly"
+          package: "cuml"
+          version_system: "CalVer"
+
+      - name: Get latest UCX-Py nightly version
+        id: ucx_py_latest
+        uses: jacobtomlinson/[email protected]
+        with:
+          org: "rapidsai-nightly"
+          package: "ucx-py"
+          version_system: "CalVer"
+
+      - name: Get old RAPIDS / UCX-Py versions
         env:
-          FULL_RAPIDS_VER: ${{ steps.latest_version.outputs.version }}
-        run: echo "RAPIDS_VER=${FULL_RAPIDS_VER::-10}" >> $GITHUB_ENV
+          FULL_CUDF_VER: ${{ steps.cudf_latest.outputs.version }}
+          FULL_CUML_VER: ${{ steps.cuml_latest.outputs.version }}
+          FULL_UCX_PY_VER: ${{ steps.ucx_py_latest.outputs.version }}
+        run: |
+          echo RAPIDS_VER=$RAPIDS_VER_0 >> $GITHUB_ENV
+          echo UCX_PY_VER=$(curl -sL https://version.gpuci.io/rapids/$RAPIDS_VER_0) >> $GITHUB_ENV
+          echo NEW_CUDF_VER=${FULL_CUDF_VER::-10} >> $GITHUB_ENV
+          echo NEW_CUML_VER=${FULL_CUML_VER::-10} >> $GITHUB_ENV
+          echo NEW_UCX_PY_VER=${FULL_UCX_PY_VER::-10} >> $GITHUB_ENV
 
-      - name: Find and Replace Release
-        uses: jacobtomlinson/gha-find-replace@0.1.4
+      - name: Update RAPIDS version
+        uses: jacobtomlinson/gha-find-replace@v2
         with:
           include: 'continuous_integration\/gpuci\/axis\.yaml'
-          find: "RAPIDS_VER:\n- .*"
-          replace: |-
-            RAPIDS_VER:
-            - "${{ env.RAPIDS_VER }}"
+          find: "${{ env.RAPIDS_VER }}"
+          replace: "${{ env.NEW_CUDF_VER }}"
+          regex: false
 
       - name: Create Pull Request
         uses: peter-evans/create-pull-request@v3
+        # make sure ucx-py nightlies are available and that cuDF/cuML nightly versions match up
+        if: |
+          env.UCX_PY_VER != env.NEW_UCX_PY_VER &&
+          env.NEW_CUDF_VER == env.NEW_CUML_VER
         with:
           token: ${{ secrets.GITHUB_TOKEN }}
           draft: true
-          commit-message: "Update gpuCI `RAPIDS_VER` to `${{ env.RAPIDS_VER }}`"
-          title: "Update gpuCI `RAPIDS_VER` to `${{ env.RAPIDS_VER }}`"
+          commit-message: "Update gpuCI `RAPIDS_VER` to `${{ env.NEW_CUDF_VER }}`"
+          title: "Update gpuCI `RAPIDS_VER` to `${{ env.NEW_CUDF_VER }}`"
           team-reviewers: "dask/gpu"
           author: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
           branch: "upgrade-gpuci-rapids"
           body: |
-            A new cuDF nightly version has been detected.
+            New cuDF and ucx-py nightly versions have been detected.
 
-            Updated `axis.yaml` to use `${{ env.RAPIDS_VER }}`.
+            Updated `axis.yaml` to use `${{ env.NEW_CUDF_VER }}`.

.gitignore

@@ -60,3 +60,4 @@ dask_sql/jar
 dask-worker-space/
 node_modules/
 docs/source/_build/
+dask_planner/Cargo.lock