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>1.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/rust.yml

@@ -0,0 +1,50 @@
+name: Rust-Datafusion
+
+on:
+  push:
+  pull_request:
+
+jobs:
+  lint:
+    name: Lint
+    runs-on: ubuntu-latest
+    container:
+      image: amd64/rust
+    steps:
+      - uses: actions/checkout@v2
+      - name: Setup toolchain
+        run: |
+          rustup toolchain install stable
+          rustup default stable
+          rustup component add rustfmt
+      - name: Run
+        run: cargo fmt --manifest-path ./dask_planner/Cargo.toml --all -- --check
+
+  clippy:
+    name: Clippy
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        arch: [amd64]
+        rust: [stable]
+    container:
+      image: ${{ matrix.arch }}/rust
+      env:
+        # Disable full debug symbol generation to speed up CI build and keep memory down
+        # "1" means line tables only, which is useful for panic tracebacks.
+        RUSTFLAGS: "-C debuginfo=1"
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          submodules: true
+      - name: Setup Rust toolchain
+        run: |
+          rustup toolchain install ${{ matrix.rust }}
+          rustup default ${{ matrix.rust }}
+          rustup component add rustfmt clippy
+      - name: Run clippy
+        run: |
+          cargo clippy --manifest-path ./dask_planner/Cargo.toml --all-targets --workspace -- -D warnings
+        env:
+          CARGO_HOME: "/github/home/.cargo"
+          CARGO_TARGET_DIR: "/github/home/target"

.github/workflows/test-upstream.yml

@@ -5,53 +5,18 @@ 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 }})"
+    name: "Test upstream dev (${{ matrix.os }}, python: ${{ matrix.python }})"
     needs: build
     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:
@@ -62,7 +27,7 @@ jobs:
         uses: actions/cache@v2
         with:
           path: ~/.m2/repository
-          key: ${{ runner.os }}-maven-v1-jdk${{ matrix.java }}-${{ hashFiles('**/pom.xml') }}
+          key: ${{ runner.os }}-maven-v1-${{ hashFiles('**/pom.xml') }}
       - name: Set up Python
         uses: conda-incubator/setup-miniconda@v2
         with:

.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 (Apache Datafusion Rust bindings) on ubuntu
+# 2. testing code (using the build Datafusion bindings) on ubuntu and windows
 name: Test Python package
 on:
   push:
@@ -39,15 +39,10 @@ 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
+    name: Build the dask_planner 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:
@@ -56,35 +51,24 @@ jobs:
           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
+          environment-file: continuous_integration/environment-3.8-dev.yaml
+      - name: Build the Rust Datafusion bindings
         run: |
-          python setup.py build_ext
-      - name: Upload the jar
-        uses: actions/upload-artifact@v1
-        with:
-          name: jar
-          path: dask_sql/jar/DaskSQL.jar
+          python setup.py build
 
   test:
-    name: "Test (${{ matrix.os }}, java: ${{ matrix.java }}, python: ${{ matrix.python }})"
+    name: "Test (${{ matrix.os }}, python: ${{ matrix.python }})"
     needs: [detect-ci-trigger, build]
     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"]
     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 +78,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: Optionally install upstream dev Dask / dask-ml
         if: needs.detect-ci-trigger.outputs.triggered == 'true'
         run: |
@@ -137,11 +112,6 @@ jobs:
     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 +120,7 @@ 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: Install dependencies
         run: |
           mamba install python-blosc lz4 -c conda-forge
@@ -188,23 +153,13 @@ jobs:
     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: |
           pip install -e .

.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 }}`.

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
   - repo: https://github.com/psf/black
-    rev: 19.10b0
+    rev: 22.3.0
     hooks:
       - id: black
         language_version: python3
@@ -16,6 +16,13 @@ repos:
         args:
           - "--profile"
           - "black"
+  - repo: https://github.com/doublify/pre-commit-rust
+    rev: master
+    hooks:
+    -   id: fmt
+        args: ['--manifest-path', './dask_planner/Cargo.toml', '--verbose', '--']
+    -   id: cargo-check
+        args: ['--manifest-path', './dask_planner/Cargo.toml', '--verbose', '--']
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v3.2.0
     hooks:

MANIFEST.in

@@ -1,4 +1,5 @@
 recursive-include dask_sql *.yaml
+recursive-include dask_planner *
 
 include versioneer.py
 include dask_sql/_version.py