diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index c24be384421..eb1b5ed20d3 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -16,30 +16,36 @@ # request a review from owners on PRs with changes to matching files. # We currently do not enforce these reviews as required so it's only a tool # for more granular notifications at the moment. For example component maintainers -# can set a rule so that they are pinged on changes to the sections of the +# can set a rule so that they are pinged on changes to the sections of the # codebase that are relevant for their component. # Only users that have write access to the repo can be added as owners. # See the official docs for more details on syntax and precedence of rules: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners#example-of-a-codeowners-file +# Velox-cuDF codeowners +* @rapidsai/velox-cudf-codeowners + # Build & CI -CMake/ @assignUser @majetideepak -*.cmake @assignUser @majetideepak -**/CMakeLists.txt @assignUser @majetideepak -scripts/ @assignUser @majetideepak -.github/ @assignUser @majetideepak +#CMake/ @assignUser @majetideepak +#*.cmake @assignUser @majetideepak +#**/CMakeLists.txt @assignUser @majetideepak +#scripts/ @assignUser @majetideepak +#.github/ @assignUser @majetideepak + +# Breeze +#velox/experimental/breeze @dreveman -# Breeze -velox/experimental/breeze @dreveman +# cuDF +velox/experimental/cudf @bdice @karthikeyann @devavret # Parquet -velox/dwio/parquet/ @majetideepak +#velox/dwio/parquet/ @majetideepak # Storage Adapters -velox/connectors/hive/storage_adapters/ @majetideepak +#velox/connectors/hive/storage_adapters/ @majetideepak # Connectors -velox/connectors/ @majetideepak +#velox/connectors/ @majetideepak # Caching -velox/common/caching/ @majetideepak +#velox/common/caching/ @majetideepak diff --git a/.github/copy-pr-bot.yaml b/.github/copy-pr-bot.yaml new file mode 100644 index 00000000000..895ba83ee54 --- /dev/null +++ b/.github/copy-pr-bot.yaml @@ -0,0 +1,4 @@ +# Configuration file for `copy-pr-bot` GitHub App +# https://docs.gha-runners.nvidia.com/apps/copy-pr-bot/ + +enabled: true diff --git a/.github/workflows/benchmark.yml b/.github/disabled-workflows/benchmark.yml similarity index 99% rename from .github/workflows/benchmark.yml rename to .github/disabled-workflows/benchmark.yml index 53d2b202700..317f83a1e25 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/disabled-workflows/benchmark.yml @@ -38,7 +38,7 @@ defaults: jobs: benchmark: if: github.repository == 'facebookincubator/velox' - runs-on: 8-core-ubuntu-22.04 + runs-on: linux-amd64-cpu8 env: CCACHE_DIR: "${{ github.workspace }}/ccache/" CCACHE_BASEDIR: "${{ github.workspace }}" diff --git a/.github/workflows/breeze.yml b/.github/disabled-workflows/breeze.yml similarity index 100% rename from .github/workflows/breeze.yml rename to .github/disabled-workflows/breeze.yml diff --git a/.github/workflows/build-metrics.yml b/.github/disabled-workflows/build-metrics.yml similarity index 99% rename from .github/workflows/build-metrics.yml rename to .github/disabled-workflows/build-metrics.yml index 3677fd20ad9..f234b7a6414 100644 --- a/.github/workflows/build-metrics.yml +++ b/.github/disabled-workflows/build-metrics.yml @@ -42,7 +42,7 @@ jobs: strategy: fail-fast: false matrix: - runner: ["16-core-ubuntu"] + runner: ["linux-amd64-cpu8"] type: ["debug", "release"] link-type: ["shared", "static"] defaults: diff --git a/.github/workflows/build_pyvelox.yml b/.github/disabled-workflows/build_pyvelox.yml similarity index 100% rename from .github/workflows/build_pyvelox.yml rename to .github/disabled-workflows/build_pyvelox.yml diff --git a/.github/workflows/docker.yml b/.github/disabled-workflows/docker.yml similarity index 100% rename from .github/workflows/docker.yml rename to .github/disabled-workflows/docker.yml diff --git a/.github/workflows/docs.yml b/.github/disabled-workflows/docs.yml similarity index 100% rename from .github/workflows/docs.yml rename to .github/disabled-workflows/docs.yml diff --git a/.github/workflows/experimental.yml b/.github/disabled-workflows/experimental.yml similarity index 99% rename from .github/workflows/experimental.yml rename to .github/disabled-workflows/experimental.yml index 73cf16c05a5..58f76bf292e 100644 --- a/.github/workflows/experimental.yml +++ b/.github/disabled-workflows/experimental.yml @@ -47,7 +47,7 @@ permissions: jobs: compile: - runs-on: 16-core-ubuntu + runs-on: linux-amd64-cpu8 timeout-minutes: 120 env: CCACHE_DIR: "${{ github.workspace }}/ccache/" @@ -110,7 +110,7 @@ jobs: path: velox/_build/debug/velox/expression/fuzzer/spark_expression_fuzzer_test presto-java-aggregation-fuzzer-run: - runs-on: 16-core-ubuntu + runs-on: linux-amd64-cpu8 container: ghcr.io/facebookincubator/velox-dev:presto-java timeout-minutes: 120 env: diff --git a/.github/workflows/linux-build-base.yml b/.github/disabled-workflows/linux-build-base.yml similarity index 100% rename from .github/workflows/linux-build-base.yml rename to .github/disabled-workflows/linux-build-base.yml diff --git a/.github/workflows/macos.yml b/.github/disabled-workflows/macos.yml similarity index 100% rename from .github/workflows/macos.yml rename to .github/disabled-workflows/macos.yml diff --git a/.github/workflows/scheduled.yml b/.github/disabled-workflows/scheduled.yml similarity index 99% rename from .github/workflows/scheduled.yml rename to .github/disabled-workflows/scheduled.yml index de383c1ad52..f60fa9fa057 100644 --- a/.github/workflows/scheduled.yml +++ b/.github/disabled-workflows/scheduled.yml @@ -93,7 +93,7 @@ jobs: name: Build # prevent errors when forks ff their main branch if: ${{ github.repository == 'facebookincubator/velox' }} - runs-on: 32-core-ubuntu + runs-on: linux-amd64-cpu8 container: ghcr.io/facebookincubator/velox-dev:centos9 timeout-minutes: 120 env: diff --git a/.github/workflows/linux-build.yml b/.github/workflows/linux-build.yml index bca6e5f9498..3b723bdc599 100644 --- a/.github/workflows/linux-build.yml +++ b/.github/workflows/linux-build.yml @@ -17,33 +17,8 @@ name: Linux Build using GCC on: push: branches: - - "main" - paths: - - "velox/**" - - "!velox/docs/**" - - "CMakeLists.txt" - - "CMake/**" - - "third_party/**" - - "scripts/setup-ubuntu.sh" - - "scripts/setup-common.sh" - - "scripts/setup-versions.sh" - - "scripts/setup-helper-functions.sh" - - ".github/workflows/linux-build.yml" - - ".github/workflows/linux-build-base.yml" - - pull_request: - paths: - - "velox/**" - - "!velox/docs/**" - - "CMakeLists.txt" - - "CMake/**" - - "third_party/**" - - "scripts/setup-common.sh" - - "scripts/setup-versions.sh" - - "scripts/setup-ubuntu.sh" - - "scripts/setup-helper-functions.sh" - - ".github/workflows/linux-build.yml" - - ".github/workflows/linux-build-base.yml" + - "velox-cudf" + - "pull-request/[0-9]+" permissions: contents: read @@ -53,6 +28,144 @@ concurrency: cancel-in-progress: true jobs: - linux-gcc: - name: Build with GCC - uses: ./.github/workflows/linux-build-base.yml + adapters: + name: Linux release with adapters + runs-on: linux-amd64-cpu16 + container: ghcr.io/facebookincubator/velox-dev:adapters + defaults: + run: + shell: bash + env: + CCACHE_DIR: "${{ github.workspace }}/ccache" + VELOX_DEPENDENCY_SOURCE: SYSTEM + GTest_SOURCE: BUNDLED + simdjson_SOURCE: BUNDLED + xsimd_SOURCE: BUNDLED + Arrow_SOURCE: BUNDLED + Thrift_SOURCE: BUNDLED + cudf_SOURCE: BUNDLED + CUDA_VERSION: "12.8" + steps: + - uses: actions/checkout@v4 + + - name: Fix git permissions + # Usually actions/checkout does this but as we run in a container + # it doesn't work + run: git config --global --add safe.directory ${GITHUB_WORKSPACE} + + - name: Install Dependencies + run: | + # Allows to install arbitrary cuda-version whithout needing to update + # docker container before. It simplifies testing new/different versions + if ! yum list installed cuda-nvcc-$(echo ${CUDA_VERSION} | tr '.' '-') 1>/dev/null; then + source scripts/setup-centos9.sh + install_cuda ${CUDA_VERSION} + fi + + # TODO: Install a newer cmake here until we update the images upstream + pip install cmake==3.30.4 + + - uses: assignUser/stash/restore@v1 + with: + path: '${{ env.CCACHE_DIR }}' + key: ccache-linux-adapters + + - name: "Zero Ccache Statistics" + run: | + ccache -sz + + - name: Make Release Build + env: + MAKEFLAGS: 'TREAT_WARNINGS_AS_ERRORS=0 NUM_THREADS=16 MAX_HIGH_MEM_JOBS=4' + CUDA_ARCHITECTURES: 70 + CUDA_COMPILER: /usr/local/cuda-${CUDA_VERSION}/bin/nvcc + # Set compiler to GCC 12 + CUDA_FLAGS: "-ccbin /opt/rh/gcc-toolset-12/root/usr/bin" + run: | + EXTRA_CMAKE_FLAGS=( + "-DVELOX_ENABLE_BENCHMARKS=ON" + "-DVELOX_ENABLE_ARROW=ON" + "-DVELOX_ENABLE_PARQUET=ON" + "-DVELOX_ENABLE_HDFS=ON" + "-DVELOX_ENABLE_S3=ON" + "-DVELOX_ENABLE_GCS=ON" + "-DVELOX_ENABLE_ABFS=ON" + "-DVELOX_ENABLE_REMOTE_FUNCTIONS=ON" + "-DVELOX_ENABLE_GPU=ON" + "-DVELOX_ENABLE_CUDF=ON" + "-DVELOX_MONO_LIBRARY=ON" + ) + make release EXTRA_CMAKE_FLAGS="${EXTRA_CMAKE_FLAGS[*]}" + + - name: Ccache after + run: ccache -s + + - uses: assignUser/stash/save@v1 + with: + path: '${{ env.CCACHE_DIR }}' + key: ccache-linux-adapters + + - name: Run Tests + # Some of the adapters dependencies are in the 'adapters' conda env + shell: mamba run --no-capture-output -n adapters /usr/bin/bash -e {0} + env: + LIBHDFS3_CONF: "${{ github.workspace }}/scripts/hdfs-client.xml" + working-directory: _build/release + run: | + export CLASSPATH=`/usr/local/hadoop/bin/hdfs classpath --glob` + ctest -j 8 --label-exclude cuda_driver --output-on-failure --no-tests=error -E "velox_exec_test|velox_hdfs_file_test|velox_s3" + +# ubuntu-debug: +# runs-on: linux-amd64-cpu16 +# name: "Ubuntu debug with resolve_dependency" +# env: +# CCACHE_DIR: "${{ github.workspace }}/ccache" +# defaults: +# run: +# shell: bash +# working-directory: velox +# steps: + +# - name: Get Ccache Stash +# uses: assignUser/stash/restore@v1 +# with: +# path: '${{ env.CCACHE_DIR }}' +# key: ccache-ubuntu-debug-default + +# - name: Ensure Stash Dirs Exists +# working-directory: ${{ github.workspace }} +# run: | +# mkdir -p '${{ env.CCACHE_DIR }}' + +# - uses: actions/checkout@v4 +# with: +# path: velox + +# - name: Install Dependencies +# run: | +# source scripts/setup-ubuntu.sh && install_apt_deps + +# - name: Clear CCache Statistics +# run: | +# ccache -sz + +# - name: Make Debug Build +# env: +# VELOX_DEPENDENCY_SOURCE: BUNDLED +# MAKEFLAGS: "TREAT_WARNINGS_AS_ERRORS=0 NUM_THREADS=16 MAX_HIGH_MEM_JOBS=4" +# EXTRA_CMAKE_FLAGS: "-DVELOX_ENABLE_ARROW=ON -DVELOX_ENABLE_PARQUET=ON" +# run: | +# make debug + +# - name: CCache after +# run: | +# ccache -vs + +# - uses: assignUser/stash/save@v1 +# with: +# path: '${{ env.CCACHE_DIR }}' +# key: ccache-ubuntu-debug-default + +# - name: Run Tests +# run: | +# cd _build/debug && ctest -j 8 --output-on-failure --no-tests=error -E "velox_exec_test" diff --git a/.github/workflows/preliminary_checks.yml b/.github/workflows/preliminary_checks.yml index 3e63212b5ec..e725ac1ca50 100644 --- a/.github/workflows/preliminary_checks.yml +++ b/.github/workflows/preliminary_checks.yml @@ -14,12 +14,9 @@ name: Run Checks on: - pull_request: - types: - - opened - - reopened - - edited - - synchronize + push: + branches: + - "pull-request/[0-9]+" permissions: contents: read @@ -37,10 +34,10 @@ jobs: fail-fast: false matrix: config: - - { name: "License Header", - command: "header-fix", - message: "Found missing License Header(s)", - } +# - { name: "License Header", +# command: "header-fix", +# message: "Found missing License Header(s)", +# } - { name: "Code Format", command: "format-fix", message: "Found format issues" @@ -76,22 +73,22 @@ jobs: exit 1 fi - title-check: - name: PR Title Format - runs-on: ubuntu-latest - steps: - - shell: python - env: - title: "${{ github.event.pull_request.title }}" - run: | - import re - import os - title = os.environ["title"] - title_re = r"^(feat|fix|build|test|docs|refactor|misc)(\(.+\))?!?: ([A-Z].+)[^.]$" - match = re.search(title_re, title) +# title-check: +# name: PR Title Format +# runs-on: ubuntu-latest +# steps: +# - shell: python +# env: +# title: "${{ github.event.pull_request.title }}" +# run: | +# import re +# import os +# title = os.environ["title"] +# title_re = r"^(feat|fix|build|test|docs|refactor|misc)(\(.+\))?!?: ([A-Z].+)[^.]$" +# match = re.search(title_re, title) - if match is None: - print("::error::Please follow conventional commit guidelines in commit titles as described in CONTRIBUTING.md: https://github.com/facebookincubator/velox/blob/main/CONTRIBUTING.md#commit-messages") - exit(1) - else: - exit(0) +# if match is None: +# print("::error::Please follow conventional commit guidelines in commit titles as described in CONTRIBUTING.md: https://github.com/facebookincubator/velox/blob/main/CONTRIBUTING.md#commit-messages") +# exit(1) +# else: +# exit(0) diff --git a/.gitignore b/.gitignore index 3ebb2942c5c..44273a09c1e 100644 --- a/.gitignore +++ b/.gitignore @@ -81,6 +81,7 @@ m4/lt~obsolete.m4 #m4/ build/ _build/ +.cache/ .ccache/ ccache/ #*.m4 @@ -310,7 +311,7 @@ third_party/imdb/data .last_format # Benchmarks .last_benchmarked_commit -benchmark_results/ +benchmark_results* duckdb_unittest_tempdir/ grammar.y.tmp src/amalgamation/ @@ -327,3 +328,9 @@ velox/docs/bindings/python/_generate/* scripts/bm-report/report.html dist/ wheelhouse/ + +# Custom ignores +aws-sdk-cpp +cufile.log +velox-tpch-*-data +xsimd diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 00000000000..6bb8dcebe16 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,20 @@ +repos: + - repo: local + hooks: + - id: check.py + name: check.py + entry: scripts/check.py format main --fix + language: python + # Note that pre-commit autoupdate does not update the versions + # of dependencies, so we'll have to update this manually. + additional_dependencies: + - clang-format==18.* + - cmakelang==0.6.13 + - pyyaml + - regex + pass_filenames: false + verbose: true + require_serial: true + +default_language_version: + python: python3 diff --git a/CMake/resolve_dependency_modules/README.md b/CMake/resolve_dependency_modules/README.md index 268617cbabe..201345e3e27 100644 --- a/CMake/resolve_dependency_modules/README.md +++ b/CMake/resolve_dependency_modules/README.md @@ -15,6 +15,7 @@ by Velox. See details on bundling below. | glog | default | Yes | | gtest (testing) | default | Yes | | libevent | default | No | +| libcudf | default | Yes | | libsodium | default | No | | lz4 | default | No | | snappy | default | No | diff --git a/CMake/resolve_dependency_modules/cudf.cmake b/CMake/resolve_dependency_modules/cudf.cmake index 56f7b9bf1d1..81a88abaade 100644 --- a/CMake/resolve_dependency_modules/cudf.cmake +++ b/CMake/resolve_dependency_modules/cudf.cmake @@ -17,35 +17,35 @@ include_guard(GLOBAL) # 3.30.4 is the minimum version required by cudf cmake_minimum_required(VERSION 3.30.4) -set(VELOX_rapids_cmake_VERSION 25.04) +set(VELOX_rapids_cmake_VERSION 25.06) set(VELOX_rapids_cmake_BUILD_SHA256_CHECKSUM - 458c14eaff9000067b32d65c8c914f4521090ede7690e16eb57035ce731386db) + 812cef3478ef2ee02742d8cba68ab754603eb21e1333f1be03b91e7820ca0e27) set(VELOX_rapids_cmake_SOURCE_URL - "https://github.com/rapidsai/rapids-cmake/archive/7828fc8ff2e9f4fa86099f3c844505c2f47ac672.tar.gz" + "https://github.com/rapidsai/rapids-cmake/archive/4f203ce0126c91ff57289dfc70555f29cd81a8ee.tar.gz" ) velox_resolve_dependency_url(rapids_cmake) -set(VELOX_rmm_VERSION 25.04) +set(VELOX_rmm_VERSION 25.06) set(VELOX_rmm_BUILD_SHA256_CHECKSUM - 294905094213a2d1fd8e024500359ff871bc52f913a3fbaca3514727c49f62de) + d2cabadc6c484620a7aaff533494920f4e8c8c027ae717cae83110883e3cb378) set(VELOX_rmm_SOURCE_URL - "https://github.com/rapidsai/rmm/archive/d8b7dacdeda302d2e37313c02d14ef5e1d1e98ea.tar.gz" + "https://github.com/rapidsai/rmm/archive/c7a331432f003bcaa8cc45d7bbcbd21a2876565f.tar.gz" ) velox_resolve_dependency_url(rmm) -set(VELOX_kvikio_VERSION 25.04) +set(VELOX_kvikio_VERSION 25.06) set(VELOX_kvikio_BUILD_SHA256_CHECKSUM - 4a0b15295d0a397433930bf9a309e4ad2361b25dc7a7b3e6a35d0c9419d0cb62) + bdf756557ea6608ec5a00cde5130ec22b6059f96208c269a83c49917967804f8) set(VELOX_kvikio_SOURCE_URL - "https://github.com/rapidsai/kvikio/archive/5c710f37236bda76e447e929e17b1efbc6c632c3.tar.gz" + "https://github.com/rapidsai/kvikio/archive/9f143867a41c56d6df4c58572311c9bee004285a.tar.gz" ) velox_resolve_dependency_url(kvikio) -set(VELOX_cudf_VERSION 25.04) +set(VELOX_cudf_VERSION 25.06) set(VELOX_cudf_BUILD_SHA256_CHECKSUM - e5a1900dfaf23dab2c5808afa17a2d04fa867d2892ecec1cb37908f3b73715c2) + 39e1b32c8491bb84f52388a3ae9cd28b1e82ce81f311e6678067741f9d0ac01d) set(VELOX_cudf_SOURCE_URL - "https://github.com/rapidsai/cudf/archive/4c1c99011da2c23856244e05adda78ba66697105.tar.gz" + "https://github.com/rapidsai/cudf/archive/191620472f3d3daeadf32003c37ee99eaa4773a9.tar.gz" ) velox_resolve_dependency_url(cudf) diff --git a/CMakeLists.txt b/CMakeLists.txt index 999ad376645..ab7be3d0241 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -388,6 +388,8 @@ if(ENABLE_ALL_WARNINGS) -Wno-unused-parameter \ -Wno-sign-compare \ -Wno-ignored-qualifiers \ + -Wno-missing-field-initializers \ + -Wno-deprecated-copy \ ${KNOWN_COMPILER_SPECIFIC_WARNINGS}") set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wextra ${KNOWN_WARNINGS}") diff --git a/README.md b/README.md index 16fd3cd265c..5633efdeb77 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,23 @@ +> [!IMPORTANT] +> # Experimental RAPIDS cuDF Backend for Velox +> This repository is a fork of +> [Velox](https://github.com/facebookincubator/velox) with support for [RAPIDS +> cuDF](https://github.com/rapidsai/cudf/) as a GPU-accelerated Velox backend. +> We are working to upstream the contents of this repository to the Velox +> public repository. That effort is +> tracked in [Velox issue +> #12412](https://github.com/facebookincubator/velox/issues/12412), with a +> description of the high level design and merge plan for upstreaming this +> work. +> +> ### Quickstart +> This repository contains scripts for quickly building and running Velox with +> the cuDF backend. To launch the CUDA container: +> ``` +> docker-compose run -e NUM_THREADS=$(nproc) --rm adapters-cuda /bin/bash +> ``` +> Then invoke `./build.sh` to build Velox with cuDF support and run tests. + Velox logo Velox is a composable execution engine distributed as an open source C++ diff --git a/benchmark.sh b/benchmark.sh new file mode 100755 index 00000000000..5ed3d87360d --- /dev/null +++ b/benchmark.sh @@ -0,0 +1,81 @@ +#!/bin/bash +# Copyright (c) Facebook, Inc. and its affiliates. +# +# Licensed 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. + +set -euo pipefail + +# To get the data, copy from /datasets/velox-tpch-sf10-data to this repo: +# cp -r /datasets/velox-tpch-sf10-data . + +# Run this to launch the CUDA container: +# docker-compose run -e NUM_THREADS=$(nproc) --rm adapters-cuda /bin/bash +# Then invoke ./build.sh to build with GPU support and run tests. + +# Run a GPU build and test +pushd "$(dirname ${0})" + +mkdir -p benchmark_results + +queries=${1:-$(seq 1 22)} +devices=${2:-"cpu gpu"} +profile=${3:-"false"} + +num_drivers=${NUM_DRIVERS:-4} +output_batch_rows=${BATCH_SIZE_ROWS:-100000} +cudf_chunk_read_limit=$((1024 * 1024 * 1024 * 1)) +cudf_pass_read_limit=0 +VELOX_CUDF_MEMORY_RESOURCE="async" + +for query_number in ${queries}; do + printf -v query_number '%02d' "${query_number}" + for device in ${devices}; do + case "${device}" in + "cpu") + num_drivers=${NUM_DRIVERS:-32} + VELOX_CUDF_ENABLED=false;; + "gpu") + VELOX_CUDF_ENABLED=true;; + esac + echo "Running query ${query_number} on ${device} with ${num_drivers} drivers." + # The benchmarks segfault after reporting results, so we disable errors + PROFILE_CMD="" + if [[ "${profile}" == "true" ]]; then + PROFILE_CMD="nsys profile -t nvtx,cuda,osrt -f true --cuda-memory-usage=true --cuda-um-cpu-page-faults=true --cuda-um-gpu-page-faults=true --output=benchmark_results/q${query_number}_${device}_${num_drivers}_drivers.nsys-rep" + # Enable GPU metrics if supported (Ampere or newer) + if [[ "$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader -i 0 | cut -d '.' -f 1)" -gt 7 ]]; then + device_id=${CUDA_VISIBLE_DEVICES:-"0"} + PROFILE_CMD="${PROFILE_CMD} --gpu-metrics-devices=${device_id}" + fi + fi + + set +e -x + ${PROFILE_CMD} \ + ./_build/release/velox/benchmarks/tpch/velox_tpch_benchmark \ + --data_path=velox-tpch-sf100-data \ + --data_format=parquet \ + --run_query_verbose=${query_number} \ + --num_repeats=1 \ + --velox_cudf_enabled=${VELOX_CUDF_ENABLED} \ + --velox_cudf_memory_resource=${VELOX_CUDF_MEMORY_RESOURCE} \ + --num_drivers=${num_drivers} \ + --preferred_output_batch_rows=${output_batch_rows} \ + --max_output_batch-rows=${output_batch_rows} 2>&1 \ + --cudf_chunk_read_limit=${cudf_chunk_read_limit} \ + --cudf_pass_read_limit=${cudf_pass_read_limit} \ + | tee benchmark_results/q${query_number}_${device}_${num_drivers}_drivers + { set -e +x; } &> /dev/null + done +done + +popd diff --git a/build.sh b/build.sh new file mode 100755 index 00000000000..12ca0786f3a --- /dev/null +++ b/build.sh @@ -0,0 +1,31 @@ +#!/bin/bash +# Copyright (c) Facebook, Inc. and its affiliates. +# +# Licensed 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. + +set -euo pipefail + +# Run this to launch the CUDA container: +# docker-compose run -e NUM_THREADS=$(nproc) --rm adapters-cuda /bin/bash +# Then invoke ./build.sh to build with GPU support and run tests. + +# Run a GPU build and test +pushd "$(dirname ${0})" + +CUDA_ARCHITECTURES="native" EXTRA_CMAKE_FLAGS="-DVELOX_ENABLE_ARROW=ON -DVELOX_ENABLE_PARQUET=ON -DVELOX_ENABLE_BENCHMARKS=ON -DVELOX_ENABLE_BENCHMARKS_BASIC=ON" make cudf + +cd _build/release + +ctest -R cudf -V + +popd diff --git a/docker-compose.yml b/docker-compose.yml index 7dfc8752d24..15f4bbc1aa6 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -29,6 +29,7 @@ services: NUM_THREADS: 8 # default value for NUM_THREADS VELOX_DEPENDENCY_SOURCE: BUNDLED # Build dependencies from source CCACHE_DIR: "/velox/.ccache" + CMAKE_EXPORT_COMPILE_COMMANDS: 1 volumes: - .:/velox:delegated command: scripts/docker/docker-command.sh diff --git a/install-aws-sdk.sh b/install-aws-sdk.sh new file mode 100755 index 00000000000..9b9c2e9bbb4 --- /dev/null +++ b/install-aws-sdk.sh @@ -0,0 +1,35 @@ +#!/bin/bash +# Copyright (c) Facebook, Inc. and its affiliates. +# +# Licensed 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. + +if [ ! -d "aws-sdk-cpp" ]; then + git clone https://github.com/aws/aws-sdk-cpp --recurse-submodules +fi +cd aws-sdk-cpp + +mkdir -p build +cd build + +cmake ../ \ + -DCMAKE_BUILD_TYPE=Debug \ + -DCMAKE_PREFIX_PATH=/usr/local \ + -DCMAKE_INSTALL_PREFIX=/usr/local \ + -DBUILD_ONLY="s3;sts;cognito-identity;identity-management" \ + -DENABLE_TESTING=OFF +cmake --build . --config=Debug +cmake --install . --config=Debug + +#cmake ../ -DCMAKE_BUILD_TYPE=Debug -DCMAKE_PREFIX_PATH=/usr/local -DCMAKE_INSTALL_PREFIX=/usr/local +#make +#sudo make install diff --git a/install-xsimd.sh b/install-xsimd.sh new file mode 100755 index 00000000000..4d227ebc237 --- /dev/null +++ b/install-xsimd.sh @@ -0,0 +1,26 @@ +#!/bin/bash +# Copyright (c) Facebook, Inc. and its affiliates. +# +# Licensed 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. + +if [ ! -d "xsimd" ]; then + git clone https://github.com/xtensor-stack/xsimd --recurse-submodules +fi +cd xsimd + +mkdir -p build +cd build + +cmake ../ +cmake --build . --config=Debug +cmake --install . --config=Debug diff --git a/scripts/docker/adapters.dockerfile b/scripts/docker/adapters.dockerfile index 47ce766b95b..f1575206591 100644 --- a/scripts/docker/adapters.dockerfile +++ b/scripts/docker/adapters.dockerfile @@ -36,7 +36,7 @@ ENV NVIDIA_VISIBLE_DEVICES=all ENV NVIDIA_DRIVER_CAPABILITIES="compute,utility" # install miniforge -RUN curl -L -o /tmp/miniforge.sh https://github.com/conda-forge/miniforge/releases/download/23.11.0-0/Mambaforge-23.11.0-0-Linux-x86_64.sh && \ +RUN curl -L -o /tmp/miniforge.sh https://github.com/conda-forge/miniforge/releases/download/23.11.0-0/Mambaforge-23.11.0-0-Linux-$(uname -m).sh && \ bash /tmp/miniforge.sh -b -p /opt/miniforge && \ rm /tmp/miniforge.sh ENV PATH=/opt/miniforge/condabin:${PATH} diff --git a/scripts/setup-centos9.sh b/scripts/setup-centos9.sh index 008f9302162..2ba94d17aff 100755 --- a/scripts/setup-centos9.sh +++ b/scripts/setup-centos9.sh @@ -85,8 +85,22 @@ function install_gflags { } function install_cuda { + dnf install -y patch + ARCH=$(uname -m) + case "$ARCH" in + x86_64) + CUDA_ARCH="x86_64" + ;; + aarch64) + CUDA_ARCH="sbsa" + ;; + *) + echo "Error: unsupported architecture $ARCH" >&2 + exit 1 + ;; + esac # See https://developer.nvidia.com/cuda-downloads - dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel9/x86_64/cuda-rhel9.repo + dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel9/${CUDA_ARCH}/cuda-rhel9.repo local dashed="$(echo $1 | tr '.' '-')" dnf install -y \ cuda-compat-$dashed \ diff --git a/scripts/setup-check.sh b/scripts/setup-check.sh index d3d6573a8ed..5c54d0fb71f 100644 --- a/scripts/setup-check.sh +++ b/scripts/setup-check.sh @@ -19,7 +19,7 @@ set -x export DEBIAN_FRONTEND=noninteractive apt update apt install --no-install-recommends -y clang-format-18 python3-pip git make ssh -pip3 install --break-system-packages cmake==3.28.3 cmake_format black pyyaml regex +pip3 install --break-system-packages cmake==3.30.4 cmake_format black pyyaml regex pip3 cache purge apt purge --auto-remove -y python3-pip update-alternatives --install /usr/bin/clang-format clang-format "$(command -v clang-format-18)" 18 diff --git a/scripts/setup-ubuntu.sh b/scripts/setup-ubuntu.sh index 33711884bc8..3ef490cfe83 100755 --- a/scripts/setup-ubuntu.sh +++ b/scripts/setup-ubuntu.sh @@ -88,7 +88,7 @@ function install_build_prerequisites { fi source ${PYTHON_VENV}/bin/activate; # Install to /usr/local to make it available to all users. - ${SUDO} pip3 install cmake==3.28.3 + ${SUDO} pip3 install cmake==3.30.4 install_gcc11_if_needed @@ -155,9 +155,22 @@ function install_conda { } function install_cuda { - # See https://developer.nvidia.com/cuda-downloads if ! dpkg -l cuda-keyring 1>/dev/null; then - wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb + ARCH=$(uname -m) + case "$ARCH" in + x86_64) + CUDA_ARCH="x86_64" + ;; + aarch64) + CUDA_ARCH="sbsa" + ;; + *) + echo "Error: unsupported architecture $ARCH" >&2 + exit 1 + ;; + esac + # See https://developer.nvidia.com/cuda-downloads + wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/${CUDA_ARCH}/cuda-keyring_1.1-1_all.deb $SUDO dpkg -i cuda-keyring_1.1-1_all.deb rm cuda-keyring_1.1-1_all.deb $SUDO apt update diff --git a/scripts/velox_env_linux.yml b/scripts/velox_env_linux.yml index 59ceeb0adb4..0fa15f13e87 100644 --- a/scripts/velox_env_linux.yml +++ b/scripts/velox_env_linux.yml @@ -26,7 +26,7 @@ dependencies: - binutils - bison - clangxx=14 - - cmake=3.28.3 + - cmake=3.30.4 - ccache - flex - gxx=12 # has to be installed to get clang to work... diff --git a/velox/benchmarks/CMakeLists.txt b/velox/benchmarks/CMakeLists.txt index 45466b1ed3f..cb235925115 100644 --- a/velox/benchmarks/CMakeLists.txt +++ b/velox/benchmarks/CMakeLists.txt @@ -60,6 +60,8 @@ target_link_libraries( velox_type_fbhive velox_caching velox_vector_test_lib + velox_cudf_exec + velox_cudf_exec_test_lib Folly::folly Folly::follybenchmark fmt::fmt) diff --git a/velox/benchmarks/QueryBenchmarkBase.cpp b/velox/benchmarks/QueryBenchmarkBase.cpp index 469b6c00542..12cc886b447 100644 --- a/velox/benchmarks/QueryBenchmarkBase.cpp +++ b/velox/benchmarks/QueryBenchmarkBase.cpp @@ -16,6 +16,11 @@ #include "velox/benchmarks/QueryBenchmarkBase.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetConfig.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetConnector.h" +#include "velox/experimental/cudf/exec/CudfConversion.h" +#include "velox/experimental/cudf/tests/utils/ParquetConnectorTestBase.h" + DEFINE_string(data_format, "parquet", "Data format"); DEFINE_validator( @@ -89,7 +94,39 @@ DEFINE_int32( "prefetch. 1 means prefetch the next row group before decoding " "the current one"); -DEFINE_int32(split_preload_per_driver, 2, "Prefetch split metadata"); +DEFINE_uint64( + cudf_chunk_read_limit, + 0, + "Output table chunk read limit for cudf::parquet_chunked_reader."); + +DEFINE_uint64( + cudf_pass_read_limit, + 0, + "Pass read limit for cudf::parquet_chunked_reader."); + +DEFINE_int32( + cudf_gpu_batch_size_rows, + 100000, + "Preferred output batch size in rows for cudf operators."); + +DEFINE_int32(split_preload_per_driver, 1, "Prefetch split metadata"); + +DEFINE_int64( + preferred_output_batch_bytes, + 10 << 20, + "Preferred output batch size in bytes"); + +DEFINE_uint64( + max_partial_aggregation_memory, + 10 << 20, + "Maximum memory usage for partial aggregation"); + +DEFINE_int32( + preferred_output_batch_rows, + 1024, + "Preferred output batch size in rows"); + +DEFINE_int32(max_output_batch_rows, 10'000, "Max output batch size in rows"); using namespace facebook::velox::exec; using namespace facebook::velox::exec::test; @@ -195,6 +232,38 @@ void QueryBenchmarkBase::initialize() { connector::registerConnector(hiveConnector); parquet::registerParquetReaderFactory(); dwrf::registerDwrfReaderFactory(); + + // Add new values into the parquet configuration... + auto parquetConfigurationValues = + std::unordered_map(); + parquetConfigurationValues + [cudf_velox::connector::parquet::ParquetConfig::kMaxChunkReadLimit] = + std::to_string(FLAGS_cudf_chunk_read_limit); + parquetConfigurationValues + [cudf_velox::connector::parquet::ParquetConfig::kMaxPassReadLimit] = + std::to_string(FLAGS_cudf_pass_read_limit); + parquetConfigurationValues[cudf_velox::connector::parquet::ParquetConfig:: + kAllowMismatchedParquetSchemas] = + std::to_string(true); + auto parquetProperties = std::make_shared( + std::move(parquetConfigurationValues)); + + // Create parquet connector with config... + connector::registerConnectorFactory( + std::make_shared< + cudf_velox::connector::parquet::ParquetConnectorFactory>()); + auto parquetConnector = + connector::getConnectorFactory( + cudf_velox::connector::parquet::ParquetConnectorFactory:: + kParquetConnectorName) + ->newConnector( + cudf_velox::exec::test::kParquetConnectorId, + parquetProperties, + ioExecutor_.get()); + connector::registerConnector(parquetConnector); + + // Enable cuDF operators + cudf_velox::registerCudf(); } std::vector> @@ -211,7 +280,27 @@ QueryBenchmarkBase::listSplits( return result; } +std::vector> +QueryBenchmarkBase::listCudfSplits( + const std::string& path, + int32_t /*numSplitsPerFile*/, + const exec::test::TpchPlan& plan) { + std::vector> result; + auto temp = cudf_velox::exec::test::ParquetConnectorTestBase:: + makeParquetConnectorSplits(path, 1); + for (auto& i : temp) { + result.push_back(i); + } + return result; +} + void QueryBenchmarkBase::shutdown() { + cudf_velox::unregisterCudf(); + facebook::velox::connector::unregisterConnector( + cudf_velox::exec::test::kParquetConnectorId); + facebook::velox::connector::unregisterConnectorFactory( + cudf_velox::connector::parquet::ParquetConnectorFactory:: + kParquetConnectorName); if (cache_) { cache_->shutdown(); } @@ -227,6 +316,16 @@ QueryBenchmarkBase::run(const TpchPlan& tpchPlan) { params.planNode = tpchPlan.plan; params.queryConfigs[core::QueryConfig::kMaxSplitPreloadPerDriver] = std::to_string(FLAGS_split_preload_per_driver); + params.queryConfigs[core::QueryConfig::kPreferredOutputBatchBytes] = + std::to_string(FLAGS_preferred_output_batch_bytes); + params.queryConfigs[core::QueryConfig::kPreferredOutputBatchRows] = + std::to_string(FLAGS_preferred_output_batch_rows); + params.queryConfigs[core::QueryConfig::kMaxOutputBatchRows] = + std::to_string(FLAGS_max_output_batch_rows); + params.queryConfigs[cudf_velox::CudfFromVelox::kGpuBatchSizeRows] = + std::to_string(FLAGS_cudf_gpu_batch_size_rows); + params.queryConfigs[core::QueryConfig::kMaxPartialAggregationMemory] = + std::to_string(FLAGS_max_partial_aggregation_memory); const int numSplitsPerFile = FLAGS_num_splits_per_file; auto addSplits = [&](TaskCursor* taskCursor) { @@ -234,7 +333,13 @@ QueryBenchmarkBase::run(const TpchPlan& tpchPlan) { if (!taskCursor->noMoreSplits()) { for (const auto& entry : tpchPlan.dataFiles) { for (const auto& path : entry.second) { - auto splits = listSplits(path, numSplitsPerFile, tpchPlan); + auto splits = facebook::velox::cudf_velox::cudfIsRegistered() && + facebook::velox::connector::getAllConnectors().count( + cudf_velox::exec::test::kParquetConnectorId) > 0 && + facebook::velox::cudf_velox::cudfTableScanEnabled() + ? listCudfSplits( + path, 1 /* numSplitsPerFile = 1 for cudf */, tpchPlan) + : listSplits(path, numSplitsPerFile, tpchPlan); for (auto split : splits) { task->addSplit(entry.first, exec::Split(std::move(split))); } diff --git a/velox/benchmarks/QueryBenchmarkBase.h b/velox/benchmarks/QueryBenchmarkBase.h index a5a172a8216..6abad9c1be6 100644 --- a/velox/benchmarks/QueryBenchmarkBase.h +++ b/velox/benchmarks/QueryBenchmarkBase.h @@ -39,6 +39,7 @@ #include "velox/exec/Split.h" #include "velox/exec/tests/utils/HiveConnectorTestBase.h" #include "velox/exec/tests/utils/TpchQueryBuilder.h" +#include "velox/experimental/cudf/exec/ToCudf.h" #include "velox/functions/prestosql/aggregates/RegisterAggregateFunctions.h" #include "velox/functions/prestosql/registration/RegistrationFunctions.h" #include "velox/parse/TypeResolver.h" @@ -94,6 +95,11 @@ class QueryBenchmarkBase { int32_t numSplitsPerFile, const exec::test::TpchPlan& plan); + std::vector> listCudfSplits( + const std::string& path, + int32_t numSplitsPerFile, + const exec::test::TpchPlan& plan); + static void ensureTaskCompletion(exec::Task* task); static bool validateDataFormat( diff --git a/velox/benchmarks/tpch/CMakeLists.txt b/velox/benchmarks/tpch/CMakeLists.txt index 1ac7c3f1aee..15f2d16e435 100644 --- a/velox/benchmarks/tpch/CMakeLists.txt +++ b/velox/benchmarks/tpch/CMakeLists.txt @@ -18,6 +18,7 @@ target_link_libraries( velox_tpch_benchmark_lib velox_query_benchmark velox_aggregates + velox_cudf_exec velox_exec velox_exec_test_lib velox_dwio_common diff --git a/velox/exec/tests/utils/CMakeLists.txt b/velox/exec/tests/utils/CMakeLists.txt index acc1119c994..55b1966602c 100644 --- a/velox/exec/tests/utils/CMakeLists.txt +++ b/velox/exec/tests/utils/CMakeLists.txt @@ -38,6 +38,8 @@ add_library( target_link_libraries( velox_exec_test_lib + velox_cudf_exec + velox_cudf_parquet_connector velox_vector_test_lib velox_vector_fuzzer velox_temp_path diff --git a/velox/exec/tests/utils/PlanBuilder.cpp b/velox/exec/tests/utils/PlanBuilder.cpp index 1ddf46c6197..1df09ad0be1 100644 --- a/velox/exec/tests/utils/PlanBuilder.cpp +++ b/velox/exec/tests/utils/PlanBuilder.cpp @@ -32,6 +32,10 @@ #include "velox/parse/Expressions.h" #include "velox/parse/TypeResolver.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetTableHandle.h" +#include "velox/experimental/cudf/exec/ToCudf.h" +#include "velox/experimental/cudf/tests/utils/ParquetConnectorTestBase.h" + using namespace facebook::velox; using namespace facebook::velox::connector; using namespace facebook::velox::connector::hive; @@ -200,6 +204,7 @@ core::PlanNodePtr PlanBuilder::TableScanBuilder::build(core::PlanNodeId id) { const RowTypePtr& parseType = dataColumns_ ? dataColumns_ : outputType_; + std::vector subfieldExprs; core::TypedExprPtr filterNodeExpr; common::SubfieldFilters filters; filters.reserve(subfieldFilters_.size()); @@ -225,10 +230,33 @@ core::PlanNodePtr PlanBuilder::TableScanBuilder::build(core::PlanNodeId id) { "Duplicate subfield: {}", subfield.toString()); + subfieldExprs.push_back(std::move(filterExpr)); filters[std::move(subfield)] = std::move(subfieldFilter); } } + // Create AND tree of subfieldExprs as combined_subfield_filter. + // replace every 2 subfieldExpr with a single AND node, until we have a single + // node. + while (subfieldExprs.size() > 1) { + std::vector combinedSubfieldExprs; + combinedSubfieldExprs.reserve(subfieldExprs.size() / 2 + 1); + for (size_t i = 0; i < subfieldExprs.size(); i += 2) { + if (i + 1 < subfieldExprs.size()) { + auto andCallExpr = std::make_shared( + BOOLEAN(), + std::vector{ + subfieldExprs[i], subfieldExprs[i + 1]}, + "and"); + combinedSubfieldExprs.push_back(andCallExpr); + } else { + combinedSubfieldExprs.push_back(subfieldExprs[i]); + } + } + subfieldExprs = std::move(combinedSubfieldExprs); + } + core::TypedExprPtr subfieldFilterExpr = + subfieldExprs.empty() ? nullptr : subfieldExprs[0]; core::TypedExprPtr remainingFilterExpr; if (remainingFilter_) { remainingFilterExpr = core::Expressions::inferTypes( @@ -241,13 +269,28 @@ core::PlanNodePtr PlanBuilder::TableScanBuilder::build(core::PlanNodeId id) { } if (!tableHandle_) { - tableHandle_ = std::make_shared( - connectorId_, - tableName_, - true, - std::move(filters), - remainingFilterExpr, - dataColumns_); + // if cudfIsRegistered, then use cudftableScan tableHandle_ here. + if (facebook::velox::cudf_velox::cudfIsRegistered() && + facebook::velox::connector::getAllConnectors().count( + cudf_velox::exec::test::kParquetConnectorId) > 0 && + facebook::velox::cudf_velox::cudfTableScanEnabled()) { + tableHandle_ = + std::make_shared( + cudf_velox::exec::test::kParquetConnectorId, + tableName_, + subfieldFilterExpr != nullptr, + subfieldFilterExpr, + remainingFilterExpr, + dataColumns_); + } else { + tableHandle_ = std::make_shared( + connectorId_, + tableName_, + true, + std::move(filters), + remainingFilterExpr, + dataColumns_); + } } core::PlanNodePtr result = std::make_shared( id, outputType_, tableHandle_, assignments_); diff --git a/velox/experimental/cudf/.clang-tidy b/velox/experimental/cudf/.clang-tidy index ec5b4e40784..61c04b1dbd5 100644 --- a/velox/experimental/cudf/.clang-tidy +++ b/velox/experimental/cudf/.clang-tidy @@ -44,11 +44,11 @@ CheckOptions: value: CamelCase - key: readability-identifier-naming.EnumConstantPrefix value: k - + # Use nullptr instead of NULL or 0 - key: modernize-use-nullptr.NullMacros value: 'NULL' - + # Prefer enum class over enum - key: modernize-use-using.IgnoreUsingStdAllocator - value: 1 \ No newline at end of file + value: 1 diff --git a/velox/experimental/cudf/CMakeLists.txt b/velox/experimental/cudf/CMakeLists.txt index 6d400056c35..96fcdb0d557 100644 --- a/velox/experimental/cudf/CMakeLists.txt +++ b/velox/experimental/cudf/CMakeLists.txt @@ -13,6 +13,8 @@ # limitations under the License. add_subdirectory(exec) +add_subdirectory(connectors) +add_subdirectory(vector) if(VELOX_BUILD_TESTING) add_subdirectory(tests) diff --git a/velox/experimental/cudf/connectors/CMakeLists.txt b/velox/experimental/cudf/connectors/CMakeLists.txt new file mode 100644 index 00000000000..37a9408221c --- /dev/null +++ b/velox/experimental/cudf/connectors/CMakeLists.txt @@ -0,0 +1,15 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# Licensed 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. + +add_subdirectory(parquet) diff --git a/velox/experimental/cudf/connectors/parquet/CMakeLists.txt b/velox/experimental/cudf/connectors/parquet/CMakeLists.txt new file mode 100644 index 00000000000..88ae38df265 --- /dev/null +++ b/velox/experimental/cudf/connectors/parquet/CMakeLists.txt @@ -0,0 +1,44 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# Licensed 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. + +add_library(velox_cudf_parquet_config ParquetConfig.cpp) + +set_target_properties( + velox_cudf_parquet_config + PROPERTIES CUDA_ARCHITECTURES native) + +target_link_libraries( + velox_cudf_parquet_config velox_core velox_exception cudf::cudf) + +add_library( + velox_cudf_parquet_connector OBJECT + ParquetConfig.cpp + ParquetConnector.cpp + ParquetConnectorSplit.cpp + ParquetDataSource.cpp + ParquetDataSink.cpp + ParquetTableHandle.cpp) + +set_target_properties( + velox_cudf_parquet_connector + PROPERTIES CUDA_ARCHITECTURES native) + +target_link_libraries( + velox_cudf_parquet_connector + PRIVATE + cudf::cudf + velox_common_io + velox_connector + velox_type_tz + velox_gcs) diff --git a/velox/experimental/cudf/connectors/parquet/ParquetConfig.cpp b/velox/experimental/cudf/connectors/parquet/ParquetConfig.cpp new file mode 100644 index 00000000000..15449837930 --- /dev/null +++ b/velox/experimental/cudf/connectors/parquet/ParquetConfig.cpp @@ -0,0 +1,179 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#include "velox/experimental/cudf/connectors/parquet/ParquetConfig.h" + +#include "velox/common/base/Exceptions.h" +#include "velox/common/config/Config.h" + +#include + +#include + +namespace facebook::velox::cudf_velox::connector::parquet { + +int64_t ParquetConfig::skipRows() const { + return config_->get(kSkipRows, 0); +} + +std::optional ParquetConfig::numRows() const { + auto numRows = config_->get(kNumRows); + return numRows.has_value() + ? std::make_optional(numRows.value()) + : std::nullopt; +} + +std::size_t ParquetConfig::maxChunkReadLimit() const { + // chunk read limit = 0 means no limit + return config_->get(kMaxChunkReadLimit, 0); +} + +std::size_t ParquetConfig::maxChunkReadLimitSession( + const config::ConfigBase* session) const { + // pass read limit = 0 means no limit + return session->get( + kMaxChunkReadLimitSession, + config_->get(kMaxChunkReadLimit, 0)); +} + +std::size_t ParquetConfig::maxPassReadLimit() const { + // pass read limit = 0 means no limit + return config_->get(kMaxPassReadLimit, 0); +} + +std::size_t ParquetConfig::maxPassReadLimitSession( + const config::ConfigBase* session) const { + // pass read limit = 0 means no limit + return session->get( + kMaxPassReadLimitSession, + config_->get(kMaxPassReadLimit, 0)); +} + +bool ParquetConfig::isConvertStringsToCategories() const { + return config_->get(kConvertStringsToCategories, false); +} + +bool ParquetConfig::isConvertStringsToCategoriesSession( + const config::ConfigBase* session) const { + return session->get( + kConvertStringsToCategoriesSession, + config_->get(kConvertStringsToCategories, false)); +} + +bool ParquetConfig::isUsePandasMetadata() const { + return config_->get(kUsePandasMetadata, true); +} + +bool ParquetConfig::isUsePandasMetadataSession( + const config::ConfigBase* session) const { + return session->get( + kUsePandasMetadataSession, config_->get(kUsePandasMetadata, true)); +} + +bool ParquetConfig::isUseArrowSchema() const { + return config_->get(kUseArrowSchema, true); +} + +bool ParquetConfig::isUseArrowSchemaSession( + const config::ConfigBase* session) const { + return session->get( + kUseArrowSchemaSession, config_->get(kUseArrowSchema, true)); +} + +bool ParquetConfig::isAllowMismatchedParquetSchemas() const { + return config_->get(kAllowMismatchedParquetSchemas, false); +} + +bool ParquetConfig::isAllowMismatchedParquetSchemasSession( + const config::ConfigBase* session) const { + return session->get( + kAllowMismatchedParquetSchemasSession, + config_->get(kAllowMismatchedParquetSchemas, false)); +} + +cudf::data_type ParquetConfig::timestampType() const { + const auto unit = config_->get( + kTimestampType, cudf::type_id::TIMESTAMP_MILLISECONDS /*milli*/); + VELOX_CHECK( + unit == cudf::type_id::TIMESTAMP_DAYS /*days*/ || + unit == cudf::type_id::TIMESTAMP_SECONDS /*seconds*/ || + unit == cudf::type_id::TIMESTAMP_MILLISECONDS /*milli*/ || + unit == cudf::type_id::TIMESTAMP_MICROSECONDS /*micro*/ || + unit == cudf::type_id::TIMESTAMP_NANOSECONDS /*nano*/, + "Invalid timestamp unit."); + return cudf::data_type(cudf::type_id{unit}); +} + +cudf::data_type ParquetConfig::timestampTypeSession( + const config::ConfigBase* session) const { + const auto unit = session->get( + kTimestampTypeSession, + config_->get( + kTimestampType, cudf::type_id::TIMESTAMP_MILLISECONDS /*milli*/)); + VELOX_CHECK( + unit == cudf::type_id::TIMESTAMP_DAYS /*days*/ || + unit == cudf::type_id::TIMESTAMP_SECONDS /*seconds*/ || + unit == cudf::type_id::TIMESTAMP_MILLISECONDS /*milli*/ || + unit == cudf::type_id::TIMESTAMP_MICROSECONDS /*micro*/ || + unit == cudf::type_id::TIMESTAMP_NANOSECONDS /*nano*/, + "Invalid timestamp unit."); + return cudf::data_type(cudf::type_id{unit}); +} + +bool ParquetConfig::immutableFiles() const { + return config_->get(kImmutableFiles, false); +} + +uint64_t ParquetConfig::sortWriterFinishTimeSliceLimitMs( + const config::ConfigBase* session) const { + return session->get( + kSortWriterFinishTimeSliceLimitMsSession, + config_->get(kSortWriterFinishTimeSliceLimitMs, 5'000)); +} + +bool ParquetConfig::writeTimestampsAsUTC() const { + return config_->get(kWriteTimestampsAsUTC, true); +} + +bool ParquetConfig::writeTimestampsAsUTCSession( + const config::ConfigBase* session) const { + return session->get( + kWriteTimestampsAsUTCSession, + config_->get(kWriteTimestampsAsUTC, true)); +} + +bool ParquetConfig::writeArrowSchema() const { + return config_->get(kWriteArrowSchema, false); +} + +bool ParquetConfig::writeArrowSchemaSession( + const config::ConfigBase* session) const { + return session->get( + kWriteArrowSchemaSession, config_->get(kWriteArrowSchema, false)); +} + +bool ParquetConfig::writev2PageHeaders() const { + return config_->get(kWritev2PageHeaders, false); +} + +bool ParquetConfig::writev2PageHeadersSession( + const config::ConfigBase* session) const { + return session->get( + kWritev2PageHeadersSession, + config_->get(kWritev2PageHeaders, false)); +} + +} // namespace facebook::velox::cudf_velox::connector::parquet diff --git a/velox/experimental/cudf/connectors/parquet/ParquetConfig.h b/velox/experimental/cudf/connectors/parquet/ParquetConfig.h new file mode 100644 index 00000000000..bf80ad8b0ec --- /dev/null +++ b/velox/experimental/cudf/connectors/parquet/ParquetConfig.h @@ -0,0 +1,165 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#pragma once + +#include "velox/common/config/Config.h" + +#include + +#include + +namespace facebook::velox::config { +class ConfigBase; +} + +namespace facebook::velox::cudf_velox::connector::parquet { + +class ParquetConfig { + public: + // Reader config options + + // Number of rows to skip from the start; Parquet stores the number of rows as + // int64_t + static constexpr const char* kSkipRows = "parquet.reader.skip-rows"; + + // Number of rows to read; `nullopt` is all + static constexpr const char* kNumRows = "parquet.reader.num-rows"; + + // This isn't a typo; parquet connector and session config names are different + // ('-' vs '_'). + static constexpr const char* kMaxChunkReadLimit = + "parquet.reader.chunk-read-limit"; + static constexpr const char* kMaxChunkReadLimitSession = + "parquet.reader.chunk_read_limit"; + + static constexpr const char* kMaxPassReadLimit = + "parquet.reader.pass-read-limit"; + static constexpr const char* kMaxPassReadLimitSession = + "parquet.reader.pass_read_limit"; + + // Whether to store string data as categorical type + static constexpr const char* kConvertStringsToCategories = + "parquet.reader.convert-strings-to-categories"; + static constexpr const char* kConvertStringsToCategoriesSession = + "parquet.reader.convert_strings_to_categories"; + + // Whether to use PANDAS metadata to load columns + static constexpr const char* kUsePandasMetadata = + "parquet.reader.use-pandas-metadata"; + static constexpr const char* kUsePandasMetadataSession = + "parquet.reader.use_pandas_metadata"; + + // Whether to read and use ARROW schema + static constexpr const char* kUseArrowSchema = + "parquet.reader.use-arrow-schema"; + static constexpr const char* kUseArrowSchemaSession = + "parquet.reader.use_arrow_schema"; + + // Whether to allow reading matching select columns from mismatched Parquet + // files. + static constexpr const char* kAllowMismatchedParquetSchemas = + "parquet.reader.allow-mismatched-parquet-schemas"; + static constexpr const char* kAllowMismatchedParquetSchemasSession = + "parquet.reader.allow_mismatched_parquet_schemas"; + + // Cast timestamp columns to a specific type + static constexpr const char* kTimestampType = "parquet.reader.timestamp-type"; + static constexpr const char* kTimestampTypeSession = + "parquet.reader.timestamp_type"; + + // Writer config options + + /// Whether new data can be inserted into a Parquet file + /// Cudf-Velox currently does not support appending data to existing files. + static constexpr const char* kImmutableFiles = "parquet.immutable-files"; + + /// Sort Writer will exit finish() method after this many milliseconds even if + /// it has not completed its work yet. Zero means no time limit. + static constexpr const char* kSortWriterFinishTimeSliceLimitMs = + "sort-writer_finish_time_slice_limit_ms"; + static constexpr const char* kSortWriterFinishTimeSliceLimitMsSession = + "sort_writer_finish_time_slice_limit_ms"; + + static constexpr const char* kWriteTimestampsAsUTC = + "parquet.writer.write-timestamps-as-utc"; + static constexpr const char* kWriteTimestampsAsUTCSession = + "parquet.writer.write_timestamps_as_utc"; + + static constexpr const char* kWriteArrowSchema = + "parquet.writer.write-arrow-schema"; + static constexpr const char* kWriteArrowSchemaSession = + "parquet.writer.write_arrow_schema"; + + static constexpr const char* kWritev2PageHeaders = + "parquet.writer.write-v2-page-headers"; + static constexpr const char* kWritev2PageHeadersSession = + "parquet.writer.write_v2_page_headers"; + + ParquetConfig(std::shared_ptr config) { + VELOX_CHECK_NOT_NULL( + config, "Config is null for ParquetConfig initialization"); + config_ = std::move(config); + } + + const std::shared_ptr& config() const { + return config_; + } + + uint64_t sortWriterFinishTimeSliceLimitMs( + const config::ConfigBase* session) const; + + std::size_t maxChunkReadLimit() const; + std::size_t maxChunkReadLimitSession(const config::ConfigBase* session) const; + + std::size_t maxPassReadLimit() const; + std::size_t maxPassReadLimitSession(const config::ConfigBase* session) const; + + int64_t skipRows() const; + std::optional numRows() const; + + bool isConvertStringsToCategories() const; + bool isConvertStringsToCategoriesSession( + const config::ConfigBase* session) const; + + bool isUsePandasMetadata() const; + bool isUsePandasMetadataSession(const config::ConfigBase* session) const; + + bool isUseArrowSchema() const; + bool isUseArrowSchemaSession(const config::ConfigBase* session) const; + + bool isAllowMismatchedParquetSchemas() const; + bool isAllowMismatchedParquetSchemasSession( + const config::ConfigBase* session) const; + + cudf::data_type timestampType() const; + cudf::data_type timestampTypeSession(const config::ConfigBase* session) const; + + bool immutableFiles() const; + + bool writeTimestampsAsUTC() const; + bool writeTimestampsAsUTCSession(const config::ConfigBase* session) const; + + bool writeArrowSchema() const; + bool writeArrowSchemaSession(const config::ConfigBase* session) const; + + bool writev2PageHeaders() const; + bool writev2PageHeadersSession(const config::ConfigBase* session) const; + + private: + std::shared_ptr config_; +}; +} // namespace facebook::velox::cudf_velox::connector::parquet diff --git a/velox/experimental/cudf/connectors/parquet/ParquetConnector.cpp b/velox/experimental/cudf/connectors/parquet/ParquetConnector.cpp new file mode 100644 index 00000000000..d53d02f0f4a --- /dev/null +++ b/velox/experimental/cudf/connectors/parquet/ParquetConnector.cpp @@ -0,0 +1,75 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#include "velox/experimental/cudf/connectors/parquet/ParquetConnector.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetDataSource.h" + +namespace facebook::velox::cudf_velox::connector::parquet { + +using namespace facebook::velox::connector; + +ParquetConnector::ParquetConnector( + const std::string& id, + std::shared_ptr config, + folly::Executor* executor) + : Connector(id), + parquetConfig_(std::make_shared(config)), + executor_(executor) { + LOG(INFO) << "cudf::Parquet connector " << connectorId() << " created."; +} + +std::unique_ptr ParquetConnector::createDataSource( + const std::shared_ptr& outputType, + const std::shared_ptr& tableHandle, + const std::unordered_map>& + columnHandles, + ConnectorQueryCtx* connectorQueryCtx) { + return std::make_unique( + outputType, + tableHandle, + columnHandles, + executor_, + connectorQueryCtx, + parquetConfig_); +} + +std::unique_ptr ParquetConnector::createDataSink( + RowTypePtr inputType, + std::shared_ptr connectorInsertTableHandle, + ConnectorQueryCtx* connectorQueryCtx, + CommitStrategy /*commitStrategy*/) { + auto parquetInsertHandle = + std::dynamic_pointer_cast( + connectorInsertTableHandle); + VELOX_CHECK_NOT_NULL( + parquetInsertHandle, "Parquet connector expecting parquet write handle!"); + return std::make_unique( + inputType, + parquetInsertHandle, + connectorQueryCtx, + CommitStrategy::kNoCommit, + parquetConfig_); +} + +std::shared_ptr ParquetConnectorFactory::newConnector( + const std::string& id, + std::shared_ptr config, + folly::Executor* ioExecutor, + folly::Executor* cpuExecutor) { + return std::make_shared(id, config, ioExecutor); +} + +} // namespace facebook::velox::cudf_velox::connector::parquet diff --git a/velox/experimental/cudf/connectors/parquet/ParquetConnector.h b/velox/experimental/cudf/connectors/parquet/ParquetConnector.h new file mode 100644 index 00000000000..2e6b24a3be6 --- /dev/null +++ b/velox/experimental/cudf/connectors/parquet/ParquetConnector.h @@ -0,0 +1,84 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#pragma once + +#include "velox/experimental/cudf/connectors/parquet/ParquetConfig.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetDataSink.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetDataSource.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetTableHandle.h" + +#include "velox/connectors/Connector.h" + +#include +#include +#include + +namespace facebook::velox::cudf_velox::connector::parquet { + +using namespace facebook::velox::connector; +using namespace facebook::velox::config; + +class ParquetConnector final : public Connector { + public: + ParquetConnector( + const std::string& id, + std::shared_ptr config, + folly::Executor* executor); + + std::unique_ptr createDataSource( + const std::shared_ptr& outputType, + const std::shared_ptr& tableHandle, + const std::unordered_map>& + columnHandles, + ConnectorQueryCtx* connectorQueryCtx) override final; + + const std::shared_ptr& connectorConfig() const override { + return parquetConfig_->config(); + } + + std::unique_ptr createDataSink( + RowTypePtr inputType, + std::shared_ptr connectorInsertTableHandle, + ConnectorQueryCtx* connectorQueryCtx, + CommitStrategy commitStrategy) override final; + + folly::Executor* executor() const override { + return executor_; + } + + protected: + const std::shared_ptr parquetConfig_; + folly::Executor* executor_; +}; + +class ParquetConnectorFactory : public ConnectorFactory { + public: + static constexpr const char* kParquetConnectorName = "parquet"; + + ParquetConnectorFactory() : ConnectorFactory(kParquetConnectorName) {} + + explicit ParquetConnectorFactory(const char* connectorName) + : ConnectorFactory(connectorName) {} + + std::shared_ptr newConnector( + const std::string& id, + std::shared_ptr config, + folly::Executor* ioExecutor = nullptr, + folly::Executor* cpuExecutor = nullptr) override; +}; + +} // namespace facebook::velox::cudf_velox::connector::parquet diff --git a/velox/experimental/cudf/connectors/parquet/ParquetConnectorSplit.cpp b/velox/experimental/cudf/connectors/parquet/ParquetConnectorSplit.cpp new file mode 100644 index 00000000000..1dc05127659 --- /dev/null +++ b/velox/experimental/cudf/connectors/parquet/ParquetConnectorSplit.cpp @@ -0,0 +1,43 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#include "velox/experimental/cudf/connectors/parquet/ParquetConnectorSplit.h" + +#include + +namespace facebook::velox::cudf_velox::connector::parquet { + +std::string ParquetConnectorSplit::toString() const { + return fmt::format("Parquet: {}", filePath); +} + +std::string ParquetConnectorSplit::getFileName() const { + const auto i = filePath.rfind('/'); + return i == std::string::npos ? filePath : filePath.substr(i + 1); +} + +// static +std::shared_ptr ParquetConnectorSplit::create( + const folly::dynamic& obj) { + const auto connectorId = obj["connectorId"].asString(); + const auto splitWeight = obj["splitWeight"].asInt(); + const auto filePath = obj["filePath"].asString(); + + return std::make_shared( + connectorId, filePath, splitWeight); +} + +} // namespace facebook::velox::cudf_velox::connector::parquet diff --git a/velox/experimental/cudf/connectors/parquet/ParquetConnectorSplit.h b/velox/experimental/cudf/connectors/parquet/ParquetConnectorSplit.h new file mode 100644 index 00000000000..72e9ba7a572 --- /dev/null +++ b/velox/experimental/cudf/connectors/parquet/ParquetConnectorSplit.h @@ -0,0 +1,80 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#pragma once + +#include "velox/connectors/Connector.h" +#include "velox/dwio/common/Options.h" + +#include + +#include + +namespace facebook::velox::cudf_velox::connector::parquet { + +struct ParquetConnectorSplit + : public facebook::velox::connector::ConnectorSplit { + const std::string filePath; + const facebook::velox::dwio::common::FileFormat fileFormat{ + facebook::velox::dwio::common::FileFormat::PARQUET}; + const cudf::io::source_info cudfSourceInfo; + + ParquetConnectorSplit( + const std::string& connectorId, + const std::string& _filePath, + int64_t _splitWeight = 0) + : facebook::velox::connector::ConnectorSplit(connectorId, _splitWeight), + filePath(_filePath), + cudfSourceInfo({filePath}) {} + + std::string toString() const override; + std::string getFileName() const; + + const cudf::io::source_info& getCudfSourceInfo() const { + return cudfSourceInfo; + } + + static std::shared_ptr create( + const folly::dynamic& obj); +}; + +class ParquetConnectorSplitBuilder { + public: + explicit ParquetConnectorSplitBuilder(std::string filePath) + : filePath_{std::move(filePath)} {} + + ParquetConnectorSplitBuilder& splitWeight(int64_t splitWeight) { + splitWeight_ = splitWeight; + return *this; + } + + ParquetConnectorSplitBuilder& connectorId(const std::string& connectorId) { + connectorId_ = connectorId; + return *this; + } + + std::shared_ptr build() const { + return std::make_shared( + connectorId_, filePath_, splitWeight_); + } + + private: + const std::string filePath_; + std::string connectorId_; + int64_t splitWeight_{0}; +}; + +} // namespace facebook::velox::cudf_velox::connector::parquet diff --git a/velox/experimental/cudf/connectors/parquet/ParquetDataSink.cpp b/velox/experimental/cudf/connectors/parquet/ParquetDataSink.cpp new file mode 100644 index 00000000000..00c27048fc9 --- /dev/null +++ b/velox/experimental/cudf/connectors/parquet/ParquetDataSink.cpp @@ -0,0 +1,569 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#include "velox/experimental/cudf/connectors/parquet/ParquetConfig.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetDataSink.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetTableHandle.h" +#include "velox/experimental/cudf/exec/Utilities.h" +#include "velox/experimental/cudf/exec/VeloxCudfInterop.h" +#include "velox/experimental/cudf/vector/CudfVector.h" + +#include "velox/common/base/Counters.h" +#include "velox/common/base/Fs.h" +#include "velox/common/base/StatsReporter.h" +#include "velox/dwio/common/Options.h" +#include "velox/exec/OperatorUtils.h" + +#include +#include +#include +#include +#include + +#include +#include +#include + +using facebook::velox::common::testutil::TestValue; + +namespace facebook::velox::cudf_velox::connector::parquet { + +namespace { + +std::unordered_map tableTypeNames() { + return { + {LocationHandle::TableType::kNew, "kNew"}, + }; +} + +template +std::unordered_map invertMap(const std::unordered_map& mapping) { + std::unordered_map inverted; + for (const auto& [key, value] : mapping) { + inverted.emplace(value, key); + } + return inverted; +} + +uint64_t getFinishTimeSliceLimitMsFromParquetConfig( + const std::shared_ptr& config, + const config::ConfigBase* sessions) { + const uint64_t flushTimeSliceLimitMsFromConfig = + config->sortWriterFinishTimeSliceLimitMs(sessions); + // NOTE: if the flush time slice limit is set to 0, then we treat it as no + // limit. + return flushTimeSliceLimitMsFromConfig == 0 + ? std::numeric_limits::max() + : flushTimeSliceLimitMsFromConfig; +} + +std::string makeUuid() { + return boost::lexical_cast(boost::uuids::random_generator()()); +} + +cudf::io::compression_type getCompressionType( + facebook::velox::common::CompressionKind name) { + using CompressionType = cudf::io::compression_type; + + static std::unordered_map< + facebook::velox::common::CompressionKind, + CompressionType> const kMap = { + {facebook::velox::common::CompressionKind::CompressionKind_NONE, + CompressionType::NONE}, + {facebook::velox::common::CompressionKind::CompressionKind_SNAPPY, + CompressionType::SNAPPY}, + {facebook::velox::common::CompressionKind::CompressionKind_LZ4, + CompressionType::LZ4}, + {facebook::velox::common::CompressionKind::CompressionKind_ZSTD, + CompressionType::ZSTD}}; + + VELOX_CHECK( + kMap.find(name) != kMap.end(), + "Unsupported compression type requested. Supported compression types are: " + "NONE, SNAPPY, LZ4, ZSTD"); + + return kMap.at(name); +} + +std::shared_ptr createSinkPool( + const std::shared_ptr& writerPool) { + return writerPool->addLeafChild(fmt::format("{}.sink", writerPool->name())); +} + +std::shared_ptr createSortPool( + const std::shared_ptr& writerPool) { + return writerPool->addLeafChild(fmt::format("{}.sort", writerPool->name())); +} + +} // namespace + +const std::string LocationHandle::tableTypeName( + LocationHandle::TableType type) { + static const auto kTableTypes = tableTypeNames(); + return kTableTypes.at(type); +} + +LocationHandle::TableType LocationHandle::tableTypeFromName( + const std::string& name) { + static const auto kNameTableTypes = invertMap(tableTypeNames()); + return kNameTableTypes.at(name); +} + +ParquetDataSink::ParquetDataSink( + RowTypePtr inputType, + std::shared_ptr insertTableHandle, + const ConnectorQueryCtx* connectorQueryCtx, + CommitStrategy commitStrategy, + const std::shared_ptr& parquetConfig) + : inputType_(std::move(inputType)), + insertTableHandle_(std::move(insertTableHandle)), + connectorQueryCtx_(connectorQueryCtx), + commitStrategy_(commitStrategy), + parquetConfig_(parquetConfig), + spillConfig_(connectorQueryCtx->spillConfig()), + sortWriterFinishTimeSliceLimitMs_( + getFinishTimeSliceLimitMsFromParquetConfig( + parquetConfig_, + connectorQueryCtx->sessionProperties())) { + VELOX_USER_CHECK( + (commitStrategy_ == CommitStrategy::kNoCommit) || + (commitStrategy_ == CommitStrategy::kTaskCommit), + "Unsupported commit strategy: {}", + commitStrategyToString(commitStrategy_)); + + const auto& writerOptions = dynamic_cast( + insertTableHandle_->writerOptions().get()); + + if (writerOptions != nullptr) { + sortingColumns_ = std::move(writerOptions->sortingColumns); + } +} + +void ParquetDataSink::appendData(RowVectorPtr input) { + checkRunning(); + + // Convert the input RowVectorPtr to cudf::table + auto stream = cudfGlobalStreamPool().get_stream(); + auto cudfInput = with_arrow::toCudfTable(input, input->pool(), stream); + stream.synchronize(); + VELOX_CHECK_NOT_NULL( + cudfInput, "Failed to convert input RowVectorPtr to cudf::table"); + + // Check if the writer doesn't already exist + if (writer_ == nullptr) { + writer_ = createCudfWriter(cudfInput->view()); + } + + // Write the table to the sink + writer_->write(cudfInput->view()); + writerInfo_->inputSizeInBytes += input->estimateFlatSize(); + writerInfo_->numWrittenRows += input->size(); +} + +std::unique_ptr +ParquetDataSink::createCudfWriter(cudf::table_view cudfTable) { + // Create a table_input_metadata from the input + auto tableInputMetadata = createCudfTableInputMetadata(cudfTable); + + auto compressionKind = + getCompressionType(insertTableHandle_->compressionKind().value_or( + facebook::velox::common::CompressionKind::CompressionKind_NONE)); + + // Create a sink and writer + const auto& locationHandle = insertTableHandle_->locationHandle(); + const auto targetFileName = locationHandle->targetFileName().empty() + ? fmt::format("{}{}", makeUuid(), ".parquet") + : locationHandle->targetFileName(); + + auto writerParameters = ParquetWriterParameters( + ParquetWriterParameters::UpdateMode::kNew, + targetFileName, + locationHandle->targetPath()); + + const auto writePath = fs::path(writerParameters.writeDirectory()) / + writerParameters.writeFileName(); + + makeWriterOptions(writerParameters); + + // Create writer options for the given sink + const auto sinkInfo = cudf::io::sink_info( + fmt::format("{}/{}", locationHandle->targetPath(), targetFileName)); + auto cudfWriterOptions = + cudf::io::chunked_parquet_writer_options::builder(sinkInfo) + .metadata(tableInputMetadata) + .utc_timestamps(parquetConfig_->writeTimestampsAsUTC()) + .write_arrow_schema(parquetConfig_->writeArrowSchema()) + .write_v2_headers(parquetConfig_->writev2PageHeaders()) + .compression(compressionKind) + .build(); + + const auto& writerOptions = dynamic_cast( + insertTableHandle_->writerOptions().get()); + + // If non-null writerOptions were passed, pass them to the chunked parquet + // writer options + if (writerOptions != nullptr) { + // Set encoding for all columns + std::for_each( + tableInputMetadata.column_metadata.begin(), + tableInputMetadata.column_metadata.end(), + [=](auto& colMeta) { colMeta.set_encoding(writerOptions->encoding); }); + + cudfWriterOptions.set_row_group_size_bytes( + writerOptions->rowGroupSizeBytes); + cudfWriterOptions.set_row_group_size_rows(writerOptions->rowGroupSizeRows); + cudfWriterOptions.set_max_page_size_bytes(writerOptions->maxPageSizeBytes); + cudfWriterOptions.set_max_page_size_rows(writerOptions->maxPageSizeRows); + cudfWriterOptions.set_dictionary_policy(writerOptions->dictionaryPolicy); + cudfWriterOptions.set_max_dictionary_size(writerOptions->maxDictionarySize); + cudfWriterOptions.enable_int96_timestamps( + writerOptions->writeTimestampsAsInt96); + + // Enable if enabled in the session or the writerOptions + cudfWriterOptions.enable_utc_timestamps( + parquetConfig_->writeTimestampsAsUTC() or + writerOptions->writeTimestampsAsUTC); + cudfWriterOptions.enable_write_arrow_schema( + parquetConfig_->writeArrowSchema() or writerOptions->writeArrowSchema); + cudfWriterOptions.enable_write_v2_headers( + parquetConfig_->writev2PageHeaders() or writerOptions->v2PageHeaders); + cudfWriterOptions.set_stats_level(writerOptions->statsLevel); + + if (writerOptions->maxPageFragmentSize.has_value()) { + cudfWriterOptions.set_max_page_fragment_size( + writerOptions->maxPageFragmentSize.value()); + } + // Get compression stats if needed + if (writerOptions->compressionStats != nullptr) { + cudfWriterOptions.set_compression_statistics( + writerOptions->compressionStats); + } + // Write sorting columns if available + if (sortingColumns_.empty()) { + cudfWriterOptions.set_sorting_columns(sortingColumns_); + } + } + + return std::make_unique(cudfWriterOptions); +} + +cudf::io::table_input_metadata ParquetDataSink::createCudfTableInputMetadata( + cudf::table_view cudfTable) { + auto tableInputMetadata = cudf::io::table_input_metadata(cudfTable); + auto inputColumns = insertTableHandle_->inputColumns(); + + // Check if equal number of columns in the input and + // ParquetInsertTableHandle + VELOX_CHECK_EQ( + tableInputMetadata.column_metadata.size(), + inputColumns.size(), + "Unequal number of columns in the input and ParquetInsertTableHandle"); + + std::function + setColumnName = [&](cudf::io::column_in_metadata& colMeta, + const ParquetColumnHandle& columnHandle) { + // Check if equal number of children + const auto& childrenHandles = columnHandle.children(); + + // Warn if the mismatch in the number of child cols in Parquet + // table_metadata and columnHandles + if (colMeta.num_children() != childrenHandles.size()) { + LOG(WARNING) << fmt::format( + "({} vs {}): Unequal number of child columns in Parquet table_metadata and ColumnHandles", + colMeta.num_children(), + childrenHandles.size()); + } + + // Set children's names + for (int32_t i = 0; i < + std::min(colMeta.num_children(), childrenHandles.size()); + ++i) { + setColumnName(colMeta.child(i), childrenHandles[i]); + } + // Set this column's name + colMeta.set_name(columnHandle.name()); + }; + + // Set names for all columns and their children + for (int32_t i = 0; i < tableInputMetadata.column_metadata.size(); ++i) { + setColumnName(tableInputMetadata.column_metadata[i], *inputColumns[i]); + } + + return tableInputMetadata; +} + +std::string ParquetDataSink::stateString(State state) { + switch (state) { + case State::kRunning: + return "RUNNING"; + case State::kFinishing: + return "FLUSHING"; + case State::kClosed: + return "CLOSED"; + case State::kAborted: + return "ABORTED"; + default: + VELOX_UNREACHABLE("BAD STATE: {}", static_cast(state)); + } +} + +DataSink::Stats ParquetDataSink::stats() const { + Stats stats; + if (state_ == State::kAborted) { + return stats; + } + + int64_t numWrittenBytes{0}; + int64_t writeIOTimeUs{0}; + + numWrittenBytes += ioStats_->rawBytesWritten(); + writeIOTimeUs += ioStats_->writeIOTimeUs(); + + stats.numWrittenBytes = numWrittenBytes; + stats.writeIOTimeUs = writeIOTimeUs; + + if (state_ != State::kClosed) { + return stats; + } + + stats.numWrittenFiles = 1; + VELOX_CHECK_NOT_NULL(writerInfo_); + const auto spillStats = writerInfo_->spillStats->rlock(); + if (!spillStats->empty()) { + stats.spillStats += *spillStats; + } + + return stats; +} + +void ParquetDataSink::setState(State newState) { + checkStateTransition(state_, newState); + state_ = newState; +} + +/// Validates the state transition from 'oldState' to 'newState'. +void ParquetDataSink::checkStateTransition(State oldState, State newState) { + switch (oldState) { + case State::kRunning: + if (newState == State::kAborted || newState == State::kFinishing) { + return; + } + break; + case State::kFinishing: + if (newState == State::kAborted || newState == State::kClosed || + // The finishing state is reentry state if we yield in the + // middle of finish processing if a single run takes too long. + newState == State::kFinishing) { + return; + } + [[fallthrough]]; + case State::kAborted: + case State::kClosed: + default: + break; + } + VELOX_FAIL("Unexpected state transition from {} to {}", oldState, newState); +} + +bool ParquetDataSink::finish() { + VELOX_CHECK_NOT_NULL(writer_, "ParquetDataSink has no writer"); + + setState(State::kFinishing); + return true; +} + +std::vector ParquetDataSink::close() { + setState(State::kClosed); + closeInternal(); + + std::vector partitionUpdates{}; + + partitionUpdates.reserve(1); + VELOX_CHECK_NOT_NULL(writerInfo_); + // clang-format off + auto partitionUpdateJson = folly::toJson( + folly::dynamic::object + ("writePath", writerInfo_->writerParameters.writeDirectory()) + ("targetPath", writerInfo_->writerParameters.targetDirectory()) + ("fileWriteInfos", folly::dynamic::array( + folly::dynamic::object + ("writeFileName", writerInfo_->writerParameters.writeFileName()) + ("targetFileName", writerInfo_->writerParameters.targetFileName()) + ("fileSize", ioStats_->rawBytesWritten()))) + ("rowCount", writerInfo_->numWrittenRows) + ("inMemoryDataSizeInBytes", writerInfo_->inputSizeInBytes) + ("onDiskDataSizeInBytes", ioStats_->rawBytesWritten()) + ("containsNumberedFileNames", true)); + // clang-format on + partitionUpdates.emplace_back(partitionUpdateJson); + + return partitionUpdates; +} + +void ParquetDataSink::abort() { + setState(State::kAborted); + closeInternal(); +} + +void ParquetDataSink::closeInternal() { + VELOX_CHECK_NE(state_, State::kRunning); + VELOX_CHECK_NE(state_, State::kFinishing); + VELOX_CHECK_NOT_NULL(writer_, "ParquetDataSink has no writer"); + + TestValue::adjust( + "facebook::velox::connector::parquet::ParquetDataSink::closeInternal", + this); + + // Close cudf writer + writer_->close(); + + // Reset the unique pointers to Cudf writer and options + writer_.reset(); +} + +std::shared_ptr ParquetDataSink::createWriterPool() { + auto* connectorPool = connectorQueryCtx_->connectorMemoryPool(); + return connectorPool->addAggregateChild( + fmt::format("{}.{}", connectorPool->name(), "parquet-writer")); +} + +void ParquetDataSink::makeWriterOptions( + ParquetWriterParameters writerParameters) { + auto writerPool = createWriterPool(); + auto sinkPool = createSinkPool(writerPool); + std::shared_ptr sortPool{nullptr}; + if (sortWrite()) { + sortPool = createSortPool(writerPool); + } + + writerInfo_ = std::make_shared( + std::move(writerParameters), + std::move(writerPool), + std::move(sinkPool), + std::move(sortPool)); + + ioStats_ = std::make_shared(); + + // Take the writer options provided by the user as a starting point, + // or allocate a new one. + auto options = insertTableHandle_->writerOptions(); + if (!options) { + options = std::make_unique(); + } + + const auto* connectorSessionProperties = + connectorQueryCtx_->sessionProperties(); + + if (options->memoryPool == nullptr) { + options->memoryPool = writerInfo_->writerPool.get(); + } + + if (!options->compressionKind) { + options->compressionKind = insertTableHandle_->compressionKind(); + } + + const auto& sessionTimeZoneName = connectorQueryCtx_->sessionTimezone(); + if (!sessionTimeZoneName.empty()) { + options->sessionTimezoneName = sessionTimeZoneName; + } + options->adjustTimestampToTimezone = + connectorQueryCtx_->adjustTimestampToTimezone(); +} + +folly::dynamic ParquetInsertTableHandle::serialize() const { + folly::dynamic obj = folly::dynamic::object; + obj["name"] = "ParquetInsertTableHandle"; + folly::dynamic arr = folly::dynamic::array; + for (const auto& ic : inputColumns_) { + arr.push_back(ic->serialize()); + } + + obj["inputColumns"] = arr; + obj["locationHandle"] = locationHandle_->serialize(); + obj["tableStorageFormat"] = dwio::common::toString(storageFormat_); + + if (compressionKind_.has_value()) { + obj["compressionKind"] = common::compressionKindToString(*compressionKind_); + } + + return obj; +} + +ParquetInsertTableHandlePtr ParquetInsertTableHandle::create( + const folly::dynamic& obj) { + auto inputColumns = + ISerializable::deserialize>( + obj["inputColumns"]); + auto locationHandle = + ISerializable::deserialize(obj["locationHandle"]); + std::optional compressionKind = std::nullopt; + if (obj.count("compressionKind") > 0) { + compressionKind = + common::stringToCompressionKind(obj["compressionKind"].asString()); + } + std::unordered_map serdeParameters; + for (const auto& pair : obj["serdeParameters"].items()) { + serdeParameters.emplace(pair.first.asString(), pair.second.asString()); + } + return std::make_shared( + inputColumns, locationHandle, compressionKind, serdeParameters); +} + +std::string ParquetInsertTableHandle::toString() const { + std::ostringstream out; + out << "ParquetInsertTableHandle [" << dwio::common::toString(storageFormat_); + if (compressionKind_.has_value()) { + out << " " << common::compressionKindToString(compressionKind_.value()); + } else { + out << " none"; + } + out << "], [inputColumns: ["; + for (const auto& i : inputColumns_) { + out << " " << i->toString(); + } + out << " ], locationHandle: " << locationHandle_->toString(); + + out << "]"; + return out.str(); +} + +void ParquetInsertTableHandle::registerSerDe() { + auto& registry = DeserializationRegistryForSharedPtr(); + registry.Register("HiveInsertTableHandle", ParquetInsertTableHandle::create); +} + +std::string LocationHandle::toString() const { + return fmt::format( + "LocationHandle [targetPath: {}, tableType: {},", + targetPath_, + tableTypeName(tableType_)); +} + +folly::dynamic LocationHandle::serialize() const { + folly::dynamic obj = folly::dynamic::object; + obj["name"] = "LocationHandle"; + obj["targetPath"] = targetPath_; + obj["tableType"] = tableTypeName(tableType_); + return obj; +} + +LocationHandlePtr LocationHandle::create(const folly::dynamic& obj) { + auto targetPath = obj["targetPath"].asString(); + auto tableType = tableTypeFromName(obj["tableType"].asString()); + return std::make_shared(targetPath, tableType); +} + +} // namespace facebook::velox::cudf_velox::connector::parquet diff --git a/velox/experimental/cudf/connectors/parquet/ParquetDataSink.h b/velox/experimental/cudf/connectors/parquet/ParquetDataSink.h new file mode 100644 index 00000000000..f1dc47a8389 --- /dev/null +++ b/velox/experimental/cudf/connectors/parquet/ParquetDataSink.h @@ -0,0 +1,390 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ +#pragma once + +#include "velox/experimental/cudf/connectors/parquet/ParquetConfig.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetConnectorSplit.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetTableHandle.h" +#include "velox/experimental/cudf/connectors/parquet/WriterOptions.h" + +#include "velox/common/compression/Compression.h" +#include "velox/connectors/Connector.h" +#include "velox/dwio/common/Options.h" +#include "velox/dwio/common/Statistics.h" +#include "velox/dwio/common/WriterFactory.h" +#include "velox/exec/MemoryReclaimer.h" +#include "velox/type/Type.h" + +#include +#include +#include + +namespace facebook::velox::cudf_velox::connector::parquet { + +using namespace facebook::velox::connector; + +class LocationHandle; +using LocationHandlePtr = std::shared_ptr; + +/// Location related properties of the Parquet table to be written. +class LocationHandle : public ISerializable { + public: + enum class TableType { + /// Write to a new table to be created. + kNew, + }; + + LocationHandle( + std::string targetPath, + TableType tableType, + std::string targetFileName = "") + : targetPath_(std::move(targetPath)), + targetFileName_(std::move(targetFileName)), + tableType_(tableType) {} + + const std::string& targetPath() const { + return targetPath_; + } + + const std::string& targetFileName() const { + return targetFileName_; + } + + TableType tableType() const { + return tableType_; + } + + std::string toString() const; + + static void registerSerDe(); + + folly::dynamic serialize() const override; + + static LocationHandlePtr create(const folly::dynamic& obj); + + static const std::string tableTypeName(LocationHandle::TableType type); + + static LocationHandle::TableType tableTypeFromName(const std::string& name); + + private: + // Target directory path. + const std::string targetPath_; + // If non-empty, use this name instead of generating our own. + const std::string targetFileName_; + // Whether the table to be written is new, already existing or temporary. + const TableType tableType_; +}; + +/// Parameters for Hive writers. +class ParquetWriterParameters { + public: + enum class UpdateMode { + kNew, // Write files to a new directory. + }; + + /// @param updateMode Write the files to a new directory, or append to an + /// existing directory or overwrite an existing directory. + /// @param targetFileName The final name of a file after committing. + /// @param targetDirectory The final directory that a file should be in after + /// committing. + /// @param writeFileName The temporary name of the file that a running writer + /// writes to. If a running writer writes directory to the target file, set + /// writeFileName to targetFileName by default. + /// @param writeDirectory The temporary directory that a running writer writes + /// to. If a running writer writes directory to the target directory, set + /// writeDirectory to targetDirectory by default. + ParquetWriterParameters( + UpdateMode updateMode, + std::string targetFileName, + std::string targetDirectory, + std::optional writeFileName = std::nullopt, + std::optional writeDirectory = std::nullopt) + : updateMode_(updateMode), + targetFileName_(std::move(targetFileName)), + targetDirectory_(std::move(targetDirectory)), + writeFileName_(writeFileName.value_or(targetFileName_)), + writeDirectory_(writeDirectory.value_or(targetDirectory_)) {} + + UpdateMode updateMode() const { + return updateMode_; + } + + static std::string updateModeToString(UpdateMode updateMode) { + switch (updateMode) { + case UpdateMode::kNew: + return "NEW"; + default: + VELOX_UNSUPPORTED("Unsupported update mode."); + } + } + + const std::string& targetFileName() const { + return targetFileName_; + } + + const std::string& writeFileName() const { + return writeFileName_; + } + + const std::string& targetDirectory() const { + return targetDirectory_; + } + + const std::string& writeDirectory() const { + return writeDirectory_; + } + + private: + const UpdateMode updateMode_; + const std::optional partitionName_; + const std::string targetFileName_; + const std::string targetDirectory_; + const std::string writeFileName_; + const std::string writeDirectory_; +}; + +struct ParquetWriterInfo { + ParquetWriterInfo( + ParquetWriterParameters parameters, + std::shared_ptr _writerPool, + std::shared_ptr _sinkPool, + std::shared_ptr _sortPool) + : writerParameters(std::move(parameters)), + nonReclaimableSectionHolder(new tsan_atomic(false)), + spillStats(std::make_unique>()), + writerPool(std::move(_writerPool)), + sinkPool(std::move(_sinkPool)), + sortPool(std::move(_sortPool)) {} + + const ParquetWriterParameters writerParameters; + const std::unique_ptr> nonReclaimableSectionHolder; + /// Collects the spill stats from sort writer if the spilling has been + /// triggered. + const std::unique_ptr> spillStats; + const std::shared_ptr writerPool; + const std::shared_ptr sinkPool; + const std::shared_ptr sortPool; + int64_t numWrittenRows = 0; + int64_t inputSizeInBytes = 0; +}; + +class ParquetInsertTableHandle; +using ParquetInsertTableHandlePtr = std::shared_ptr; + +/// Represents a request for Parquet write. +class ParquetInsertTableHandle : public ConnectorInsertTableHandle { + public: + ParquetInsertTableHandle( + std::vector> inputColumns, + std::shared_ptr locationHandle, + std::optional compressionKind = {}, + const std::unordered_map& serdeParameters = {}, + const std::shared_ptr& writerOptions = + nullptr) + : inputColumns_(std::move(inputColumns)), + locationHandle_(std::move(locationHandle)), + compressionKind_(compressionKind), + serdeParameters_(serdeParameters), + writerOptions_(writerOptions) { + if (compressionKind.has_value()) { + VELOX_CHECK( + compressionKind.value() != common::CompressionKind_MAX, + "Unsupported compression type: CompressionKind_MAX"); + VELOX_CHECK( + compressionKind.value() == common::CompressionKind_NONE or + compressionKind.value() == common::CompressionKind_SNAPPY or + compressionKind.value() == common::CompressionKind_LZ4 or + compressionKind.value() == common::CompressionKind_ZSTD, + "Parquet DataSink only supports NONE, SNAPPY, LZ4, and ZSTD compressions."); + } + } + + virtual ~ParquetInsertTableHandle() = default; + + const std::vector>& inputColumns() + const { + return inputColumns_; + } + + const std::shared_ptr& locationHandle() const { + return locationHandle_; + } + + std::optional compressionKind() const { + return compressionKind_; + } + + const dwio::common::FileFormat storageFormat() const { + return storageFormat_; + } + + const std::unordered_map& serdeParameters() const { + return serdeParameters_; + } + + const std::shared_ptr& writerOptions() const { + return writerOptions_; + } + + bool supportsMultiThreading() const override { + return true; // TODO: Needs more testing if this is ok + } + + bool isExistingTable() const { + return false; // This is always false as cudf's Parquet writer doesn't yet + // support updating existing Parquet files + } + + folly::dynamic serialize() const override; + + static ParquetInsertTableHandlePtr create(const folly::dynamic& obj); + + static void registerSerDe(); + + std::string toString() const override; + + private: + const std::vector> inputColumns_; + const std::shared_ptr locationHandle_; + const std::optional compressionKind_; + const dwio::common::FileFormat storageFormat_ = + dwio::common::FileFormat::PARQUET; + const std::unordered_map serdeParameters_; + const std::shared_ptr writerOptions_; +}; + +class ParquetDataSink : public DataSink { + public: + /// The list of runtime stats reported by parquet data sink + static constexpr const char* kEarlyFlushedRawBytes = "earlyFlushedRawBytes"; + + /// Defines the execution states of a parquet data sink running internally. + enum class State { + /// The data sink accepts new append data in this state. + kRunning = 0, + /// The data sink flushes any buffered data to the underlying file writer + /// but no more data can be appended. + kFinishing = 1, + /// The data sink is aborted on error and no more data can be appended. + kAborted = 2, + /// The data sink is closed on error and no more data can be appended. + kClosed = 3 + }; + static std::string stateString(State state); + + ParquetDataSink( + RowTypePtr inputType, + std::shared_ptr insertTableHandle, + const ConnectorQueryCtx* connectorQueryCtx, + CommitStrategy commitStrategy, + const std::shared_ptr& parquetConfig); + + void appendData(RowVectorPtr input) override; + + bool finish() override; + + Stats stats() const override; + + std::vector close() override; + + void abort() override; + + bool canReclaim() const { + return false; + }; + + private: + // Creates a new cudf chunked parquet writer. + std::unique_ptr createCudfWriter( + cudf::table_view cudfTable); + cudf::io::table_input_metadata createCudfTableInputMetadata( + cudf::table_view cudfTable); + + // Validates the state transition from 'oldState' to 'newState'. + void checkStateTransition(State oldState, State newState); + void setState(State newState); + + std::shared_ptr createWriterPool(); + + FOLLY_ALWAYS_INLINE bool sortWrite() const { + return not sortingColumns_.empty(); + } + + FOLLY_ALWAYS_INLINE bool isCommitRequired() const { + return false; // Since we always immediately write + } + + FOLLY_ALWAYS_INLINE void checkRunning() const { + VELOX_CHECK_EQ(state_, State::kRunning, "Parquet data sink is not running"); + } + + void closeInternal(); + void makeWriterOptions(ParquetWriterParameters writerParameters); + + const RowTypePtr inputType_; + const std::shared_ptr insertTableHandle_; + const ConnectorQueryCtx* const connectorQueryCtx_; + const CommitStrategy commitStrategy_; + const std::shared_ptr parquetConfig_; + const common::SpillConfig* const spillConfig_; + const uint64_t sortWriterFinishTimeSliceLimitMs_{0}; + State state_{State::kRunning}; + + // Below are structures for partitions from all inputs. writerInfo_ and + // writers_ are both indexed by partitionId. + std::unique_ptr writer_; + + std::vector sortingColumns_; + + std::shared_ptr writerInfo_; + + // IO statistics collected for writer. + std::shared_ptr ioStats_; +}; + +FOLLY_ALWAYS_INLINE std::ostream& operator<<( + std::ostream& os, + ParquetDataSink::State state) { + os << ParquetDataSink::stateString(state); + return os; +} +} // namespace facebook::velox::cudf_velox::connector::parquet + +template <> +struct fmt::formatter< + facebook::velox::cudf_velox::connector::parquet::ParquetDataSink::State> + : formatter { + auto format( + facebook::velox::cudf_velox::connector::parquet::ParquetDataSink::State s, + format_context& ctx) const { + return formatter::format( + facebook::velox::cudf_velox::connector::parquet::ParquetDataSink:: + stateString(s), + ctx); + } +}; + +template <> +struct fmt::formatter< + facebook::velox::cudf_velox::connector::parquet::LocationHandle::TableType> + : formatter { + auto format( + facebook::velox::cudf_velox::connector::parquet::LocationHandle::TableType + s, + format_context& ctx) const { + return formatter::format(static_cast(s), ctx); + } +}; diff --git a/velox/experimental/cudf/connectors/parquet/ParquetDataSource.cpp b/velox/experimental/cudf/connectors/parquet/ParquetDataSource.cpp new file mode 100644 index 00000000000..57da2d68bd9 --- /dev/null +++ b/velox/experimental/cudf/connectors/parquet/ParquetDataSource.cpp @@ -0,0 +1,255 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#include "velox/experimental/cudf/connectors/parquet/ParquetConfig.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetConnectorSplit.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetDataSource.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetTableHandle.h" +#include "velox/experimental/cudf/exec/ToCudf.h" +#include "velox/experimental/cudf/exec/Utilities.h" +#include "velox/experimental/cudf/exec/VeloxCudfInterop.h" +#include "velox/experimental/cudf/vector/CudfVector.h" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace facebook::velox::cudf_velox::connector::parquet { + +using namespace facebook::velox::connector; + +ParquetDataSource::ParquetDataSource( + const std::shared_ptr& outputType, + const std::shared_ptr& tableHandle, + const std::unordered_map>& + columnHandles, + folly::Executor* executor, + const ConnectorQueryCtx* connectorQueryCtx, + const std::shared_ptr& parquetConfig) + : NvtxHelper( + nvtx3::rgb{80, 171, 241}, // Parquet blue, + std::nullopt, + fmt::format("[{}]", tableHandle->name())), + parquetConfig_(parquetConfig), + executor_(executor), + connectorQueryCtx_(connectorQueryCtx), + pool_(connectorQueryCtx->memoryPool()), + outputType_(outputType), + expressionEvaluator_(connectorQueryCtx->expressionEvaluator()) { + // Set up column projection if needed + auto readColumnTypes = outputType_->children(); + for (const auto& outputName : outputType_->names()) { + auto it = columnHandles.find(outputName); + VELOX_CHECK( + it != columnHandles.end(), + "ColumnHandle is missing for output column: {}", + outputName); + + auto* handle = static_cast(it->second.get()); + readColumnNames_.emplace_back(handle->name()); + } + + // Dynamic cast tableHandle to ParquetTableHandle + tableHandle_ = std::dynamic_pointer_cast(tableHandle); + VELOX_CHECK_NOT_NULL( + tableHandle_, "TableHandle must be an instance of ParquetTableHandle"); + + // Create empty IOStats for later use + ioStats_ = std::make_shared(); + + // Create subfield filter + auto subfieldFilter = tableHandle_->subfieldFilterExpr(); + if (subfieldFilter) { + subfieldFilterExprSet_ = expressionEvaluator_->compile(subfieldFilter); + } + + // Create remaining filter + auto remainingFilter = tableHandle_->remainingFilter(); + if (remainingFilter) { + remainingFilterExprSet_ = expressionEvaluator_->compile(remainingFilter); + cudfExpressionEvaluator_ = velox::cudf_velox::ExpressionEvaluator( + remainingFilterExprSet_->exprs(), outputType_); + // TODO(kn): Get column names and subfields from remaining filter and add to + // readColumnNames_ + } +} + +std::optional ParquetDataSource::next( + uint64_t /*size*/, + velox::ContinueFuture& /* future */) { + VELOX_NVTX_OPERATOR_FUNC_RANGE(); + // Basic sanity checks + VELOX_CHECK_NOT_NULL(split_, "No split to process. Call addSplit first."); + VELOX_CHECK_NOT_NULL(splitReader_, "No split reader present"); + + if (not splitReader_->has_next()) { + return nullptr; + } + + // Read a table chunk + auto [table, metadata] = splitReader_->read_chunk(); + auto cudfTable = std::move(table); + // Fill in the column names if reading the first chunk. + if (columnNames_.empty()) { + for (auto schema : metadata.schema_info) { + columnNames_.emplace_back(schema.name); + } + } + + // Apply remaining filter if present + if (remainingFilterExprSet_) { + auto cudfTableColumns = cudfTable->release(); + const auto originalNumColumns = cudfTableColumns.size(); + // Filter may need addtional computed columns which are added to + // cudfTableColumns + auto filterResult = cudfExpressionEvaluator_.compute( + cudfTableColumns, stream_, cudf::get_current_device_resource_ref()); + // discard computed columns + std::vector> originalColumns; + originalColumns.reserve(originalNumColumns); + std::move( + cudfTableColumns.begin(), + cudfTableColumns.begin() + originalNumColumns, + std::back_inserter(originalColumns)); + auto originalTable = + std::make_unique(std::move(originalColumns)); + // Keep only rows where the filter is true + cudfTable = cudf::apply_boolean_mask( + *originalTable, + *filterResult[0], + stream_, + cudf::get_current_device_resource_ref()); + } + + // Output RowVectorPtr + const auto nRows = cudfTable->num_rows(); + + // keep only outputType_.size() columns in cudfTable_ + if (outputType_->size() < cudfTable->num_columns()) { + auto cudfTableColumns = cudfTable->release(); + std::vector> originalColumns; + originalColumns.reserve(outputType_->size()); + std::move( + cudfTableColumns.begin(), + cudfTableColumns.begin() + outputType_->size(), + std::back_inserter(originalColumns)); + cudfTable = std::make_unique(std::move(originalColumns)); + } + + auto output = cudfIsRegistered() + ? std::make_shared( + pool_, outputType_, nRows, std::move(cudfTable), stream_) + : with_arrow::toVeloxColumn( + cudfTable->view(), pool_, outputType_->names(), stream_); + stream_.synchronize(); + + // Check if conversion yielded a nullptr + VELOX_CHECK_NOT_NULL(output, "Cudf to Velox conversion yielded a nullptr"); + + // Update completedRows_. + completedRows_ += output->size(); + + // TODO: Update `completedBytes_` here instead of in `addSplit()` + + return output; +} + +void ParquetDataSource::addSplit(std::shared_ptr split) { + // Dynamic cast split to `ParquetConnectorSplit` + split_ = std::dynamic_pointer_cast(split); + VLOG(1) << "Adding split " << split_->toString(); + + // Split reader already exists, reset + if (splitReader_) { + splitReader_.reset(); + } + + // Clear columnNames if not empty + if (not columnNames_.empty()) { + columnNames_.clear(); + } + + // Create a `cudf::io::chunked_parquet_reader` SplitReader + splitReader_ = createSplitReader(); + + // TODO: `completedBytes_` should be updated in `next()` as we read more and + // more table bytes + const auto& filePaths = split_->getCudfSourceInfo().filepaths(); + for (const auto& filePath : filePaths) { + completedBytes_ += std::filesystem::file_size(filePath); + } +} + +std::unique_ptr +ParquetDataSource::createSplitReader() { + // Reader options + auto readerOptions = + cudf::io::parquet_reader_options::builder(split_->getCudfSourceInfo()) + .skip_rows(parquetConfig_->skipRows()) + .use_pandas_metadata(parquetConfig_->isUsePandasMetadata()) + .use_arrow_schema(parquetConfig_->isUseArrowSchema()) + .allow_mismatched_pq_schemas( + parquetConfig_->isAllowMismatchedParquetSchemas()) + .timestamp_type(parquetConfig_->timestampType()) + .build(); + + // Set num_rows only if available + if (parquetConfig_->numRows().has_value()) { + readerOptions.set_num_rows(parquetConfig_->numRows().value()); + } + + // Set column projection if needed + if (readColumnNames_.size()) { + readerOptions.set_columns(readColumnNames_); + } + if (subfieldFilterExprSet_) { + auto subfieldFilterExpr = subfieldFilterExprSet_->expr(0); + // non-ast instructions in filter is not supported for SubFieldFilter. + // precomputeInstructions which are non-ast instructions should be empty. + std::vector precomputeInstructions; + createAstTree( + subfieldFilterExpr, + subfieldTree_, + subfieldScalars_, + outputType_, + precomputeInstructions); + VELOX_CHECK_EQ(precomputeInstructions.size(), 0); + readerOptions.set_filter(subfieldTree_.back()); + } + stream_ = cudfGlobalStreamPool().get_stream(); + // Create a parquet reader + return std::make_unique( + parquetConfig_->maxChunkReadLimit(), + parquetConfig_->maxPassReadLimit(), + readerOptions, + stream_, + cudf::get_current_device_resource_ref()); +} + +void ParquetDataSource::resetSplit() { + split_.reset(); + splitReader_.reset(); + columnNames_.clear(); +} + +} // namespace facebook::velox::cudf_velox::connector::parquet diff --git a/velox/experimental/cudf/connectors/parquet/ParquetDataSource.h b/velox/experimental/cudf/connectors/parquet/ParquetDataSource.h new file mode 100644 index 00000000000..4de52278da4 --- /dev/null +++ b/velox/experimental/cudf/connectors/parquet/ParquetDataSource.h @@ -0,0 +1,139 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#pragma once + +#include "velox/experimental/cudf/connectors/parquet/ParquetConfig.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetConnectorSplit.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetTableHandle.h" +#include "velox/experimental/cudf/exec/ExpressionEvaluator.h" +#include "velox/experimental/cudf/exec/NvtxHelper.h" + +#include "velox/common/base/RandomUtil.h" +#include "velox/common/io/IoStatistics.h" +#include "velox/connectors/Connector.h" +#include "velox/dwio/common/Statistics.h" +#include "velox/type/Type.h" + +#include +#include + +namespace facebook::velox::cudf_velox::connector::parquet { + +using namespace facebook::velox::connector; + +class ParquetDataSource : public DataSource, public NvtxHelper { + public: + ParquetDataSource( + const std::shared_ptr& outputType, + const std::shared_ptr& tableHandle, + const std::unordered_map>& + columnHandles, + folly::Executor* executor, + const ConnectorQueryCtx* connectorQueryCtx, + const std::shared_ptr& ParquetConfig); + + void addSplit(std::shared_ptr split) override; + + void addDynamicFilter( + column_index_t /*outputChannel*/, + const std::shared_ptr& /*filter*/) + override { + VELOX_NYI("Dynamic filters not yet implemented by cudf::ParquetConnector."); + } + + std::optional next( + uint64_t size, + velox::ContinueFuture& /* future */) override; + + uint64_t getCompletedRows() override { + return completedRows_; + } + + uint64_t getCompletedBytes() override { + return completedBytes_; + } + + std::unordered_map runtimeStats() override { + // TODO: Which stats do we want to expose here? + return {}; + } + + private: + // Create a cudf::io::chunked_parquet_reader with the given split. + std::unique_ptr createSplitReader(); + // Clear split_ and splitReader after split has been fully processed. Keep + // readers around to hold adaptation. + void resetSplit(); + // Clear cudfTable_ and currentCudfTableView_ once we have successfully + // converted it to `RowVectorPtr` and returned. + void resetCudfTableAndView(); + const RowVectorPtr& getEmptyOutput() { + if (!emptyOutput_) { + emptyOutput_ = RowVector::createEmpty(outputType_, pool_); + } + return emptyOutput_; + } + RowVectorPtr emptyOutput_; + + std::shared_ptr split_; + std::shared_ptr tableHandle_; + + const std::shared_ptr parquetConfig_; + + folly::Executor* const executor_; + const ConnectorQueryCtx* const connectorQueryCtx_; + + memory::MemoryPool* const pool_; + + // cuDF Parquet reader stuff. + cudf::io::parquet_reader_options readerOptions_; + std::unique_ptr splitReader_; + rmm::cuda_stream_view stream_; + + // Table column names read from the Parquet file + std::vector columnNames_; + + // Output type from file reader. This is different from outputType_ that it + // contains column names before assignment, and columns that only used in + // remaining filter. + RowTypePtr readerOutputType_; + + // Columns to read. + std::vector readColumnNames_; + + std::shared_ptr ioStats_; + + size_t completedRows_{0}; + size_t completedBytes_{0}; + + // The row type for the data source output, not including filter-only columns + const RowTypePtr outputType_; + + // Expression evaluator for remaining filter. + core::ExpressionEvaluator* const expressionEvaluator_; + std::unique_ptr remainingFilterExprSet_; + velox::cudf_velox::ExpressionEvaluator cudfExpressionEvaluator_; + + // Expression evaluator for subfield filter. + std::vector> subfieldScalars_; + cudf::ast::tree subfieldTree_; + std::unique_ptr subfieldFilterExprSet_; + + dwio::common::RuntimeStatistics runtimeStats_; +}; + +} // namespace facebook::velox::cudf_velox::connector::parquet diff --git a/velox/experimental/cudf/connectors/parquet/ParquetTableHandle.cpp b/velox/experimental/cudf/connectors/parquet/ParquetTableHandle.cpp new file mode 100644 index 00000000000..dcb51868516 --- /dev/null +++ b/velox/experimental/cudf/connectors/parquet/ParquetTableHandle.cpp @@ -0,0 +1,64 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#include "velox/experimental/cudf/connectors/parquet/ParquetTableHandle.h" + +#include "velox/connectors/Connector.h" +#include "velox/type/Type.h" + +#include + +namespace facebook::velox::cudf_velox::connector::parquet { + +using namespace facebook::velox::connector; + +std::string ParquetColumnHandle::toString() const { + std::ostringstream out; + out << fmt::format( + "ParquetColumnHandle [name: {}, Type: {},", name_, type_->toString()); + return out.str(); +} + +ParquetTableHandle::ParquetTableHandle( + std::string connectorId, + const std::string& tableName, + bool filterPushdownEnabled, + const core::TypedExprPtr& subfieldFilterExpr, + const core::TypedExprPtr& remainingFilter, + const RowTypePtr& dataColumns) + : ConnectorTableHandle(std::move(connectorId)), + tableName_(tableName), + filterPushdownEnabled_(filterPushdownEnabled), + subfieldFilterExpr_(subfieldFilterExpr), + remainingFilter_(remainingFilter), + dataColumns_(dataColumns) {} + +std::string ParquetTableHandle::toString() const { + std::stringstream out; + out << "table: " << tableName_; + if (dataColumns_) { + out << ", data columns: " << dataColumns_->toString(); + } + return out.str(); +} + +ConnectorTableHandlePtr ParquetTableHandle::create( + const folly::dynamic& obj, + void* context) { + VELOX_NYI("ParquetTableHandle::create() not yet implemented"); +} + +} // namespace facebook::velox::cudf_velox::connector::parquet diff --git a/velox/experimental/cudf/connectors/parquet/ParquetTableHandle.h b/velox/experimental/cudf/connectors/parquet/ParquetTableHandle.h new file mode 100644 index 00000000000..b8e1fae7b88 --- /dev/null +++ b/velox/experimental/cudf/connectors/parquet/ParquetTableHandle.h @@ -0,0 +1,119 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#pragma once + +#include "velox/connectors/Connector.h" +#include "velox/core/Expressions.h" +#include "velox/expression/Expr.h" +#include "velox/type/Type.h" + +#include + +#include +#include + +namespace facebook::velox::cudf_velox::connector::parquet { + +using namespace facebook::velox::connector; + +// Parquet column handle only needs the column name (all columns are generated +// in the same way). +class ParquetColumnHandle : public ColumnHandle { + public: + explicit ParquetColumnHandle( + const std::string& name, + const TypePtr type, + const cudf::data_type cudfDataType, + std::vector children = {}) + : name_(name), + type_(type), + cudfDataType_(cudfDataType), + children_(std::move(children)) {} + + const std::string& name() const { + return name_; + } + + const TypePtr& type() const { + return type_; + } + + const cudf::data_type cudfDataType() const { + return cudfDataType_; + } + + const std::vector& children() const { + return children_; + } + + std::string toString() const; + + private: + const std::string name_; + const TypePtr type_; + const cudf::data_type cudfDataType_; + const std::vector children_; +}; + +class ParquetTableHandle : public ConnectorTableHandle { + public: + ParquetTableHandle( + std::string connectorId, + const std::string& tableName, + bool filterPushdownEnabled, + const core::TypedExprPtr& subfieldFilterExpr, + const core::TypedExprPtr& remainingFilter = nullptr, + const RowTypePtr& dataColumns = nullptr); + + const std::string& name() const override { + return tableName_; + } + + bool isFilterPushdownEnabled() const { + return filterPushdownEnabled_; + } + + const core::TypedExprPtr& subfieldFilterExpr() const { + return subfieldFilterExpr_; + } + + const core::TypedExprPtr& remainingFilter() const { + return remainingFilter_; + } + + // Schema of the table. Need this for reading TEXTFILE. + const RowTypePtr& dataColumns() const { + return dataColumns_; + } + + std::string toString() const override; + + static ConnectorTableHandlePtr create( + const folly::dynamic& obj, + void* context); + + private: + const std::string tableName_; + const bool filterPushdownEnabled_; + // This expression is used for predicate pushdown. + const core::TypedExprPtr subfieldFilterExpr_; + // This expression is used for post-scan filtering. + const core::TypedExprPtr remainingFilter_; + const RowTypePtr dataColumns_; +}; + +} // namespace facebook::velox::cudf_velox::connector::parquet diff --git a/velox/experimental/cudf/connectors/parquet/WriterOptions.h b/velox/experimental/cudf/connectors/parquet/WriterOptions.h new file mode 100644 index 00000000000..09d9527ddae --- /dev/null +++ b/velox/experimental/cudf/connectors/parquet/WriterOptions.h @@ -0,0 +1,90 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#pragma once + +#include "velox/dwio/common/Options.h" + +#include +#include +#include + +#include + +namespace facebook::velox::cudf_velox::connector::parquet { + +using namespace cudf::io; + +/** + * @brief Struct to 1:1 correspond with cudf::io::chunked_parquet_reader_options + * except sink_info and a few others which are provided to the ParquetDataSink + * from elsewhere. + */ +struct ParquetWriterOptions + : public facebook::velox::dwio::common::WriterOptions { + // Specify the level of statistics in the output file + statistics_freq statsLevel = statistics_freq::STATISTICS_ROWGROUP; + + // Parquet writer can write INT96 or TIMESTAMP_MICROS. Defaults to + // TIMESTAMPMICROS. If true then overrides any per-column setting in + // Metadata. + bool writeTimestampsAsInt96 = false; + + // Parquet writer can write timestamps as UTC + // Defaults to true because libcudf timestamps are implicitly UTC + bool writeTimestampsAsUTC = true; + + // Whether to write ARROW schema + bool writeArrowSchema = false; + + // Maximum size of each row group (unless smaller than a single page) + size_t rowGroupSizeBytes = default_row_group_size_bytes; + + // Maximum number of rows in row group (unless smaller than a single page) + cudf::size_type rowGroupSizeRows = default_row_group_size_rows; + + // Maximum size of each page (uncompressed) - Velox uses 1KB (2 x cudf limit) + size_t maxPageSizeBytes = 2 * default_max_page_size_bytes; + + // Maximum number of rows in a page + cudf::size_type maxPageSizeRows = default_max_page_size_rows; + + // Maximum size of min or max values in column index + int32_t columnIndexTruncateLength = default_column_index_truncate_length; + + // When to use dictionary encoding for data + dictionary_policy dictionaryPolicy = dictionary_policy::ADAPTIVE; + + // Maximum size of column chunk dictionary (in bytes) + size_t maxDictionarySize = default_max_dictionary_size; + + // Maximum number of rows in a page fragment + std::optional maxPageFragmentSize; + + // Optional compression statistics + std::shared_ptr compressionStats; + + // Write V2 page headers? + bool v2PageHeaders = false; + + // Encoding to use for columns + column_encoding encoding = column_encoding::PLAIN; + + // Sorting columns + std::vector sortingColumns; +}; + +} // namespace facebook::velox::cudf_velox::connector::parquet diff --git a/velox/experimental/cudf/exec/CMakeLists.txt b/velox/experimental/cudf/exec/CMakeLists.txt index c5c2711f812..7d1ca8b7557 100644 --- a/velox/experimental/cudf/exec/CMakeLists.txt +++ b/velox/experimental/cudf/exec/CMakeLists.txt @@ -15,19 +15,28 @@ add_library( velox_cudf_exec CudfConversion.cpp + CudfFilterProject.cpp CudfHashAggregation.cpp + CudfHashJoin.cpp + CudfLimit.cpp + CudfLocalPartition.cpp CudfOrderBy.cpp + DebugUtil.cpp + ExpressionEvaluator.cpp + PrestoAggregates.cpp ToCudf.cpp Utilities.cpp VeloxCudfInterop.cpp) target_link_libraries( velox_cudf_exec - cudf::cudf - arrow - velox_arrow_bridge - velox_exception - velox_common_base - velox_exec) + PUBLIC cudf::cudf + PRIVATE + arrow + velox_arrow_bridge + velox_exception + velox_common_base + velox_cudf_vector + velox_exec) target_compile_options(velox_cudf_exec PRIVATE -Wno-missing-field-initializers) diff --git a/velox/experimental/cudf/exec/CudfConversion.cpp b/velox/experimental/cudf/exec/CudfConversion.cpp index 903b01fe814..09c8f4774fd 100644 --- a/velox/experimental/cudf/exec/CudfConversion.cpp +++ b/velox/experimental/cudf/exec/CudfConversion.cpp @@ -13,8 +13,10 @@ * See the License for the specific language governing permissions and * limitations under the License. */ + #include "velox/experimental/cudf/exec/CudfConversion.h" #include "velox/experimental/cudf/exec/NvtxHelper.h" +#include "velox/experimental/cudf/exec/ToCudf.h" #include "velox/experimental/cudf/exec/Utilities.h" #include "velox/experimental/cudf/exec/VeloxCudfInterop.h" #include "velox/experimental/cudf/vector/CudfVector.h" @@ -23,7 +25,7 @@ #include "velox/exec/Operator.h" #include "velox/vector/ComplexVector.h" -#include +#include #include #include @@ -172,6 +174,11 @@ CudfToVelox::CudfToVelox( operatorId, fmt::format("[{}]", planNodeId)) {} +bool CudfToVelox::isPassthroughMode() const { + return operatorCtx_->driverCtx()->queryConfig().get( + kPassthroughMode, true); +} + void CudfToVelox::addInput(RowVectorPtr input) { // Accumulate inputs if (input->size() > 0) { @@ -181,6 +188,14 @@ void CudfToVelox::addInput(RowVectorPtr input) { } } +std::optional CudfToVelox::averageRowSize() { + if (!averageRowSize_) { + averageRowSize_ = + inputs_.front()->estimateFlatSize() / inputs_.front()->size(); + } + return averageRowSize_; +} + RowVectorPtr CudfToVelox::getOutput() { VELOX_NVTX_OPERATOR_FUNC_RANGE(); if (finished_ || inputs_.empty()) { @@ -188,16 +203,97 @@ RowVectorPtr CudfToVelox::getOutput() { return nullptr; } + // Get the target batch size + const auto targetBatchSize = outputBatchRows(averageRowSize()); auto stream = inputs_.front()->stream(); - std::unique_ptr tbl = inputs_.front()->release(); - inputs_.pop_front(); - VELOX_CHECK_NOT_NULL(tbl); - if (tbl->num_rows() == 0) { + // Process single input directly in these cases: + // 1. In passthrough mode + // 2. If we only have one input and it's smaller than or equal to the target + // batch size + if (isPassthroughMode() || + (inputs_.size() == 1 && inputs_.front()->size() <= targetBatchSize)) { + std::unique_ptr tbl = inputs_.front()->release(); + inputs_.pop_front(); + + VELOX_CHECK_NOT_NULL(tbl); + if (tbl->num_rows() == 0) { + finished_ = noMoreInput_ && inputs_.empty(); + return nullptr; + } + RowVectorPtr output = + with_arrow::toVeloxColumn(tbl->view(), pool(), "", stream); + stream.synchronize(); + finished_ = noMoreInput_ && inputs_.empty(); + output->setType(outputType_); + return output; + } + + // Calculate how many tables we need to concatenate to reach the target batch + // size and collect them in a vector + std::vector selectedInputs; + vector_size_t totalSize = 0; + + while (!inputs_.empty() && totalSize < targetBatchSize) { + auto& input = inputs_.front(); + if (totalSize + input->size() <= targetBatchSize) { + totalSize += input->size(); + selectedInputs.push_back(std::move(input)); + inputs_.pop_front(); + } else { + // If the next input would exceed targetBatchSize, + // we need to split it and only take what we need + auto cudfTableView = input->getTableView(); + auto partitions = std::vector{ + static_cast(targetBatchSize - totalSize)}; + auto tableSplits = cudf::split(cudfTableView, partitions); + + // Create new CudfVector from the first part + auto firstPart = std::make_unique(tableSplits[0], stream); + auto firstPartSize = firstPart->num_rows(); + auto firstPartVector = std::make_shared( + pool(), input->type(), firstPartSize, std::move(firstPart), stream); + + // Create new CudfVector from the second part + auto secondPart = std::make_unique(tableSplits[1], stream); + auto secondPartSize = secondPart->num_rows(); + auto secondPartVector = std::make_shared( + pool(), input->type(), secondPartSize, std::move(secondPart), stream); + + // Replace the original input with the second part + input = std::move(secondPartVector); + + // Add the first part to selectedInputs + selectedInputs.push_back(std::move(firstPartVector)); + totalSize += firstPartSize; + break; + } + } + + finished_ = noMoreInput_ && inputs_.empty(); + + // If we have no inputs to process, return nullptr + if (selectedInputs.empty()) { + return nullptr; + } + + // Concatenate the selected tables on the GPU + std::unique_ptr resultTable; + if (selectedInputs.size() == 1) { + resultTable = selectedInputs[0]->release(); + } else { + resultTable = getConcatenatedTable(selectedInputs, stream); + } + + // Convert the concatenated table to a RowVector + const auto size = resultTable->num_rows(); + VELOX_CHECK_NOT_NULL(resultTable); + if (size == 0) { return nullptr; } + RowVectorPtr output = - with_arrow::toVeloxColumn(tbl->view(), pool(), "", stream); + with_arrow::toVeloxColumn(resultTable->view(), pool(), "", stream); stream.synchronize(); finished_ = noMoreInput_ && inputs_.empty(); output->setType(outputType_); diff --git a/velox/experimental/cudf/exec/CudfConversion.h b/velox/experimental/cudf/exec/CudfConversion.h index 16ca33d786c..879f21d512b 100644 --- a/velox/experimental/cudf/exec/CudfConversion.h +++ b/velox/experimental/cudf/exec/CudfConversion.h @@ -23,8 +23,6 @@ #include "velox/exec/Operator.h" #include "velox/vector/ComplexVector.h" -#include - #include #include #include @@ -68,6 +66,9 @@ class CudfFromVelox : public exec::Operator, public NvtxHelper { class CudfToVelox : public exec::Operator, public NvtxHelper { public: + static constexpr const char* kPassthroughMode = + "velox.cudf.to_velox.passthrough_mode"; + CudfToVelox( int32_t operatorId, RowTypePtr outputType, @@ -93,6 +94,9 @@ class CudfToVelox : public exec::Operator, public NvtxHelper { void close() override; private: + bool isPassthroughMode() const; + std::optional averageRowSize(); + std::optional averageRowSize_; std::deque inputs_; bool finished_ = false; }; diff --git a/velox/experimental/cudf/exec/CudfFilterProject.cpp b/velox/experimental/cudf/exec/CudfFilterProject.cpp new file mode 100644 index 00000000000..0959d78e800 --- /dev/null +++ b/velox/experimental/cudf/exec/CudfFilterProject.cpp @@ -0,0 +1,206 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#include "velox/experimental/cudf/exec/CudfFilterProject.h" +#include "velox/experimental/cudf/exec/ToCudf.h" +#include "velox/experimental/cudf/vector/CudfVector.h" + +#include "velox/expression/Expr.h" + +#include +#include +#include + +#include + +namespace facebook::velox::cudf_velox { + +namespace { + +void debugPrintTree( + const std::shared_ptr& expr, + int indent = 0) { + std::cout << std::string(indent, ' ') << expr->name() << std::endl; + for (auto& input : expr->inputs()) { + debugPrintTree(input, indent + 2); + } +} +} // namespace + +CudfFilterProject::CudfFilterProject( + int32_t operatorId, + velox::exec::DriverCtx* driverCtx, + const velox::exec::FilterProject::Export& info, + std::vector identityProjections, + const std::shared_ptr& filter, + const std::shared_ptr& project) + : Operator( + driverCtx, + project ? project->outputType() : filter->outputType(), + operatorId, + project ? project->id() : filter->id(), + "CudfFilterProject"), + NvtxHelper( + nvtx3::rgb{220, 20, 60}, // Crimson + operatorId, + fmt::format("[{}]", project ? project->id() : filter->id())), + hasFilter_(info.hasFilter), + project_(project), + filter_(filter) { + resultProjections_ = *(info.resultProjections); + identityProjections_ = std::move(identityProjections); + const auto inputType = project_ ? project_->sources()[0]->outputType() + : filter_->sources()[0]->outputType(); + + // convert to AST + if (cudfDebugEnabled()) { + int i = 0; + for (auto expr : info.exprs->exprs()) { + std::cout << "expr[" << i++ << "] " << expr->toString() << std::endl; + debugPrintTree(expr); + } + } + std::vector> projectExprs; + if (hasFilter_) { + // First expr is Filter, rest are Project + filterEvaluator_ = ExpressionEvaluator({info.exprs->exprs()[0]}, inputType); + projectExprs = {info.exprs->exprs().begin() + 1, info.exprs->exprs().end()}; + } + projectEvaluator_ = ExpressionEvaluator( + hasFilter_ ? projectExprs : info.exprs->exprs(), inputType); +} + +void CudfFilterProject::addInput(RowVectorPtr input) { + input_ = std::move(input); +} + +RowVectorPtr CudfFilterProject::getOutput() { + VELOX_NVTX_OPERATOR_FUNC_RANGE(); + + if (allInputProcessed()) { + return nullptr; + } + if (input_->size() == 0) { + input_.reset(); + return nullptr; + } + + auto cudfInput = std::dynamic_pointer_cast(input_); + VELOX_CHECK_NOT_NULL(cudfInput); + auto stream = cudfInput->stream(); + auto inputTableColumns = cudfInput->release()->release(); + + if (hasFilter_) { + filter(inputTableColumns, stream); + } + auto outputColumns = project(inputTableColumns, stream); + + auto outputTable = std::make_unique(std::move(outputColumns)); + stream.synchronize(); + auto const numColumns = outputTable->num_columns(); + auto const size = outputTable->num_rows(); + if (cudfDebugEnabled()) { + std::cout << "cudfProject Output: " << size << " rows, " << numColumns + << " columns " << std::endl; + } + + auto cudfOutput = std::make_shared( + input_->pool(), outputType_, size, std::move(outputTable), stream); + input_.reset(); + if (numColumns == 0 or size == 0) { + return nullptr; + } + return cudfOutput; +} + +void CudfFilterProject::filter( + std::vector>& inputTableColumns, + rmm::cuda_stream_view stream) { + // Evaluate the Filter + auto filterColumns = filterEvaluator_.compute( + inputTableColumns, stream, cudf::get_current_device_resource_ref()); + auto filterColumn = filterColumns[0]->view(); + // is all true in filter_column + auto isAllTrue = cudf::reduce( + filterColumn, + *cudf::make_all_aggregation(), + cudf::data_type(cudf::type_id::BOOL8), + stream, + cudf::get_current_device_resource_ref()); + using ScalarType = cudf::scalar_type_t; + auto result = static_cast(isAllTrue.get()); + // If filter is not all true, apply the filter + if (!(result->is_valid(stream) && result->value(stream))) { + // Apply the Filter + auto filterTable = + std::make_unique(std::move(inputTableColumns)); + auto filteredTable = + cudf::apply_boolean_mask(*filterTable, filterColumn, stream); + inputTableColumns = filteredTable->release(); + } +} + +std::vector> CudfFilterProject::project( + std::vector>& inputTableColumns, + rmm::cuda_stream_view stream) { + auto columns = projectEvaluator_.compute( + inputTableColumns, stream, cudf::get_current_device_resource_ref()); + + // Rearrange columns to match outputType_ + std::vector> outputColumns(outputType_->size()); + // computed resultProjections + for (int i = 0; i < resultProjections_.size(); i++) { + VELOX_CHECK_NOT_NULL(columns[i]); + outputColumns[resultProjections_[i].outputChannel] = std::move(columns[i]); + } + + // Count occurrences of each inputChannel, and move columns if they occur only + // once + std::unordered_map inputChannelCount; + for (const auto& identity : identityProjections_) { + inputChannelCount[identity.inputChannel]++; + } + + // identityProjections (input to output copy) + for (auto const& identity : identityProjections_) { + VELOX_CHECK_NOT_NULL(inputTableColumns[identity.inputChannel]); + if (inputChannelCount[identity.inputChannel] == 1) { + // Move the column if it occurs only once + outputColumns[identity.outputChannel] = + std::move(inputTableColumns[identity.inputChannel]); + } else { + // Otherwise, copy the column and decrement the count + outputColumns[identity.outputChannel] = std::make_unique( + *inputTableColumns[identity.inputChannel], + stream, + cudf::get_current_device_resource_ref()); + } + VELOX_CHECK_GT(inputChannelCount[identity.inputChannel], 0); + inputChannelCount[identity.inputChannel]--; + } + + return outputColumns; +} + +bool CudfFilterProject::allInputProcessed() { + return !input_; +} + +bool CudfFilterProject::isFinished() { + return noMoreInput_ && allInputProcessed(); +} + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/exec/CudfFilterProject.h b/velox/experimental/cudf/exec/CudfFilterProject.h new file mode 100644 index 00000000000..23853ddb47e --- /dev/null +++ b/velox/experimental/cudf/exec/CudfFilterProject.h @@ -0,0 +1,84 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#pragma once + +#include "velox/experimental/cudf/exec/ExpressionEvaluator.h" +#include "velox/experimental/cudf/exec/NvtxHelper.h" + +#include "velox/core/Expressions.h" +#include "velox/core/PlanNode.h" +#include "velox/exec/FilterProject.h" +#include "velox/exec/Operator.h" +#include "velox/vector/ComplexVector.h" + +namespace facebook::velox::cudf_velox { + +// TODO: Does not support Filter yet. +class CudfFilterProject : public exec::Operator, public NvtxHelper { + public: + CudfFilterProject( + int32_t operatorId, + velox::exec::DriverCtx* driverCtx, + const velox::exec::FilterProject::Export& info, + std::vector identityProjections, + const std::shared_ptr& filter, + const std::shared_ptr& project); + + bool needsInput() const override { + return !input_; + } + + void addInput(RowVectorPtr input) override; + + RowVectorPtr getOutput() override; + + void filter( + std::vector>& inputTableColumns, + rmm::cuda_stream_view stream); + + std::vector> project( + std::vector>& inputTableColumns, + rmm::cuda_stream_view stream); + + exec::BlockingReason isBlocked(ContinueFuture* /*future*/) override { + return exec::BlockingReason::kNotBlocked; + } + + bool isFinished() override; + + void close() override { + Operator::close(); + projectEvaluator_.close(); + filterEvaluator_.close(); + } + + private: + bool allInputProcessed(); + // If true exprs_[0] is a filter and the other expressions are projections + const bool hasFilter_{false}; + // Cached filter and project node for lazy initialization. After + // initialization, they will be reset, and initialized_ will be set to true. + std::shared_ptr project_; + std::shared_ptr filter_; + ExpressionEvaluator projectEvaluator_; + ExpressionEvaluator filterEvaluator_; + + std::vector resultProjections_; + std::vector identityProjections_; +}; + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/exec/CudfHashAggregation.cpp b/velox/experimental/cudf/exec/CudfHashAggregation.cpp index dfc03a1b258..db3aadd88b1 100644 --- a/velox/experimental/cudf/exec/CudfHashAggregation.cpp +++ b/velox/experimental/cudf/exec/CudfHashAggregation.cpp @@ -30,6 +30,8 @@ #include #include +#include + namespace { using namespace facebook::velox; @@ -386,32 +388,37 @@ struct MeanAggregator : cudf_velox::CudfHashAggregation::Aggregator { uint32_t countIdx_; }; -std::unique_ptr createAggregator( - core::AggregationNode::Step step, +static const std::unordered_map + companionStep = { + {"_partial", core::AggregationNode::Step::kPartial}, + {"_merge", core::AggregationNode::Step::kIntermediate}, + {"_merge_extract", core::AggregationNode::Step::kFinal}}; + +/// \brief Convert companion function to step for the aggregation function +/// +/// Companion functions are functions that are registered in velox along with +/// their main aggregation functions. These are designed to always function +/// with a fixed `step`. This is to allow spark style planNodes where `step` is +/// the property of the aggregation function rather than the planNode. +/// Companion functions allow us to override the planNode's step and use +/// aggregations of different steps in the same planNode +core::AggregationNode::Step getCompanionStep( std::string const& kind, - uint32_t inputIndex, - VectorPtr constant, - bool isGlobal) { - // Companion function may be count_merge_extract or count_partial or others, - // so use this to map - if (kind.rfind("sum", 0) == 0) { - return std::make_unique( - step, inputIndex, constant, isGlobal); - } else if (kind.rfind("count", 0) == 0) { - return std::make_unique( - step, inputIndex, constant, isGlobal); - } else if (kind.rfind("min", 0) == 0) { - return std::make_unique( - step, inputIndex, constant, isGlobal); - } else if (kind.rfind("max", 0) == 0) { - return std::make_unique( - step, inputIndex, constant, isGlobal); - } else if (kind.rfind("avg", 0) == 0) { - return std::make_unique( - step, inputIndex, constant, isGlobal); - } else { - VELOX_NYI("Aggregation not yet supported"); + core::AggregationNode::Step step) { + for (const auto& [k, v] : companionStep) { + if (folly::StringPiece(kind).endsWith(k)) { + step = v; + break; + } } + return step; +} + +bool hasFinalAggs( + std::vector const& aggregates) { + return std::any_of(aggregates.begin(), aggregates.end(), [](auto const& agg) { + return folly::StringPiece(agg.call->name()).endsWith("_merge_extract"); + }); } auto toAggregators( @@ -453,8 +460,9 @@ auto toAggregators( auto const kind = aggregate.call->name(); auto const inputIndex = aggInputs[0]; auto const constant = aggConstants.empty() ? nullptr : aggConstants[0]; - aggregators.push_back( - createAggregator(step, kind, inputIndex, constant, isGlobal)); + auto const companionStep = getCompanionStep(kind, step); + aggregators.push_back(facebook::velox::cudf_velox::createAggregator( + kind, companionStep, inputIndex, constant, isGlobal)); } return aggregators; } @@ -475,8 +483,8 @@ auto toIntermediateAggregators( auto const inputIndex = aggregationNode.groupingKeys().size() + i; auto const kind = aggregate.call->name(); auto const constant = nullptr; - aggregators.push_back( - createAggregator(step, kind, inputIndex, constant, isGlobal)); + aggregators.push_back(facebook::velox::cudf_velox::createAggregator( + kind, step, inputIndex, constant, isGlobal)); } return aggregators; } @@ -505,7 +513,9 @@ CudfHashAggregation::CudfHashAggregation( operatorId, fmt::format("[{}]", aggregationNode->id())), aggregationNode_(aggregationNode), - isPartialOutput_(exec::isPartialOutput(aggregationNode->step())), + isPartialOutput_( + exec::isPartialOutput(aggregationNode->step()) && + !hasFinalAggs(aggregationNode->aggregates())), isGlobal_(aggregationNode->groupingKeys().empty()), isDistinct_(!isGlobal_ && aggregationNode->aggregates().empty()), maxPartialAggregationMemoryUsage_( @@ -855,4 +865,150 @@ bool CudfHashAggregation::isFinished() { return finished_; } +std::unique_ptr createAggregator( + const std::string& kind, + core::AggregationNode::Step step, + uint32_t inputIndex, + VectorPtr constant, + bool isGlobal) { + // Ensure basic cudf aggregators are registered + static std::once_flag registrationFlag; + std::call_once(registrationFlag, []() { + facebook::velox::cudf_velox::registerCudfAggregators( + false /* withCompanionFunctions */, false /* overwrite */); + }); + + if (auto entry = facebook::velox::cudf_velox::getAggregatorEntry(kind)) { + return entry->factory(step, inputIndex, constant, isGlobal); + } + + VELOX_NYI("Aggregation not yet supported: {}", kind); +} + +AggregatorMap& aggregators() { + static AggregatorMap aggregators; + return aggregators; +} + +const AggregatorEntry* FOLLY_NULLABLE +getAggregatorEntry(const std::string& name) { + return aggregators().withRLock( + [&](const auto& aggregatorsMap) -> const AggregatorEntry* { + auto it = aggregatorsMap.find(name); + if (it != aggregatorsMap.end()) { + return &it->second; + } + return nullptr; + }); +} + +bool registerAggregator( + const std::string& name, + const AggregatorFactory& factory, + bool overwrite) { + if (overwrite) { + aggregators().withWLock( + [&](auto& aggregatorsMap) { aggregatorsMap[name] = {factory}; }); + return true; + } else { + return aggregators().withWLock([&](auto& aggregatorsMap) { + auto [_, inserted] = aggregatorsMap.insert({name, {factory}}); + return inserted; + }); + } +} + +// Registration functions for CUDF aggregators +template +void registerAggregatorImpl( + const std::string& name, + bool withCompanionFunctions, + bool overwrite) { + registerAggregator( + name, + [](core::AggregationNode::Step step, + uint32_t inputIndex, + VectorPtr constant, + bool isGlobal) -> std::unique_ptr { + return std::make_unique( + step, inputIndex, constant, isGlobal); + }, + overwrite); + if (withCompanionFunctions) { + registerAggregator( + name + "_partial", + [](core::AggregationNode::Step, + uint32_t inputIndex, + VectorPtr constant, + bool isGlobal) -> std::unique_ptr { + return std::make_unique( + core::AggregationNode::Step::kPartial, + inputIndex, + constant, + isGlobal); + }, + overwrite); + registerAggregator( + name + "_merge", + [](core::AggregationNode::Step, + uint32_t inputIndex, + VectorPtr constant, + bool isGlobal) -> std::unique_ptr { + return std::make_unique( + core::AggregationNode::Step::kIntermediate, + inputIndex, + constant, + isGlobal); + }, + overwrite); + registerAggregator( + name + "_merge_extract", + [](core::AggregationNode::Step, + uint32_t inputIndex, + VectorPtr constant, + bool isGlobal) -> std::unique_ptr { + return std::make_unique( + core::AggregationNode::Step::kFinal, + inputIndex, + constant, + isGlobal); + }, + overwrite); + } +} + +void registerSumAggregator(bool withCompanionFunctions, bool overwrite) { + registerAggregatorImpl( + "sum", withCompanionFunctions, overwrite); +} + +void registerCountAggregator(bool withCompanionFunctions, bool overwrite) { + registerAggregatorImpl( + "count", withCompanionFunctions, overwrite); +} + +void registerMinAggregator(bool withCompanionFunctions, bool overwrite) { + registerAggregatorImpl( + "min", withCompanionFunctions, overwrite); +} + +void registerMaxAggregator(bool withCompanionFunctions, bool overwrite) { + registerAggregatorImpl( + "max", withCompanionFunctions, overwrite); +} + +void registerAvgAggregator(bool withCompanionFunctions, bool overwrite) { + registerAggregatorImpl( + "avg", withCompanionFunctions, overwrite); +} + +// Register all CUDF aggregators +void registerCudfAggregators(bool withCompanionFunctions, bool overwrite) { + registerSumAggregator(withCompanionFunctions, overwrite); + registerCountAggregator(withCompanionFunctions, overwrite); + registerMinAggregator(withCompanionFunctions, overwrite); + registerMaxAggregator(withCompanionFunctions, overwrite); + registerAvgAggregator(withCompanionFunctions, overwrite); +} + } // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/exec/CudfHashAggregation.h b/velox/experimental/cudf/exec/CudfHashAggregation.h index 3b828e18e3c..f1129f38c54 100644 --- a/velox/experimental/cudf/exec/CudfHashAggregation.h +++ b/velox/experimental/cudf/exec/CudfHashAggregation.h @@ -147,4 +147,54 @@ class CudfHashAggregation : public exec::Operator, public NvtxHelper { CudfVectorPtr partialOutput_; }; +using AggregatorFactory = + std::function( + core::AggregationNode::Step step, + uint32_t inputIndex, + VectorPtr constant, + bool isGlobal)>; + +struct AggregatorEntry { + AggregatorFactory factory; +}; + +using AggregatorMap = + folly::Synchronized>; + +AggregatorMap& aggregators(); + +const AggregatorEntry* FOLLY_NULLABLE +getAggregatorEntry(const std::string& name); + +/// Register an aggregator function with the specified name and factory. +/// When function with `name` already exists, if overwrite is true, existing +/// registration will be replaced. Otherwise, return false without overwriting. +bool registerAggregator( + const std::string& name, + const AggregatorFactory& factory, + bool overwrite = false); + +/// Creates an aggregator instance using the registered factory. +/// Returns nullptr if no factory is registered for the given name. +std::unique_ptr createAggregator( + const std::string& kind, + core::AggregationNode::Step step, + uint32_t inputIndex, + VectorPtr constant, + bool isGlobal); + +/// Registration functions for CUDF aggregators +void registerSumAggregator(bool withCompanionFunctions, bool overwrite = false); +void registerCountAggregator( + bool withCompanionFunctions, + bool overwrite = false); +void registerMinAggregator(bool withCompanionFunctions, bool overwrite = false); +void registerMaxAggregator(bool withCompanionFunctions, bool overwrite = false); +void registerAvgAggregator(bool withCompanionFunctions, bool overwrite = false); + +/// Register all CUDF aggregators +void registerCudfAggregators( + bool withCompanionFunctions, + bool overwrite = false); + } // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/exec/CudfHashJoin.cpp b/velox/experimental/cudf/exec/CudfHashJoin.cpp new file mode 100644 index 00000000000..f4e59a16873 --- /dev/null +++ b/velox/experimental/cudf/exec/CudfHashJoin.cpp @@ -0,0 +1,739 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#include "velox/experimental/cudf/exec/CudfHashJoin.h" +#include "velox/experimental/cudf/exec/ExpressionEvaluator.h" +#include "velox/experimental/cudf/exec/ToCudf.h" +#include "velox/experimental/cudf/exec/Utilities.h" +#include "velox/experimental/cudf/exec/VeloxCudfInterop.h" + +#include "velox/exec/Task.h" + +#include +#include +#include +#include +#include + +#include + +namespace facebook::velox::cudf_velox { + +void CudfHashJoinBridge::setHashTable( + std::optional hashObject) { + if (cudfDebugEnabled()) { + std::cout << "Calling CudfHashJoinBridge::setHashTable" << std::endl; + } + std::vector promises; + { + std::lock_guard l(mutex_); + VELOX_CHECK( + !hashObject_.has_value(), + "CudfHashJoinBridge already has a hash table"); + hashObject_ = std::move(hashObject); + promises = std::move(promises_); + } + notify(std::move(promises)); +} + +std::optional CudfHashJoinBridge::hashOrFuture( + ContinueFuture* future) { + if (cudfDebugEnabled()) { + std::cout << "Calling CudfHashJoinBridge::hashOrFuture" << std::endl; + } + std::lock_guard l(mutex_); + if (hashObject_.has_value()) { + return hashObject_; + } + if (cudfDebugEnabled()) { + std::cout << "Calling CudfHashJoinBridge::hashOrFuture constructing promise" + << std::endl; + } + promises_.emplace_back("CudfHashJoinBridge::hashOrFuture"); + if (cudfDebugEnabled()) { + std::cout << "Calling CudfHashJoinBridge::hashOrFuture getSemiFuture" + << std::endl; + } + *future = promises_.back().getSemiFuture(); + if (cudfDebugEnabled()) { + std::cout << "Calling CudfHashJoinBridge::hashOrFuture returning nullopt" + << std::endl; + } + return std::nullopt; +} + +CudfHashJoinBuild::CudfHashJoinBuild( + int32_t operatorId, + exec::DriverCtx* driverCtx, + std::shared_ptr joinNode) + // TODO check outputType should be set or not? + : exec::Operator( + driverCtx, + nullptr, // joinNode->sources(), + operatorId, + joinNode->id(), + "CudfHashJoinBuild"), + NvtxHelper( + nvtx3::rgb{65, 105, 225}, // Royal Blue + operatorId, + fmt::format("[{}]", joinNode->id())), + joinNode_(joinNode) { + if (cudfDebugEnabled()) { + std::cout << "CudfHashJoinBuild constructor" << std::endl; + } +} + +void CudfHashJoinBuild::addInput(RowVectorPtr input) { + if (cudfDebugEnabled()) { + std::cout << "Calling CudfHashJoinBuild::addInput" << std::endl; + } + // Queue inputs, process all at once. + if (input->size() > 0) { + auto cudfInput = std::dynamic_pointer_cast(input); + VELOX_CHECK_NOT_NULL(cudfInput); + // Count nulls in join key columns + cudf::size_type null_count{}; + std::tie(std::ignore, null_count) = cudf::bitmask_and( + cudfInput->getTableView(), + cudfInput->stream(), + cudf::get_current_device_resource_ref()); + { + // Update statistics for null keys in join operator. + auto lockedStats = stats_.wlock(); + lockedStats->numNullKeys += null_count; + } + inputs_.push_back(std::move(cudfInput)); + } +} + +bool CudfHashJoinBuild::needsInput() const { + if (cudfDebugEnabled()) { + std::cout << "Calling CudfHashJoinBuild::needsInput" << std::endl; + } + return !noMoreInput_; +} + +RowVectorPtr CudfHashJoinBuild::getOutput() { + return nullptr; +} + +void CudfHashJoinBuild::noMoreInput() { + if (cudfDebugEnabled()) { + std::cout << "Calling CudfHashJoinBuild::noMoreInput" << std::endl; + } + VELOX_NVTX_OPERATOR_FUNC_RANGE(); + Operator::noMoreInput(); + std::vector promises; + std::vector> peers; + // Only last driver collects all answers + if (!operatorCtx_->task()->allPeersFinished( + planNodeId(), operatorCtx_->driver(), &future_, promises, peers)) { + return; + } + // Collect results from peers + for (auto& peer : peers) { + auto op = peer->findOperator(planNodeId()); + auto* build = dynamic_cast(op); + VELOX_CHECK_NOT_NULL(build); + inputs_.insert(inputs_.end(), build->inputs_.begin(), build->inputs_.end()); + } + + SCOPE_EXIT { + // Realize the promises so that the other Drivers (which were not + // the last to finish) can continue from the barrier and finish. + peers.clear(); + for (auto& promise : promises) { + promise.setValue(); + } + }; + + auto stream = cudfGlobalStreamPool().get_stream(); + std::unique_ptr tbl; + if (inputs_.size() == 0) { + auto emptyRowVector = RowVector::createEmpty( + joinNode_->sources()[1]->outputType(), operatorCtx_->pool()); + tbl = facebook::velox::cudf_velox::with_arrow::toCudfTable( + emptyRowVector, operatorCtx_->pool(), stream); + } else { + tbl = getConcatenatedTable(inputs_, stream); + } + + // Release input data after synchronizing + stream.synchronize(); + inputs_.clear(); + + VELOX_CHECK_NOT_NULL(tbl); + if (cudfDebugEnabled()) { + std::cout << "Build table number of columns: " << tbl->num_columns() + << std::endl; + std::cout << "Build table number of rows: " << tbl->num_rows() << std::endl; + } + + auto buildType = joinNode_->sources()[1]->outputType(); + auto rightKeys = joinNode_->rightKeys(); + + auto buildKeyIndices = std::vector(rightKeys.size()); + for (size_t i = 0; i < buildKeyIndices.size(); i++) { + buildKeyIndices[i] = static_cast( + buildType->getChildIdx(rightKeys[i]->name())); + } + + // Only need to construct hash_join object if it's an inner join or left join + // and doesn't have a filter. All other cases use a standalone function in + // cudf + bool buildHashJoin = (joinNode_->isInnerJoin() || joinNode_->isLeftJoin()) && + !joinNode_->filter(); + auto hashObject = (buildHashJoin) ? std::make_shared( + tbl->view().select(buildKeyIndices), + cudf::null_equality::UNEQUAL, + stream) + : nullptr; + if (buildHashJoin) { + VELOX_CHECK_NOT_NULL(hashObject); + } + + if (cudfDebugEnabled()) { + if (hashObject != nullptr) { + printf("hashObject is not nullptr %p\n", hashObject.get()); + } else { + printf("hashObject is *** nullptr\n"); + } + } + + // set hash table to CudfHashJoinBridge + auto joinBridge = operatorCtx_->task()->getCustomJoinBridge( + operatorCtx_->driverCtx()->splitGroupId, planNodeId()); + auto cudfHashJoinBridge = + std::dynamic_pointer_cast(joinBridge); + cudfHashJoinBridge->setHashTable(std::make_optional( + std::make_pair(std::shared_ptr(std::move(tbl)), std::move(hashObject)))); +} + +exec::BlockingReason CudfHashJoinBuild::isBlocked(ContinueFuture* future) { + if (!future_.valid()) { + return exec::BlockingReason::kNotBlocked; + } + *future = std::move(future_); + return exec::BlockingReason::kWaitForJoinBuild; +} + +bool CudfHashJoinBuild::isFinished() { + return !future_.valid() && noMoreInput_; +} + +CudfHashJoinProbe::CudfHashJoinProbe( + int32_t operatorId, + exec::DriverCtx* driverCtx, + std::shared_ptr joinNode) + : exec::Operator( + driverCtx, + joinNode->outputType(), + operatorId, + joinNode->id(), + "CudfHashJoinProbe"), + NvtxHelper( + nvtx3::rgb{0, 128, 128}, // Teal + operatorId, + fmt::format("[{}]", joinNode->id())), + joinNode_(joinNode) { + if (cudfDebugEnabled()) { + std::cout << "CudfHashJoinProbe constructor" << std::endl; + } + auto probeType = joinNode_->sources()[0]->outputType(); + auto buildType = joinNode_->sources()[1]->outputType(); + auto const& leftKeys = joinNode_->leftKeys(); // probe keys + auto const& rightKeys = joinNode_->rightKeys(); // build keys + + if (cudfDebugEnabled()) { + for (int i = 0; i < probeType->names().size(); i++) { + std::cout << "Left column " << i << ": " << probeType->names()[i] + << std::endl; + } + + for (int i = 0; i < buildType->names().size(); i++) { + std::cout << "Right column " << i << ": " << buildType->names()[i] + << std::endl; + } + + for (int i = 0; i < leftKeys.size(); i++) { + std::cout << "Left key " << i << ": " << leftKeys[i]->name() << " " + << leftKeys[i]->type()->kind() << std::endl; + } + + for (int i = 0; i < rightKeys.size(); i++) { + std::cout << "Right key " << i << ": " << rightKeys[i]->name() << " " + << rightKeys[i]->type()->kind() << std::endl; + } + } + + auto const probeTableNumColumns = probeType->size(); + leftKeyIndices_ = std::vector(leftKeys.size()); + for (size_t i = 0; i < leftKeyIndices_.size(); i++) { + leftKeyIndices_[i] = static_cast( + probeType->getChildIdx(leftKeys[i]->name())); + VELOX_CHECK_LT(leftKeyIndices_[i], probeTableNumColumns); + } + auto const buildTableNumColumns = buildType->size(); + rightKeyIndices_ = std::vector(rightKeys.size()); + for (size_t i = 0; i < rightKeyIndices_.size(); i++) { + rightKeyIndices_[i] = static_cast( + buildType->getChildIdx(rightKeys[i]->name())); + VELOX_CHECK_LT(rightKeyIndices_[i], buildTableNumColumns); + } + + auto outputType = joinNode_->outputType(); + leftColumnIndicesToGather_ = std::vector(); + rightColumnIndicesToGather_ = std::vector(); + leftColumnOutputIndices_ = std::vector(); + rightColumnOutputIndices_ = std::vector(); + for (int i = 0; i < outputType->names().size(); i++) { + auto const outputName = outputType->names()[i]; + if (cudfDebugEnabled()) { + std::cout << "Output column " << i << ": " << outputName << std::endl; + } + auto channel = probeType->getChildIdxIfExists(outputName); + if (channel.has_value()) { + leftColumnIndicesToGather_.push_back( + static_cast(channel.value())); + leftColumnOutputIndices_.push_back(i); + continue; + } + channel = buildType->getChildIdxIfExists(outputName); + if (channel.has_value()) { + rightColumnIndicesToGather_.push_back( + static_cast(channel.value())); + rightColumnOutputIndices_.push_back(i); + continue; + } + VELOX_FAIL( + "Join field {} not in probe or build input", outputType->children()[i]); + } + + if (cudfDebugEnabled()) { + for (int i = 0; i < leftColumnIndicesToGather_.size(); i++) { + std::cout << "Left index to gather " << i << ": " + << leftColumnIndicesToGather_[i] << std::endl; + } + + for (int i = 0; i < rightColumnIndicesToGather_.size(); i++) { + std::cout << "Right index to gather " << i << ": " + << rightColumnIndicesToGather_[i] << std::endl; + } + } + + // Setup filter in case it exists + if (joinNode_->filter()) { + // simplify expression + exec::ExprSet exprs({joinNode_->filter()}, operatorCtx_->execCtx()); + VELOX_CHECK_EQ(exprs.exprs().size(), 1); + + // We don't need to get tables that contain conditional comparison columns + // We'll pass the entire table. The ast will handle finding the required + // columns. This is required because we build the ast with whole row schema + // and the column locations in that schema translate to column locations + // in whole tables + + // create ast tree + std::vector rightPrecomputeInstructions; + std::vector leftPrecomputeInstructions; + if (joinNode_->isRightJoin() || joinNode_->isRightSemiFilterJoin()) { + createAstTree( + exprs.exprs()[0], + tree_, + scalars_, + buildType, + probeType, + rightPrecomputeInstructions, + leftPrecomputeInstructions); + } else { + createAstTree( + exprs.exprs()[0], + tree_, + scalars_, + probeType, + buildType, + leftPrecomputeInstructions, + rightPrecomputeInstructions); + } + if (leftPrecomputeInstructions.size() > 0 || + rightPrecomputeInstructions.size() > 0) { + VELOX_NYI("Filters that require precomputation are not yet supported"); + } + } +} + +bool CudfHashJoinProbe::needsInput() const { + return !finished_ && input_ == nullptr; +} + +void CudfHashJoinProbe::addInput(RowVectorPtr input) { + if (skipInput_) { + VELOX_CHECK_NULL(input_); + return; + } + auto cudfInput = std::dynamic_pointer_cast(input); + VELOX_CHECK_NOT_NULL(cudfInput); + // Count nulls in join key columns + cudf::size_type null_count{}; + std::tie(std::ignore, null_count) = cudf::bitmask_and( + cudfInput->getTableView(), + cudfInput->stream(), + cudf::get_current_device_resource_ref()); + { + // Update statistics for null keys in join operator. + auto lockedStats = stats_.wlock(); + lockedStats->numNullKeys += null_count; + } + input_ = std::move(input); +} + +RowVectorPtr CudfHashJoinProbe::getOutput() { + if (cudfDebugEnabled()) { + std::cout << "Calling CudfHashJoinProbe::getOutput" << std::endl; + } + VELOX_NVTX_OPERATOR_FUNC_RANGE(); + + if (finished_ or !hashObject_.has_value()) { + return nullptr; + } + + if (!input_) { + return nullptr; + } + auto cudfInput = std::dynamic_pointer_cast(input_); + VELOX_CHECK_NOT_NULL(cudfInput); + auto stream = cudfInput->stream(); + auto leftTable = cudfInput->release(); // probe table + if (cudfDebugEnabled()) { + std::cout << "Probe table number of columns: " << leftTable->num_columns() + << std::endl; + std::cout << "Probe table number of rows: " << leftTable->num_rows() + << std::endl; + } + + // TODO pass the input pool !!! + // TODO: We should probably subset columns before calling to_cudf_table? + // Maybe that isn't a problem if we fuse operators together. + auto& rightTable = hashObject_.value().first; + auto& hb = hashObject_.value().second; + VELOX_CHECK_NOT_NULL(rightTable); + if (cudfDebugEnabled()) { + if (rightTable != nullptr) + printf( + "right_table is not nullptr %p hasValue(%d)\n", + rightTable.get(), + hashObject_.has_value()); + if (hb != nullptr) + printf( + "hb is not nullptr %p hasValue(%d)\n", + hb.get(), + hashObject_.has_value()); + } + + // Special case for null-aware anti join where + // build table is not empty, no nulls, and probe table has nulls + if (joinNode_->isNullAware() and !joinNode_->filter()) { + auto const rightTableHasNulls = + cudf::has_nulls(rightTable->view().select(rightKeyIndices_)); + auto const leftTableHasNulls = + cudf::has_nulls(leftTable->view().select(leftKeyIndices_)); + if (rightTable->num_rows() > 0 and !rightTableHasNulls and + leftTableHasNulls) { + // drop nulls on probe table + leftTable = cudf::drop_nulls( + leftTable->view(), + leftKeyIndices_, + stream, + cudf::get_current_device_resource_ref()); + } + } + + std::unique_ptr> leftJoinIndices; + std::unique_ptr> rightJoinIndices; + + auto leftTableView = leftTable->view(); + auto rightTableView = rightTable->view(); + + if (joinNode_->isInnerJoin()) { + // left = probe, right = build + if (joinNode_->filter()) { + std::tie(leftJoinIndices, rightJoinIndices) = cudf::mixed_inner_join( + leftTableView.select(leftKeyIndices_), + rightTableView.select(rightKeyIndices_), + leftTableView, + rightTableView, + tree_.back(), + cudf::null_equality::UNEQUAL, + std::nullopt, + stream); + } else { + VELOX_CHECK_NOT_NULL(hb); + std::tie(leftJoinIndices, rightJoinIndices) = hb->inner_join( + leftTableView.select(leftKeyIndices_), std::nullopt, stream); + } + } else if (joinNode_->isLeftJoin()) { + if (joinNode_->filter()) { + std::tie(leftJoinIndices, rightJoinIndices) = cudf::mixed_left_join( + leftTableView.select(leftKeyIndices_), + rightTableView.select(rightKeyIndices_), + leftTableView, + rightTableView, + tree_.back(), + cudf::null_equality::UNEQUAL, + std::nullopt, + stream); + } else { + VELOX_CHECK_NOT_NULL(hb); + std::tie(leftJoinIndices, rightJoinIndices) = hb->left_join( + leftTableView.select(leftKeyIndices_), std::nullopt, stream); + } + } else if (joinNode_->isRightJoin()) { + if (joinNode_->filter()) { + std::tie(rightJoinIndices, leftJoinIndices) = cudf::mixed_left_join( + rightTableView.select(rightKeyIndices_), + leftTableView.select(leftKeyIndices_), + rightTableView, + leftTableView, + tree_.back(), + cudf::null_equality::UNEQUAL, + std::nullopt, + stream); + } else { + std::tie(rightJoinIndices, leftJoinIndices) = cudf::left_join( + rightTableView.select(rightKeyIndices_), + leftTableView.select(leftKeyIndices_), + cudf::null_equality::UNEQUAL, + stream, + cudf::get_current_device_resource_ref()); + } + } else if (joinNode_->isAntiJoin()) { + if (joinNode_->filter()) { + leftJoinIndices = cudf::mixed_left_anti_join( + leftTableView.select(leftKeyIndices_), + rightTableView.select(rightKeyIndices_), + leftTableView, + rightTableView, + tree_.back(), + cudf::null_equality::UNEQUAL, + stream, + cudf::get_current_device_resource_ref()); + } else { + auto const rightTableHasNulls = + cudf::has_nulls(rightTableView.select(rightKeyIndices_)); + if (joinNode_->isNullAware() and rightTableHasNulls) { + // empty result + leftJoinIndices = + std::make_unique>( + 0, stream, cudf::get_current_device_resource_ref()); + } else { + leftJoinIndices = cudf::left_anti_join( + leftTableView.select(leftKeyIndices_), + rightTableView.select(rightKeyIndices_), + cudf::null_equality::UNEQUAL, + stream, + cudf::get_current_device_resource_ref()); + } + } + } else if (joinNode_->isLeftSemiFilterJoin()) { + if (joinNode_->filter()) { + leftJoinIndices = cudf::mixed_left_semi_join( + leftTableView.select(leftKeyIndices_), + rightTableView.select(rightKeyIndices_), + leftTableView, + rightTableView, + tree_.back(), + cudf::null_equality::UNEQUAL, + stream, + cudf::get_current_device_resource_ref()); + } else { + leftJoinIndices = cudf::left_semi_join( + leftTableView.select(leftKeyIndices_), + rightTableView.select(rightKeyIndices_), + cudf::null_equality::UNEQUAL, + stream, + cudf::get_current_device_resource_ref()); + } + } else if (joinNode_->isRightSemiFilterJoin()) { + if (joinNode_->filter()) { + rightJoinIndices = cudf::mixed_left_semi_join( + rightTableView.select(rightKeyIndices_), + leftTableView.select(leftKeyIndices_), + rightTableView, + leftTableView, + tree_.back(), + cudf::null_equality::UNEQUAL, + stream, + cudf::get_current_device_resource_ref()); + } else { + rightJoinIndices = cudf::left_semi_join( + rightTableView.select(rightKeyIndices_), + leftTableView.select(leftKeyIndices_), + cudf::null_equality::UNEQUAL, + stream, + cudf::get_current_device_resource_ref()); + } + } else { + VELOX_FAIL("Unsupported join type: ", joinNode_->joinType()); + } + auto leftIndicesSpan = leftJoinIndices + ? cudf::device_span{*leftJoinIndices} + : cudf::device_span{}; + auto rightIndicesSpan = rightJoinIndices + ? cudf::device_span{*rightJoinIndices} + : cudf::device_span{}; + + auto leftInput = leftTableView.select(leftColumnIndicesToGather_); + auto rightInput = rightTableView.select(rightColumnIndicesToGather_); + + auto leftIndicesCol = cudf::column_view{leftIndicesSpan}; + auto rightIndicesCol = cudf::column_view{rightIndicesSpan}; + auto constexpr oobPolicy = cudf::out_of_bounds_policy::NULLIFY; + auto leftResult = cudf::gather(leftInput, leftIndicesCol, oobPolicy, stream); + auto rightResult = + cudf::gather(rightInput, rightIndicesCol, oobPolicy, stream); + + if (cudfDebugEnabled()) { + std::cout << "Left result number of columns: " << leftResult->num_columns() + << std::endl; + std::cout << "Right result number of columns: " + << rightResult->num_columns() << std::endl; + } + + auto leftCols = leftResult->release(); + auto rightCols = rightResult->release(); + auto joinedCols = + std::vector>(outputType_->names().size()); + for (int i = 0; i < leftColumnOutputIndices_.size(); i++) { + joinedCols[leftColumnOutputIndices_[i]] = std::move(leftCols[i]); + } + for (int i = 0; i < rightColumnOutputIndices_.size(); i++) { + joinedCols[rightColumnOutputIndices_[i]] = std::move(rightCols[i]); + } + auto cudfOutput = std::make_unique(std::move(joinedCols)); + stream.synchronize(); + + input_.reset(); + finished_ = noMoreInput_; + + auto const size = cudfOutput->num_rows(); + if (cudfOutput->num_columns() == 0 or size == 0) { + return nullptr; + } + return std::make_shared( + pool(), outputType_, size, std::move(cudfOutput), stream); +} + +bool CudfHashJoinProbe::skipProbeOnEmptyBuild() const { + auto const joinType = joinNode_->joinType(); + return isInnerJoin(joinType) || isLeftSemiFilterJoin(joinType) || + isRightJoin(joinType) || isRightSemiFilterJoin(joinType) || + isRightSemiProjectJoin(joinType); +} + +exec::BlockingReason CudfHashJoinProbe::isBlocked(ContinueFuture* future) { + if (hashObject_.has_value()) { + return exec::BlockingReason::kNotBlocked; + } + + auto joinBridge = operatorCtx_->task()->getCustomJoinBridge( + operatorCtx_->driverCtx()->splitGroupId, planNodeId()); + auto cudfJoinBridge = + std::dynamic_pointer_cast(joinBridge); + VELOX_CHECK_NOT_NULL(cudfJoinBridge); + VELOX_CHECK_NOT_NULL(future); + auto hashObject = cudfJoinBridge->hashOrFuture(future); + + if (!hashObject.has_value()) { + if (cudfDebugEnabled()) { + std::cout << "CudfHashJoinProbe is blocked, waiting for join build" + << std::endl; + } + return exec::BlockingReason::kWaitForJoinBuild; + } + hashObject_ = std::move(hashObject); + + auto& rightTable = hashObject_.value().first; + // should be rightTable->numDistinct() but it needs compute, + // so we use num_rows() + if (rightTable->num_rows() == 0) { + if (skipProbeOnEmptyBuild()) { + if (operatorCtx_->driverCtx() + ->queryConfig() + .hashProbeFinishEarlyOnEmptyBuild()) { + noMoreInput(); + } else { + skipInput_ = true; + } + } + } + return exec::BlockingReason::kNotBlocked; +} + +bool CudfHashJoinProbe::isFinished() { + auto const isFinished = finished_ || (noMoreInput_ && input_ == nullptr); + + // Release hashObject_ if finished + if (isFinished) { + hashObject_.reset(); + } + return isFinished; +} + +std::unique_ptr CudfHashJoinBridgeTranslator::toOperator( + exec::DriverCtx* ctx, + int32_t id, + const core::PlanNodePtr& node) { + if (cudfDebugEnabled()) { + std::cout << "Calling CudfHashJoinBridgeTranslator::toOperator" + << std::endl; + } + if (auto joinNode = + std::dynamic_pointer_cast(node)) { + return std::make_unique(id, ctx, joinNode); + } + return nullptr; +} + +std::unique_ptr CudfHashJoinBridgeTranslator::toJoinBridge( + const core::PlanNodePtr& node) { + if (cudfDebugEnabled()) { + std::cout << "Calling CudfHashJoinBridgeTranslator::toJoinBridge" + << std::endl; + } + if (auto joinNode = + std::dynamic_pointer_cast(node)) { + auto joinBridge = std::make_unique(); + return joinBridge; + } + return nullptr; +} + +exec::OperatorSupplier CudfHashJoinBridgeTranslator::toOperatorSupplier( + const core::PlanNodePtr& node) { + if (cudfDebugEnabled()) { + std::cout << "Calling CudfHashJoinBridgeTranslator::toOperatorSupplier" + << std::endl; + } + if (auto joinNode = + std::dynamic_pointer_cast(node)) { + return [joinNode](int32_t operatorId, exec::DriverCtx* ctx) { + return std::make_unique(operatorId, ctx, joinNode); + }; + } + return nullptr; +} + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/exec/CudfHashJoin.h b/velox/experimental/cudf/exec/CudfHashJoin.h new file mode 100644 index 00000000000..6d7955ddb10 --- /dev/null +++ b/velox/experimental/cudf/exec/CudfHashJoin.h @@ -0,0 +1,138 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#pragma once + +#include "velox/experimental/cudf/exec/NvtxHelper.h" +#include "velox/experimental/cudf/vector/CudfVector.h" + +#include "velox/core/Expressions.h" +#include "velox/core/PlanNode.h" +#include "velox/exec/JoinBridge.h" +#include "velox/exec/Operator.h" +#include "velox/vector/ComplexVector.h" + +#include +#include +#include + +namespace facebook::velox::cudf_velox { + +class CudfHashJoinBridge : public exec::JoinBridge { + public: + using hash_type = + std::pair, std::shared_ptr>; + + void setHashTable(std::optional hashObject); + + std::optional hashOrFuture(ContinueFuture* future); + + private: + std::optional hashObject_; +}; + +class CudfHashJoinBuild : public exec::Operator, public NvtxHelper { + public: + CudfHashJoinBuild( + int32_t operatorId, + exec::DriverCtx* driverCtx, + std::shared_ptr joinNode); + + void addInput(RowVectorPtr input) override; + + bool needsInput() const override; + + RowVectorPtr getOutput() override; + + void noMoreInput() override; + + exec::BlockingReason isBlocked(ContinueFuture* future) override; + + bool isFinished() override; + + private: + std::shared_ptr joinNode_; + std::vector inputs_; + ContinueFuture future_{ContinueFuture::makeEmpty()}; +}; + +class CudfHashJoinProbe : public exec::Operator, public NvtxHelper { + public: + using hash_type = CudfHashJoinBridge::hash_type; + + CudfHashJoinProbe( + int32_t operatorId, + exec::DriverCtx* driverCtx, + std::shared_ptr joinNode); + + bool needsInput() const override; + + void addInput(RowVectorPtr input) override; + + RowVectorPtr getOutput() override; + + bool skipProbeOnEmptyBuild() const; + + exec::BlockingReason isBlocked(ContinueFuture* future) override; + + static bool isSupportedJoinType(core::JoinType joinType) { + return joinType == core::JoinType::kInner || + joinType == core::JoinType::kLeft || + joinType == core::JoinType::kAnti || + joinType == core::JoinType::kLeftSemiFilter; + } + + bool isFinished() override; + + private: + std::shared_ptr joinNode_; + std::optional hashObject_; + + // Filter related members + cudf::ast::tree tree_; + std::vector> scalars_; + + bool rightPrecomputed_{false}; + + std::vector leftKeyIndices_; + std::vector rightKeyIndices_; + std::vector leftColumnIndicesToGather_; + std::vector rightColumnIndicesToGather_; + std::vector leftColumnOutputIndices_; + std::vector rightColumnOutputIndices_; + bool finished_{false}; + + // Copied from HashProbe.h + // Indicates whether to skip probe input data processing or not. It only + // applies for a specific set of join types (see skipProbeOnEmptyBuild()), and + // the build table is empty and the probe input is read from non-spilled + // source. This ensures the hash probe operator keeps running until all the + // probe input from the sources have been processed. It prevents the exchange + // hanging problem at the producer side caused by the early query finish. + bool skipInput_{false}; +}; + +class CudfHashJoinBridgeTranslator : public exec::Operator::PlanNodeTranslator { + public: + std::unique_ptr + toOperator(exec::DriverCtx* ctx, int32_t id, const core::PlanNodePtr& node); + + std::unique_ptr toJoinBridge(const core::PlanNodePtr& node); + + exec::OperatorSupplier toOperatorSupplier(const core::PlanNodePtr& node); +}; + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/exec/CudfLimit.cpp b/velox/experimental/cudf/exec/CudfLimit.cpp new file mode 100644 index 00000000000..7cf320c89a6 --- /dev/null +++ b/velox/experimental/cudf/exec/CudfLimit.cpp @@ -0,0 +1,139 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#include "velox/experimental/cudf/exec/CudfLimit.h" +#include "velox/experimental/cudf/vector/CudfVector.h" + +#include + +namespace facebook::velox::cudf_velox { + +CudfLimit::CudfLimit( + int32_t operatorId, + exec::DriverCtx* driverCtx, + const std::shared_ptr& limitNode) + : Operator( + driverCtx, + limitNode->outputType(), + operatorId, + limitNode->id(), + "CudfLimit"), + NvtxHelper( + nvtx3::rgb{112, 128, 144}, // Slate Gray + operatorId, + fmt::format("[{}]", limitNode->id())), + remainingOffset_{limitNode->offset()}, + remainingLimit_{limitNode->count()} { + isIdentityProjection_ = true; + + const auto numColumns = limitNode->outputType()->size(); + identityProjections_.reserve(numColumns); + for (column_index_t i = 0; i < numColumns; ++i) { + identityProjections_.emplace_back(i, i); + } +} + +bool CudfLimit::needsInput() const { + return !finished_ && input_ == nullptr; +} + +void CudfLimit::addInput(RowVectorPtr input) { + VELOX_CHECK_NULL(input_); + input_ = input; +} + +RowVectorPtr CudfLimit::getOutput() { + VELOX_NVTX_OPERATOR_FUNC_RANGE(); + if (input_ == nullptr || (remainingOffset_ == 0 && remainingLimit_ == 0)) { + return nullptr; + } + + const auto inputSize = input_->size(); + + if (remainingOffset_ >= inputSize) { + remainingOffset_ -= inputSize; + input_ = nullptr; + return nullptr; + } + + auto cudfInput = std::dynamic_pointer_cast(input_); + + // This is the case where the offset lies in the middle of the current batch + // we want to start outputting rows from the middle of the input. + if (remainingOffset_ > 0) { + // Return a subset of input_ rows. + const auto outputSize = + std::min(inputSize - remainingOffset_, remainingLimit_); + + auto slicedTable = cudf::slice( + cudfInput->getTableView(), + {static_cast(remainingOffset_), + static_cast(remainingOffset_ + outputSize)}, + cudfInput->stream()); + + auto materializedTable = + std::make_unique(slicedTable[0], cudfInput->stream()); + + remainingOffset_ = 0; + remainingLimit_ -= outputSize; + if (remainingLimit_ == 0) { + finished_ = true; + } + auto output = std::make_shared( + input_->pool(), + input_->type(), + outputSize, + std::move(materializedTable), + cudfInput->stream()); + input_.reset(); + return output; + } + + if (remainingLimit_ <= inputSize) { + finished_ = true; + } + + // This is the case where we want to output all rows from the input because + // the range we want to output exceeds the input in both directions. + if (remainingLimit_ >= inputSize) { + remainingLimit_ -= inputSize; + auto output = input_; + input_.reset(); + return output; + } + + // At this point, we have no offset but the limit is less than the input size. + // We want to slice from the beginning but till the middle of the input. + auto slicedTable = cudf::slice( + cudfInput->getTableView(), + {0, static_cast(remainingLimit_)}, + cudfInput->stream()); + + auto materializedTable = + std::make_unique(slicedTable[0], cudfInput->stream()); + + auto output = std::make_shared( + input_->pool(), + input_->type(), + remainingLimit_, + std::move(materializedTable), + cudfInput->stream()); + input_.reset(); + remainingLimit_ = 0; + return output; +} + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/exec/CudfLimit.h b/velox/experimental/cudf/exec/CudfLimit.h new file mode 100644 index 00000000000..c6220e26f92 --- /dev/null +++ b/velox/experimental/cudf/exec/CudfLimit.h @@ -0,0 +1,50 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#pragma once + +#include "velox/experimental/cudf/exec/NvtxHelper.h" + +#include "velox/exec/Operator.h" + +namespace facebook::velox::cudf_velox { +class CudfLimit : public exec::Operator, public NvtxHelper { + public: + CudfLimit( + int32_t operatorId, + exec::DriverCtx* driverCtx, + const std::shared_ptr& limitNode); + + bool needsInput() const override; + + void addInput(RowVectorPtr input) override; + + RowVectorPtr getOutput() override; + + exec::BlockingReason isBlocked(ContinueFuture* /*future*/) override { + return exec::BlockingReason::kNotBlocked; + } + + bool isFinished() override { + return finished_ || (noMoreInput_ && input_ == nullptr); + } + + private: + int64_t remainingOffset_; + int64_t remainingLimit_; + bool finished_{false}; +}; +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/exec/CudfLocalPartition.cpp b/velox/experimental/cudf/exec/CudfLocalPartition.cpp new file mode 100644 index 00000000000..bb62137a3f3 --- /dev/null +++ b/velox/experimental/cudf/exec/CudfLocalPartition.cpp @@ -0,0 +1,206 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#include "velox/experimental/cudf/exec/CudfLocalPartition.h" +#include "velox/experimental/cudf/vector/CudfVector.h" + +#include "velox/exec/Task.h" + +#include +#include + +namespace facebook::velox::cudf_velox { + +CudfLocalPartition::CudfLocalPartition( + int32_t operatorId, + exec::DriverCtx* ctx, + const std::shared_ptr& planNode) + : Operator( + ctx, + planNode->outputType(), + operatorId, + planNode->id(), + "CudfLocalPartition"), + NvtxHelper( + nvtx3::rgb{255, 215, 0}, // Gold + operatorId, + fmt::format("[{}]", planNode->id())), + queues_{ + ctx->task->getLocalExchangeQueues(ctx->splitGroupId, planNode->id())}, + numPartitions_{queues_.size()} { + // Following is IMO a hacky way to get the partition key indices. It is to + // workaround the fact that the partition spec constructs the hash function + // directly and has no public methods to get the partition key indices. + + // When the operator is of type kRepartition, the partition spec is a string + // in the format "HASH(key1, key2, ...)" + // We're going to extract the keys between HASH( and ) and find their indices + // in the output row type. + + // When operator is of type kGather, we don't need to store any partition key + // indices because we're going to merge all the incoming streams together. + + // Get partition function specification string + std::string spec = planNode->partitionFunctionSpec().toString(); + + // Only parse keys if it's a hash function + if (spec.find("HASH(") != std::string::npos) { + // Extract keys between HASH( and ) + size_t start = spec.find("HASH(") + 5; + size_t end = spec.find(")", start); + if (start != std::string::npos && end != std::string::npos) { + std::string keysStr = spec.substr(start, end - start); + + // Split by comma to get individual keys. + std::vector keys; + size_t pos = 0; + while ((pos = keysStr.find(",")) != std::string::npos) { + std::string key = keysStr.substr(0, pos); + keys.push_back(key); + keysStr.erase(0, pos + 1); + } + keys.push_back(keysStr); // Add the last key. + + // Find field indices for each key. + const auto& rowType = planNode->outputType(); + for (const auto& key : keys) { + auto trimmedKey = key; + // Trim whitespace + trimmedKey.erase(0, trimmedKey.find_first_not_of(" ")); + trimmedKey.erase(trimmedKey.find_last_not_of(" ") + 1); + + auto fieldIndex = rowType->getChildIdx(trimmedKey); + partitionKeyIndices_.push_back(fieldIndex); + } + } + } + VELOX_CHECK(numPartitions_ == 1 || partitionKeyIndices_.size() > 0); + + // Since we're replacing the LocalPartition with CudfLocalPartition, the + // number of producers is already set. Adding producer only adds to a counter + // which we don't have to do again. + // Normally, this is what we'd have to do: + // for (auto& queue : queues_) { + // queue->addProducer(); + // } +} + +void CudfLocalPartition::recordOutputStats(RowVectorPtr& input) { + { + auto lockedStats = stats_.wlock(); + lockedStats->addOutputVector(input->estimateFlatSize(), input->size()); + } +} + +void CudfLocalPartition::addInput(RowVectorPtr input) { + VELOX_NVTX_OPERATOR_FUNC_RANGE(); + recordOutputStats(input); + auto cudfVector = std::dynamic_pointer_cast(input); + VELOX_CHECK(cudfVector, "Input must be a CudfVector"); + auto stream = cudfVector->stream(); + + if (numPartitions_ > 1) { + // Use cudf hash partitioning + auto tableView = cudfVector->getTableView(); + std::vector partitionKeyIndices; + for (const auto& idx : partitionKeyIndices_) { + partitionKeyIndices.push_back(static_cast(idx)); + } + + auto [partitionedTable, partitionOffsets] = cudf::hash_partition( + tableView, + partitionKeyIndices, + numPartitions_, + cudf::hash_id::HASH_MURMUR3, + cudf::DEFAULT_HASH_SEED, + stream); + + VELOX_CHECK(partitionOffsets.size() == numPartitions_); + VELOX_CHECK(partitionOffsets[0] == 0); + + // Erase first element since it's always 0 and we don't need it. + partitionOffsets.erase(partitionOffsets.begin()); + + auto partitionedTables = + cudf::split(partitionedTable->view(), partitionOffsets); + + for (int i = 0; i < numPartitions_; ++i) { + auto partitionData = partitionedTables[i]; + if (partitionData.num_rows() == 0) { + // Skip empty partitions. + continue; + } + + ContinueFuture future; + // DM: We should investigate if keeping partitionedTables alive and using + // the table view in partitonedData is more efficient than creating a new + // table each time. Currently out of scope because it would need a new + // type of RowVector that can hold a table view and shared_ptr to the + // table. + auto blockingReason = queues_[i]->enqueue( + std::make_shared( + pool(), + outputType_, + partitionData.num_rows(), + std::make_unique(partitionData), + stream), + partitionData.num_rows(), + &future); + if (blockingReason != exec::BlockingReason::kNotBlocked) { + blockingReasons_.push_back(blockingReason); + futures_.push_back(std::move(future)); + } + } + } else { + // Single partition case. + ContinueFuture future; + auto blockingReason = + queues_[0]->enqueue(input, input->retainedSize(), &future); + if (blockingReason != exec::BlockingReason::kNotBlocked) { + blockingReasons_.push_back(blockingReason); + futures_.push_back(std::move(future)); + } + } +} + +exec::BlockingReason CudfLocalPartition::isBlocked(ContinueFuture* future) { + if (!futures_.empty()) { + auto blockingReason = blockingReasons_.front(); + *future = folly::collectAll(futures_.begin(), futures_.end()).unit(); + futures_.clear(); + blockingReasons_.clear(); + return blockingReason; + } + + return exec::BlockingReason::kNotBlocked; +} + +void CudfLocalPartition::noMoreInput() { + Operator::noMoreInput(); + for (const auto& queue : queues_) { + queue->noMoreData(); + } +} + +bool CudfLocalPartition::isFinished() { + if (!futures_.empty() || !noMoreInput_) { + return false; + } + + return true; +} + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/exec/CudfLocalPartition.h b/velox/experimental/cudf/exec/CudfLocalPartition.h new file mode 100644 index 00000000000..b443dba8097 --- /dev/null +++ b/velox/experimental/cudf/exec/CudfLocalPartition.h @@ -0,0 +1,66 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ +#pragma once + +#include "velox/experimental/cudf/exec/NvtxHelper.h" + +#include "velox/exec/LocalPartition.h" +#include "velox/exec/Operator.h" + +namespace facebook::velox::cudf_velox { + +class CudfLocalPartition : public exec::Operator, public NvtxHelper { + public: + CudfLocalPartition( + int32_t operatorId, + exec::DriverCtx* driverCtx, + const std::shared_ptr& planNode); + + std::string toString() const override { + return fmt::format("CudfLocalPartition({})", numPartitions_); + } + + void recordOutputStats(RowVectorPtr& input); + + void addInput(RowVectorPtr input) override; + + RowVectorPtr getOutput() override { + return nullptr; + } + + /// Always true but the caller will check isBlocked before adding input, hence + /// the blocked state does not accumulate input. + bool needsInput() const override { + return true; + } + + exec::BlockingReason isBlocked(ContinueFuture* future) override; + + void noMoreInput() override; + + bool isFinished() override; + + protected: + const std::vector> queues_; + const size_t numPartitions_; + + std::vector blockingReasons_; + std::vector futures_; + + std::vector partitionKeyIndices_; +}; + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/exec/DebugUtil.cpp b/velox/experimental/cudf/exec/DebugUtil.cpp new file mode 100644 index 00000000000..b1153c17c0d --- /dev/null +++ b/velox/experimental/cudf/exec/DebugUtil.cpp @@ -0,0 +1,32 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#include "velox/experimental/cudf/exec/DebugUtil.h" +#include "velox/experimental/cudf/exec/Utilities.h" +#include "velox/experimental/cudf/exec/VeloxCudfInterop.h" + +namespace facebook::velox::cudf_velox { + +std::string DebugUtil::toString( + const cudf::table_view& table, + rmm::cuda_stream_view stream, + vector_size_t from, + vector_size_t to) { + auto rowVector = with_arrow::toVeloxColumn(table, pool_.get(), "", stream); + stream.synchronize(); + return rowVector->toString(from, to); +} +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/exec/DebugUtil.h b/velox/experimental/cudf/exec/DebugUtil.h new file mode 100644 index 00000000000..7fc4de0ea4c --- /dev/null +++ b/velox/experimental/cudf/exec/DebugUtil.h @@ -0,0 +1,37 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#pragma once + +#include "velox/experimental/cudf/vector/CudfVector.h" + +namespace facebook::velox::cudf_velox { +class DebugUtil { + public: + std::string toString( + const cudf::table_view& table, + rmm::cuda_stream_view stream, + vector_size_t from, + vector_size_t to); + + private: + std::shared_ptr rootPool_{ + memory::memoryManager()->addRootPool()}; + std::shared_ptr pool_{ + rootPool_->addLeafChild("debug_util")}; +}; + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/exec/ExpressionEvaluator.cpp b/velox/experimental/cudf/exec/ExpressionEvaluator.cpp new file mode 100644 index 00000000000..112c0e9d76d --- /dev/null +++ b/velox/experimental/cudf/exec/ExpressionEvaluator.cpp @@ -0,0 +1,644 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ +#include "velox/experimental/cudf/exec/ExpressionEvaluator.h" + +#include "velox/expression/ConstantExpr.h" +#include "velox/expression/FieldReference.h" +#include "velox/type/Type.h" +#include "velox/vector/BaseVector.h" +#include "velox/vector/ComplexVector.h" +#include "velox/vector/ConstantVector.h" +#include "velox/vector/VectorTypeUtils.h" + +#include +#include +#include +#include +#include +#include + +namespace facebook::velox::cudf_velox { +namespace { +template +cudf::ast::literal makeScalarAndLiteral( + const VectorPtr& vector, + std::vector>& scalars, + size_t atIndex = 0) { + using T = typename facebook::velox::KindToFlatVector::WrapperType; + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + const auto& type = vector->type(); + + if constexpr (cudf::is_fixed_width()) { + auto constVector = vector->as>(); + VELOX_CHECK_NOT_NULL(constVector, "ConstantVector is null"); + T value = constVector->valueAt(atIndex); + if (type->isShortDecimal()) { + VELOX_FAIL("Short decimal not supported"); + /* TODO: enable after rewriting using binary ops + using CudfDecimalType = cudf::numeric::decimal64; + using cudfScalarType = cudf::fixed_point_scalar; + auto scalar = std::make_unique(value, + type->scale(), + true, + stream, + mr); + scalars.emplace_back(std::move(scalar)); + return cudf::ast::literal{ + *static_cast(scalars.back().get())}; + */ + } else if (type->isLongDecimal()) { + VELOX_FAIL("Long decimal not supported"); + /* TODO: enable after rewriting using binary ops + using CudfDecimalType = cudf::numeric::decimal128; + using cudfScalarType = cudf::fixed_point_scalar; + auto scalar = std::make_unique(value, + type->scale(), + true, + stream, + mr); + scalars.emplace_back(std::move(scalar)); + return cudf::ast::literal{ + *static_cast(scalars.back().get())}; + */ + } else if (type->isIntervalYearMonth()) { + // no support for interval year month in cudf + VELOX_FAIL("Interval year month not supported"); + } else if (type->isIntervalDayTime()) { + using CudfDurationType = cudf::duration_ms; + if constexpr (std::is_same_v) { + using CudfScalarType = cudf::duration_scalar; + auto scalar = std::make_unique(value, true, stream, mr); + scalars.emplace_back(std::move(scalar)); + return cudf::ast::literal{ + *static_cast(scalars.back().get())}; + } + } else if (type->isDate()) { + using CudfDateType = cudf::timestamp_D; + if constexpr (std::is_same_v) { + using CudfScalarType = cudf::timestamp_scalar; + auto scalar = std::make_unique(value, true, stream, mr); + scalars.emplace_back(std::move(scalar)); + return cudf::ast::literal{ + *static_cast(scalars.back().get())}; + } + } else { + // Create a numeric scalar of type T, store it in the scalars vector, + // and use its reference in the literal expression. + using CudfScalarType = cudf::numeric_scalar; + scalars.emplace_back( + std::make_unique(value, true, stream, mr)); + return cudf::ast::literal{ + *static_cast(scalars.back().get())}; + } + VELOX_FAIL("Unsupported base type for literal"); + } else if (kind == TypeKind::VARCHAR) { + auto constVector = vector->as>(); + auto value = constVector->valueAt(atIndex); + std::string_view stringValue = static_cast(value); + scalars.emplace_back( + std::make_unique(stringValue, true, stream, mr)); + return cudf::ast::literal{ + *static_cast(scalars.back().get())}; + } else { + // TODO for non-numeric types too. + VELOX_NYI( + "Non-numeric types not yet implemented for kind " + + mapTypeKindToName(kind)); + } +} + +cudf::ast::literal createLiteral( + const VectorPtr& vector, + std::vector>& scalars, + size_t atIndex = 0) { + const auto kind = vector->typeKind(); + return VELOX_DYNAMIC_TYPE_DISPATCH_ALL( + makeScalarAndLiteral, kind, std::move(vector), scalars, atIndex); +} + +// Helper function to extract literals from array elements based on type +void extractArrayLiterals( + const ArrayVector* arrayVector, + std::vector& literals, + std::vector>& scalars, + vector_size_t offset, + vector_size_t size) { + auto elements = arrayVector->elements(); + + for (auto i = offset; i < offset + size; ++i) { + if (elements->isNullAt(i)) { + // Skip null values for IN expressions + continue; + } else { + literals.emplace_back(createLiteral(elements, scalars, i)); + } + } +} + +// Function to create literals from an array vector +std::vector createLiteralsFromArray( + const VectorPtr& vector, + std::vector>& scalars) { + std::vector literals; + + // Check if it's a constant vector containing an array + if (vector->isConstantEncoding()) { + auto constantVector = vector->asUnchecked>(); + if (constantVector->isNullAt(0)) { + // Return empty vector for null array + return literals; + } + + auto valueVector = constantVector->valueVector(); + if (valueVector->encoding() == VectorEncoding::Simple::ARRAY) { + auto arrayVector = valueVector->as(); + auto index = constantVector->index(); + auto size = arrayVector->sizeAt(index); + if (size == 0) { + // Return empty vector for empty array + return literals; + } + + auto offset = arrayVector->offsetAt(index); + auto elements = arrayVector->elements(); + + // Handle different element types + if (elements->isScalar()) { + literals.reserve(size); + extractArrayLiterals(arrayVector, literals, scalars, offset, size); + } else if (elements->typeKind() == TypeKind::ARRAY) { + // Nested arrays not supported in IN expressions + VELOX_FAIL("Nested arrays not supported in IN expressions"); + } else { + VELOX_FAIL( + "Unsupported element type in array: {}", + elements->type()->toString()); + } + } else { + VELOX_FAIL("Expected ARRAY encoding"); + } + } else { + VELOX_FAIL("Expected constant vector for IN list"); + } + + return literals; +} +} // namespace + +using Op = cudf::ast::ast_operator; +const std::unordered_map prestoBinaryOps = { + {"plus", Op::ADD}, + {"minus", Op::SUB}, + {"multiply", Op::MUL}, + {"divide", Op::DIV}, + {"eq", Op::EQUAL}, + {"neq", Op::NOT_EQUAL}, + {"lt", Op::LESS}, + {"gt", Op::GREATER}, + {"lte", Op::LESS_EQUAL}, + {"gte", Op::GREATER_EQUAL}, + {"and", Op::NULL_LOGICAL_AND}, + {"or", Op::NULL_LOGICAL_OR}, + {"mod", Op::MOD}, +}; + +const std::unordered_map sparkBinaryOps = { + {"add", Op::ADD}, + {"subtract", Op::SUB}, + {"multiply", Op::MUL}, + {"divide", Op::DIV}, + {"equalto", Op::EQUAL}, + {"lessthan", Op::LESS}, + {"greaterthan", Op::GREATER}, + {"lessthanorequal", Op::LESS_EQUAL}, + {"greaterthanorequal", Op::GREATER_EQUAL}, + {"and", Op::NULL_LOGICAL_AND}, + {"or", Op::NULL_LOGICAL_OR}}; + +const std::unordered_map binaryOps = [] { + std::unordered_map merged( + sparkBinaryOps.begin(), sparkBinaryOps.end()); + merged.insert(prestoBinaryOps.begin(), prestoBinaryOps.end()); + return merged; +}(); + +const std::map unaryOps = { + {"not", Op::NOT}, + {"is_null", Op::IS_NULL}}; + +const std::unordered_set supportedOps = { + "literal", + "between", + "in", + "cast", + "switch", + "year", + "length", + "substr", + "like"}; + +namespace detail { + +bool canBeEvaluated(const std::shared_ptr& expr) { + const auto& name = expr->name(); + if (supportedOps.count(name) || binaryOps.count(name) || + unaryOps.count(name)) { + return std::all_of( + expr->inputs().begin(), expr->inputs().end(), canBeEvaluated); + } + return std::dynamic_pointer_cast(expr) != + nullptr; +} + +} // namespace detail + +struct AstContext { + cudf::ast::tree& tree; + std::vector>& scalars; + const std::vector inputRowSchema; + const std::vector>> + precomputeInstructions; + cudf::ast::expression const& pushExprToTree( + const std::shared_ptr& expr); + cudf::ast::expression const& addPrecomputeInstruction( + std::string const& name, + std::string const& instruction); + cudf::ast::expression const& multipleInputsToPairWise( + const std::shared_ptr& expr); + static bool canBeEvaluated(const std::shared_ptr& expr); +}; + +// Create tree from Expr +// and collect precompute instructions for non-ast operations +cudf::ast::expression const& createAstTree( + const std::shared_ptr& expr, + cudf::ast::tree& tree, + std::vector>& scalars, + const RowTypePtr& inputRowSchema, + std::vector& precomputeInstructions) { + AstContext context{tree, scalars, {inputRowSchema}, {precomputeInstructions}}; + return context.pushExprToTree(expr); +} + +cudf::ast::expression const& createAstTree( + const std::shared_ptr& expr, + cudf::ast::tree& tree, + std::vector>& scalars, + const RowTypePtr& leftRowSchema, + const RowTypePtr& rightRowSchema, + std::vector& leftPrecomputeInstructions, + std::vector& rightPrecomputeInstructions) { + AstContext context{ + tree, + scalars, + {leftRowSchema, rightRowSchema}, + {leftPrecomputeInstructions, rightPrecomputeInstructions}}; + return context.pushExprToTree(expr); +} + +cudf::ast::expression const& AstContext::addPrecomputeInstruction( + std::string const& name, + std::string const& instruction) { + for (size_t sideIdx = 0; sideIdx < inputRowSchema.size(); ++sideIdx) { + if (inputRowSchema[sideIdx].get()->containsChild(name)) { + auto columnIndex = inputRowSchema[sideIdx].get()->getChildIdx(name); + auto newColumnIndex = inputRowSchema[sideIdx].get()->size() + + precomputeInstructions[sideIdx].get().size(); + // This custom op should be added to input columns. + precomputeInstructions[sideIdx].get().emplace_back( + columnIndex, instruction, newColumnIndex); + auto side = static_cast(sideIdx); + return tree.push(cudf::ast::column_reference(newColumnIndex, side)); + } + } + VELOX_FAIL("Field not found, " + name); +} + +/// Handles logical AND/OR expressions with multiple inputs by converting them +/// into a chain of binary operations. For example, "a AND b AND c" becomes +/// "(a AND b) AND c". +/// +/// @param expr The expression containing multiple inputs for AND/OR operation +/// @return A reference to the resulting AST expression +cudf::ast::expression const& AstContext::multipleInputsToPairWise( + const std::shared_ptr& expr) { + using Operation = cudf::ast::operation; + + const auto& name = expr->name(); + auto len = expr->inputs().size(); + // Create a simple chain of operations + auto result = &pushExprToTree(expr->inputs()[0]); + + // Chain the rest of the inputs sequentially + for (size_t i = 1; i < len; i++) { + auto const& nextInput = pushExprToTree(expr->inputs()[i]); + result = &tree.push(Operation{binaryOps.at(name), *result, nextInput}); + } + return *result; +} + +/// Pushes an expression into the AST tree and returns a reference to the +/// resulting expression. +/// +/// @param expr The expression to push into the AST tree +/// @return A reference to the resulting AST expression +cudf::ast::expression const& AstContext::pushExprToTree( + const std::shared_ptr& expr) { + using Op = cudf::ast::ast_operator; + using Operation = cudf::ast::operation; + using velox::exec::ConstantExpr; + using velox::exec::FieldReference; + + auto& name = expr->name(); + auto len = expr->inputs().size(); + + if (name == "literal") { + auto c = dynamic_cast(expr.get()); + VELOX_CHECK_NOT_NULL(c, "literal expression should be ConstantExpr"); + auto value = c->value(); + VELOX_CHECK(value->isConstantEncoding()); + // convert to cudf scalar + return tree.push(createLiteral(value, scalars)); + } else if (binaryOps.find(name) != binaryOps.end()) { + if (len > 2 and (name == "and" or name == "or")) { + return multipleInputsToPairWise(expr); + } + VELOX_CHECK_EQ(len, 2); + auto const& op1 = pushExprToTree(expr->inputs()[0]); + auto const& op2 = pushExprToTree(expr->inputs()[1]); + return tree.push(Operation{binaryOps.at(name), op1, op2}); + } else if (unaryOps.find(name) != unaryOps.end()) { + VELOX_CHECK_EQ(len, 1); + auto const& op1 = pushExprToTree(expr->inputs()[0]); + return tree.push(Operation{unaryOps.at(name), op1}); + } else if (name == "between") { + VELOX_CHECK_EQ(len, 3); + auto const& value = pushExprToTree(expr->inputs()[0]); + auto const& lower = pushExprToTree(expr->inputs()[1]); + auto const& upper = pushExprToTree(expr->inputs()[2]); + // construct between(op2, op3) using >= and <= + auto const& geLower = tree.push(Operation{Op::GREATER_EQUAL, value, lower}); + auto const& leUpper = tree.push(Operation{Op::LESS_EQUAL, value, upper}); + return tree.push(Operation{Op::NULL_LOGICAL_AND, geLower, leUpper}); + } else if (name == "in") { + // number of inputs is variable. >=2 + VELOX_CHECK_EQ(len, 2); + // actually len is 2, second input is ARRAY + auto const& op1 = pushExprToTree(expr->inputs()[0]); + auto c = dynamic_cast(expr->inputs()[1].get()); + VELOX_CHECK_NOT_NULL(c, "literal expression should be ConstantExpr"); + auto value = c->value(); + VELOX_CHECK_NOT_NULL(value, "ConstantExpr value is null"); + + // Use the new createLiteralsFromArray function to get literals + auto literals = createLiteralsFromArray(value, scalars); + + // Create equality expressions for each literal and OR them together + std::vector exprVec; + for (auto& literal : literals) { + auto const& opi = tree.push(std::move(literal)); + auto const& logicalNode = tree.push(Operation{Op::EQUAL, op1, opi}); + exprVec.push_back(&logicalNode); + } + + // Handle empty IN list case + if (exprVec.empty()) { + // FAIL + VELOX_FAIL("Empty IN list"); + // Return FALSE for empty IN list + // auto falseValue = std::make_shared>( + // value->pool(), 1, false, TypeKind::BOOLEAN, false); + // return tree.push(createLiteral(falseValue, scalars)); + } + + // OR all logical nodes + auto* result = exprVec[0]; + for (size_t i = 1; i < exprVec.size(); i++) { + auto const& treeNode = + tree.push(Operation{Op::NULL_LOGICAL_OR, *result, *exprVec[i]}); + result = &treeNode; + } + return *result; + } else if (name == "cast") { + VELOX_CHECK_EQ(len, 1); + auto const& op1 = pushExprToTree(expr->inputs()[0]); + if (expr->type()->kind() == TypeKind::INTEGER) { + // No int32 cast in cudf ast + return tree.push(Operation{Op::CAST_TO_INT64, op1}); + } else if (expr->type()->kind() == TypeKind::BIGINT) { + return tree.push(Operation{Op::CAST_TO_INT64, op1}); + } else if (expr->type()->kind() == TypeKind::DOUBLE) { + return tree.push(Operation{Op::CAST_TO_FLOAT64, op1}); + } else { + VELOX_FAIL("Unsupported type for cast operation"); + } + } else if (name == "switch") { + VELOX_CHECK_EQ(len, 3); + // check if input[1], input[2] are literals 1 and 0. + // then simplify as typecast bool to int + auto c1 = dynamic_cast(expr->inputs()[1].get()); + auto c2 = dynamic_cast(expr->inputs()[2].get()); + if (c1 and c1->toString() == "1:BIGINT" and c2 and + c2->toString() == "0:BIGINT") { + auto const& op1 = pushExprToTree(expr->inputs()[0]); + return tree.push(Operation{Op::CAST_TO_INT64, op1}); + } else if (c2 and c2->toString() == "0:DOUBLE") { + auto const& op1 = pushExprToTree(expr->inputs()[0]); + auto const& op1d = tree.push(Operation{Op::CAST_TO_FLOAT64, op1}); + auto const& op2 = pushExprToTree(expr->inputs()[1]); + return tree.push(Operation{Op::MUL, op1d, op2}); + } else { + VELOX_NYI("Unsupported switch complex operation " + expr->toString()); + } + } else if (name == "year") { + VELOX_CHECK_EQ(len, 1); + + auto fieldExpr = + std::dynamic_pointer_cast(expr->inputs()[0]); + VELOX_CHECK_NOT_NULL(fieldExpr, "Expression is not a field"); + + auto const& colRef = addPrecomputeInstruction(fieldExpr->name(), "year"); + + return tree.push(Operation{Op::CAST_TO_INT64, colRef}); + } else if (name == "length") { + VELOX_CHECK_EQ(len, 1); + + auto fieldExpr = + std::dynamic_pointer_cast(expr->inputs()[0]); + VELOX_CHECK_NOT_NULL(fieldExpr, "Expression is not a field"); + + auto const& colRef = addPrecomputeInstruction(fieldExpr->name(), "length"); + + return tree.push(Operation{Op::CAST_TO_INT64, colRef}); + } else if (name == "substr") { + // Extract the start and length parameters from the substr function call + // and create a precomputed column with the substring operation. + // This will be handled during AST evaluation with special column reference. + VELOX_CHECK_EQ(len, 3); + auto fieldExpr = + std::dynamic_pointer_cast(expr->inputs()[0]); + VELOX_CHECK_NOT_NULL(fieldExpr, "Expression is not a field"); + + auto c1 = dynamic_cast(expr->inputs()[1].get()); + auto c2 = dynamic_cast(expr->inputs()[2].get()); + std::string substrExpr = + "substr " + c1->value()->toString(0) + " " + c2->value()->toString(0); + + return addPrecomputeInstruction(fieldExpr->name(), substrExpr); + } else if (name == "like") { + VELOX_CHECK_EQ(len, 2); + + auto fieldExpr = + std::dynamic_pointer_cast(expr->inputs()[0]); + VELOX_CHECK_NOT_NULL(fieldExpr, "Expression is not a field"); + auto literalExpr = + std::dynamic_pointer_cast(expr->inputs()[1]); + VELOX_CHECK_NOT_NULL(literalExpr, "Expression is not a literal"); + + createLiteral(literalExpr->value(), scalars); + + std::string likeExpr = "like " + std::to_string(scalars.size() - 1); + + return addPrecomputeInstruction(fieldExpr->name(), likeExpr); + } else if (auto fieldExpr = std::dynamic_pointer_cast(expr)) { + // Refer to the appropriate side + for (size_t sideIdx = 0; sideIdx < inputRowSchema.size(); ++sideIdx) { + auto& schema = inputRowSchema[sideIdx]; + if (schema.get()->containsChild(name)) { + auto columnIndex = schema.get()->getChildIdx(name); + auto side = static_cast(sideIdx); + return tree.push(cudf::ast::column_reference(columnIndex, side)); + } + } + VELOX_FAIL("Field not found, " + name); + } else { + std::cerr << "Unsupported expression: " << expr->toString() << std::endl; + VELOX_FAIL("Unsupported expression: " + name); + } +} + +void addPrecomputedColumns( + std::vector>& input_table_columns, + const std::vector& precompute_instructions, + const std::vector>& scalars, + rmm::cuda_stream_view stream) { + for (const auto& instruction : precompute_instructions) { + auto [dependent_column_index, ins_name, new_column_index] = instruction; + if (ins_name == "year") { + auto newColumn = cudf::datetime::extract_datetime_component( + input_table_columns[dependent_column_index]->view(), + cudf::datetime::datetime_component::YEAR, + stream, + cudf::get_current_device_resource_ref()); + input_table_columns.emplace_back(std::move(newColumn)); + } else if (ins_name == "length") { + auto newColumn = cudf::strings::count_characters( + input_table_columns[dependent_column_index]->view(), + stream, + cudf::get_current_device_resource_ref()); + input_table_columns.emplace_back(std::move(newColumn)); + } else if (ins_name.rfind("substr", 0) == 0) { + std::istringstream iss(ins_name.substr(6)); + int beginValue, length; + iss >> beginValue >> length; + auto beginScalar = cudf::numeric_scalar( + beginValue - 1, + true, + stream, + cudf::get_current_device_resource_ref()); + auto endScalar = cudf::numeric_scalar( + beginValue - 1 + length, + true, + stream, + cudf::get_current_device_resource_ref()); + auto stepScalar = cudf::numeric_scalar( + 1, true, stream, cudf::get_current_device_resource_ref()); + auto newColumn = cudf::strings::slice_strings( + input_table_columns[dependent_column_index]->view(), + beginScalar, + endScalar, + stepScalar, + stream, + cudf::get_current_device_resource_ref()); + input_table_columns.emplace_back(std::move(newColumn)); + } else if (ins_name.rfind("like", 0) == 0) { + auto scalarIndex = std::stoi(ins_name.substr(4)); + auto newColumn = cudf::strings::like( + input_table_columns[dependent_column_index]->view(), + *static_cast(scalars[scalarIndex].get()), + cudf::string_scalar( + "", true, stream, cudf::get_current_device_resource_ref()), + stream, + cudf::get_current_device_resource_ref()); + input_table_columns.emplace_back(std::move(newColumn)); + } else { + VELOX_FAIL("Unsupported precompute operation " + ins_name); + } + } +} + +ExpressionEvaluator::ExpressionEvaluator( + const std::vector>& exprs, + const RowTypePtr& inputRowSchema) { + exprAst_.reserve(exprs.size()); + for (const auto& expr : exprs) { + cudf::ast::tree tree; + createAstTree( + expr, tree, scalars_, inputRowSchema, precomputeInstructions_); + exprAst_.emplace_back(std::move(tree)); + } +} + +void ExpressionEvaluator::close() { + exprAst_.clear(); + scalars_.clear(); + precomputeInstructions_.clear(); +} + +std::vector> ExpressionEvaluator::compute( + std::vector>& inputTableColumns, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto numColumns = inputTableColumns.size(); + addPrecomputedColumns( + inputTableColumns, precomputeInstructions_, scalars_, stream); + auto astInputTable = + std::make_unique(std::move(inputTableColumns)); + auto astInputTableView = astInputTable->view(); + std::vector> columns; + for (auto& tree : exprAst_) { + if (auto colRefPtr = + dynamic_cast(&tree.back())) { + auto col = std::make_unique( + astInputTableView.column(colRefPtr->get_column_index()), stream, mr); + columns.emplace_back(std::move(col)); + } else { + auto col = + cudf::compute_column(astInputTableView, tree.back(), stream, mr); + columns.emplace_back(std::move(col)); + } + } + inputTableColumns = astInputTable->release(); + inputTableColumns.resize(numColumns); + return columns; +} + +bool ExpressionEvaluator::canBeEvaluated( + const std::vector>& exprs) { + return std::all_of(exprs.begin(), exprs.end(), detail::canBeEvaluated); +} +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/exec/ExpressionEvaluator.h b/velox/experimental/cudf/exec/ExpressionEvaluator.h new file mode 100644 index 00000000000..648215aa7e5 --- /dev/null +++ b/velox/experimental/cudf/exec/ExpressionEvaluator.h @@ -0,0 +1,97 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#pragma once + +#include "velox/core/Expressions.h" +#include "velox/expression/Expr.h" +#include "velox/type/Type.h" + +#include + +#include +#include +#include + +namespace facebook::velox::cudf_velox { + +// Pre-compute instructions for the expression, +// for ops that are not supported by cudf::ast +struct PrecomputeInstruction { + int dependent_column_index; + std::string ins_name; + int new_column_index; + + // Constructor to initialize the struct with values + PrecomputeInstruction(int depIndex, const std::string& name, int newIndex) + : dependent_column_index(depIndex), + ins_name(name), + new_column_index(newIndex) {} +}; + +cudf::ast::expression const& createAstTree( + const std::shared_ptr& expr, + cudf::ast::tree& tree, + std::vector>& scalars, + const RowTypePtr& inputRowSchema, + std::vector& precomputeInstructions); + +cudf::ast::expression const& createAstTree( + const std::shared_ptr& expr, + cudf::ast::tree& tree, + std::vector>& scalars, + const RowTypePtr& leftRowSchema, + const RowTypePtr& rightRowSchema, + std::vector& leftPrecomputeInstructions, + std::vector& rightPrecomputeInstructions); + +void addPrecomputedColumns( + std::vector>& inputTableColumns, + const std::vector& precomputeInstructions, + const std::vector>& scalars, + rmm::cuda_stream_view stream); + +// Evaluates the expression tree +class ExpressionEvaluator { + public: + ExpressionEvaluator() = default; + // Converts velox expressions to cudf::ast::tree, scalars and + // precompute instructions and stores them + ExpressionEvaluator( + const std::vector>& exprs, + const RowTypePtr& inputRowSchema); + + // Evaluates the expression tree for the given input columns + std::vector> compute( + std::vector>& inputTableColumns, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + + void close(); + + static bool canBeEvaluated( + const std::vector>& exprs); + + private: + std::vector exprAst_; + std::vector> scalars_; + // instruction on dependent column to get new column index on non-ast + // supported operations in expressions + // + std::vector precomputeInstructions_; +}; + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/exec/PrestoAggregates.cpp b/velox/experimental/cudf/exec/PrestoAggregates.cpp new file mode 100644 index 00000000000..81092a29b3f --- /dev/null +++ b/velox/experimental/cudf/exec/PrestoAggregates.cpp @@ -0,0 +1,79 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#include "velox/experimental/cudf/exec/CudfHashAggregation.h" +#include "velox/experimental/cudf/exec/PrestoAggregates.h" + +#include "velox/functions/prestosql/aggregates/AggregateNames.h" + +namespace facebook::velox::cudf_velox::presto { + +namespace { +// Use constants from AggregateNames.h +using facebook::velox::aggregate::kAvg; +using facebook::velox::aggregate::kCount; +using facebook::velox::aggregate::kMax; +using facebook::velox::aggregate::kMin; +using facebook::velox::aggregate::kSum; +} // namespace + +void registerPrestoAggregate( + const std::string& prefix, + const std::string& aggregateName, + bool overwrite) { + auto name = prefix + aggregateName; + registerAggregator( + name, + [aggregateName]( + core::AggregationNode::Step step, + uint32_t inputIndex, + VectorPtr constant, + bool isGlobal) -> std::unique_ptr { + return facebook::velox::cudf_velox::createAggregator( + aggregateName, step, inputIndex, constant, isGlobal); + }, + overwrite); +} + +void registerPrestoSumAggregate(const std::string& prefix, bool overwrite) { + registerPrestoAggregate(prefix, kSum, overwrite); +} + +void registerPrestoCountAggregate(const std::string& prefix, bool overwrite) { + registerPrestoAggregate(prefix, kCount, overwrite); +} + +void registerPrestoMinAggregate(const std::string& prefix, bool overwrite) { + registerPrestoAggregate(prefix, kMin, overwrite); +} + +void registerPrestoMaxAggregate(const std::string& prefix, bool overwrite) { + registerPrestoAggregate(prefix, kMax, overwrite); +} + +void registerPrestoAvgAggregate(const std::string& prefix, bool overwrite) { + registerPrestoAggregate(prefix, kAvg, overwrite); +} + +void registerAllPrestoAggregates(const std::string& prefix, bool overwrite) { + registerPrestoSumAggregate(prefix, overwrite); + registerPrestoCountAggregate(prefix, overwrite); + registerPrestoMinAggregate(prefix, overwrite); + registerPrestoMaxAggregate(prefix, overwrite); + registerPrestoAvgAggregate(prefix, overwrite); +} + +} // namespace facebook::velox::cudf_velox::presto diff --git a/velox/experimental/cudf/exec/PrestoAggregates.h b/velox/experimental/cudf/exec/PrestoAggregates.h new file mode 100644 index 00000000000..dbbd6fde048 --- /dev/null +++ b/velox/experimental/cudf/exec/PrestoAggregates.h @@ -0,0 +1,50 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#pragma once + +#include + +namespace facebook::velox::cudf_velox::presto { + +/// Register presto-style aggregate functions that use CUDF aggregators +/// with the specified prefix. +void registerPrestoSumAggregate( + const std::string& prefix, + bool overwrite = false); + +void registerPrestoCountAggregate( + const std::string& prefix, + bool overwrite = false); + +void registerPrestoMinAggregate( + const std::string& prefix, + bool overwrite = false); + +void registerPrestoMaxAggregate( + const std::string& prefix, + bool overwrite = false); + +void registerPrestoAvgAggregate( + const std::string& prefix, + bool overwrite = false); + +/// Register all presto-style CUDF aggregators with the specified prefix +void registerAllPrestoAggregates( + const std::string& prefix, + bool overwrite = false); + +} // namespace facebook::velox::cudf_velox::presto diff --git a/velox/experimental/cudf/exec/ToCudf.cpp b/velox/experimental/cudf/exec/ToCudf.cpp index 2c1fc91cd41..11a7d2b4e9f 100644 --- a/velox/experimental/cudf/exec/ToCudf.cpp +++ b/velox/experimental/cudf/exec/ToCudf.cpp @@ -15,8 +15,13 @@ */ #include "velox/experimental/cudf/exec/CudfConversion.h" +#include "velox/experimental/cudf/exec/CudfFilterProject.h" #include "velox/experimental/cudf/exec/CudfHashAggregation.h" +#include "velox/experimental/cudf/exec/CudfHashJoin.h" +#include "velox/experimental/cudf/exec/CudfLimit.h" +#include "velox/experimental/cudf/exec/CudfLocalPartition.h" #include "velox/experimental/cudf/exec/CudfOrderBy.h" +#include "velox/experimental/cudf/exec/ExpressionEvaluator.h" #include "velox/experimental/cudf/exec/ToCudf.h" #include "velox/experimental/cudf/exec/Utilities.h" @@ -25,8 +30,10 @@ #include "velox/exec/HashAggregation.h" #include "velox/exec/HashBuild.h" #include "velox/exec/HashProbe.h" +#include "velox/exec/Limit.h" #include "velox/exec/Operator.h" #include "velox/exec/OrderBy.h" +#include "velox/exec/TableScan.h" #include @@ -37,6 +44,7 @@ DEFINE_bool(velox_cudf_enabled, true, "Enable cuDF-Velox acceleration"); DEFINE_string(velox_cudf_memory_resource, "async", "Memory resource for cuDF"); DEFINE_bool(velox_cudf_debug, false, "Enable debug printing"); +DEFINE_bool(velox_cudf_table_scan, true, "Enable cuDF table scan"); namespace facebook::velox::cudf_velox { @@ -82,23 +90,81 @@ bool CompileState::compile() { return driverFactory_.consumerNode; }; - auto isSupportedGpuOperator = [](const exec::Operator* op) { - return isAnyOf(op); + const bool isParquetConnectorRegistered = + facebook::velox::connector::getAllConnectors().count("test-parquet") > 0; + auto isTableScanSupported = + [isParquetConnectorRegistered](const exec::Operator* op) { + return isAnyOf(op) && isParquetConnectorRegistered && + cudfTableScanEnabled(); + }; + + auto isFilterProjectSupported = [](const exec::Operator* op) { + if (auto filterProjectOp = dynamic_cast(op)) { + auto info = filterProjectOp->exprsAndProjection(); + return ExpressionEvaluator::canBeEvaluated(info.exprs->exprs()); + } + return false; }; + auto isJoinSupported = [getPlanNode](const exec::Operator* op) { + if (!isAnyOf(op)) { + return false; + } + auto planNode = std::dynamic_pointer_cast( + getPlanNode(op->planNodeId())); + if (!planNode) { + return false; + } + if (!CudfHashJoinProbe::isSupportedJoinType(planNode->joinType())) { + return false; + } + // disabling null-aware anti join with filter until we implement it right + if (planNode->joinType() == core::JoinType::kAnti and + planNode->isNullAware() and planNode->filter()) { + return false; + } + return true; + }; + + auto isSupportedGpuOperator = + [isFilterProjectSupported, isJoinSupported, isTableScanSupported]( + const exec::Operator* op) { + return isAnyOf< + exec::OrderBy, + exec::HashAggregation, + exec::Limit, + exec::LocalPartition, + exec::LocalExchange>(op) || + isFilterProjectSupported(op) || isJoinSupported(op) || + isTableScanSupported(op); + }; + std::vector isSupportedGpuOperators(operators.size()); std::transform( operators.begin(), operators.end(), isSupportedGpuOperators.begin(), isSupportedGpuOperator); - - auto acceptsGpuInput = [](const exec::Operator* op) { - return isAnyOf(op); + auto acceptsGpuInput = [isFilterProjectSupported, + isJoinSupported](const exec::Operator* op) { + return isAnyOf< + exec::OrderBy, + exec::HashAggregation, + exec::Limit, + exec::LocalPartition>(op) || + isFilterProjectSupported(op) || isJoinSupported(op); }; - - auto producesGpuOutput = [](const exec::Operator* op) { - return isAnyOf(op); + auto producesGpuOutput = [isFilterProjectSupported, + isJoinSupported, + isTableScanSupported](const exec::Operator* op) { + return isAnyOf< + exec::OrderBy, + exec::HashAggregation, + exec::Limit, + exec::LocalExchange>(op) || + isFilterProjectSupported(op) || + (isAnyOf(op) && isJoinSupported(op)) || + (isTableScanSupported(op)); }; int32_t operatorsOffset = 0; @@ -126,7 +192,34 @@ bool CompileState::compile() { replaceOp.back()->initialize(); } - if (auto* orderByOp = dynamic_cast(oper)) { + // This is used to denote if the current operator is kept or replaced. + auto keepOperator = 0; + // TableScan + if (isTableScanSupported(oper)) { + auto planNode = std::dynamic_pointer_cast( + getPlanNode(oper->planNodeId())); + VELOX_CHECK(planNode != nullptr); + keepOperator = 1; + } else if (isJoinSupported(oper)) { + if (auto joinBuildOp = dynamic_cast(oper)) { + auto planNode = std::dynamic_pointer_cast( + getPlanNode(joinBuildOp->planNodeId())); + VELOX_CHECK(planNode != nullptr); + // From-Velox (optional) + replaceOp.push_back( + std::make_unique(id, ctx, planNode)); + replaceOp.back()->initialize(); + } else if (auto joinProbeOp = dynamic_cast(oper)) { + auto planNode = std::dynamic_pointer_cast( + getPlanNode(joinProbeOp->planNodeId())); + VELOX_CHECK(planNode != nullptr); + // From-Velox (optional) + replaceOp.push_back( + std::make_unique(id, ctx, planNode)); + replaceOp.back()->initialize(); + // To-Velox (optional) + } + } else if (auto orderByOp = dynamic_cast(oper)) { auto id = orderByOp->operatorId(); auto planNode = std::dynamic_pointer_cast( getPlanNode(orderByOp->planNodeId())); @@ -140,6 +233,37 @@ bool CompileState::compile() { replaceOp.push_back( std::make_unique(id, ctx, planNode)); replaceOp.back()->initialize(); + } else if (isFilterProjectSupported(oper)) { + auto filterProjectOp = dynamic_cast(oper); + auto info = filterProjectOp->exprsAndProjection(); + auto& idProjections = filterProjectOp->identityProjections(); + auto projectPlanNode = std::dynamic_pointer_cast( + getPlanNode(filterProjectOp->planNodeId())); + auto filterPlanNode = std::dynamic_pointer_cast( + getPlanNode(filterProjectOp->planNodeId())); + // If filter only, filter node only exists. + // If project only, or filter and project, project node only exists. + VELOX_CHECK(projectPlanNode != nullptr or filterPlanNode != nullptr); + replaceOp.push_back(std::make_unique( + id, ctx, info, idProjections, filterPlanNode, projectPlanNode)); + replaceOp.back()->initialize(); + } else if (auto limitOp = dynamic_cast(oper)) { + auto planNode = std::dynamic_pointer_cast( + getPlanNode(limitOp->planNodeId())); + VELOX_CHECK(planNode != nullptr); + replaceOp.push_back(std::make_unique(id, ctx, planNode)); + replaceOp.back()->initialize(); + } else if ( + auto localPartitionOp = dynamic_cast(oper)) { + auto planNode = std::dynamic_pointer_cast( + getPlanNode(localPartitionOp->planNodeId())); + VELOX_CHECK(planNode != nullptr); + replaceOp.push_back( + std::make_unique(id, ctx, planNode)); + replaceOp.back()->initialize(); + } else if ( + auto localExchangeOp = dynamic_cast(oper)) { + keepOperator = 1; } if (producesGpuOutput(oper) and @@ -151,10 +275,10 @@ bool CompileState::compile() { } if (not replaceOp.empty()) { - operatorsOffset += replaceOp.size() - 1; + operatorsOffset += replaceOp.size() - 1 + keepOperator; [[maybe_unused]] auto replaced = driverFactory_.replaceOperators( driver_, - replacingOperatorIndex, + replacingOperatorIndex + keepOperator, replacingOperatorIndex + 1, std::move(replaceOp)); replacementsMade = true; @@ -204,6 +328,9 @@ void registerCudf(const CudfOptions& options) { const std::string mrMode = options.cudfMemoryResource; auto mr = cudf_velox::createMemoryResource(mrMode); cudf::set_current_device_resource(mr.get()); + + exec::Operator::registerOperator( + std::make_unique()); CudfDriverAdapter cda{mr}; exec::DriverAdapter cudfAdapter{kCudfAdapterName, {}, cda}; exec::DriverFactory::registerAdapter(cudfAdapter); @@ -227,4 +354,12 @@ bool cudfIsRegistered() { return isCudfRegistered; } +bool cudfDebugEnabled() { + return FLAGS_velox_cudf_debug; +} + +bool cudfTableScanEnabled() { + return FLAGS_velox_cudf_table_scan; +} + } // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/exec/ToCudf.h b/velox/experimental/cudf/exec/ToCudf.h index 63fcf0d5dd7..433bb27579c 100644 --- a/velox/experimental/cudf/exec/ToCudf.h +++ b/velox/experimental/cudf/exec/ToCudf.h @@ -24,6 +24,7 @@ DECLARE_bool(velox_cudf_enabled); DECLARE_string(velox_cudf_memory_resource); DECLARE_bool(velox_cudf_debug); +DECLARE_bool(velox_cudf_table_scan); namespace facebook::velox::cudf_velox { @@ -61,4 +62,14 @@ void unregisterCudf(); /// Returns true if cuDF is registered. bool cudfIsRegistered(); +/** + * @brief Returns true if the velox_cudf_debug flag is set to true. + */ +bool cudfDebugEnabled(); + +/** + * @brief Returns true if the velox_cudf_table_scan flag is set to true. + */ +bool cudfTableScanEnabled(); + } // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/tests/AggregationTest.cpp b/velox/experimental/cudf/tests/AggregationTest.cpp index 2f405b672e2..5d4489ce969 100644 --- a/velox/experimental/cudf/tests/AggregationTest.cpp +++ b/velox/experimental/cudf/tests/AggregationTest.cpp @@ -14,6 +14,7 @@ * limitations under the License. */ +#include "velox/experimental/cudf/exec/CudfHashAggregation.h" #include "velox/experimental/cudf/exec/ToCudf.h" #include "velox/dwio/common/tests/utils/BatchMaker.h" @@ -479,6 +480,36 @@ TEST_F(AggregationTest, countPartialFinalGlobal) { assertQuery(op, "SELECT count(*) FROM tmp"); } +/// Tests the spark scenario of having different types of aggs in the same +/// planNode Specific example being tested is +/// https://github.com/facebookincubator/velox/issues/12830#issuecomment-2783340233 +TEST_F(AggregationTest, CompanionAggs) { + std::vector keys0{1, 1, 1, 2, 1, 1, 2, 2}; + std::vector keys1{1, 2, 1, 2, 1, 2, 1, 2}; + std::vector values{1, 2, 3, 4, 5, 6, 7, 8}; + auto rowVector = makeRowVector( + {makeFlatVector(keys0), + makeFlatVector(keys1), + makeFlatVector(values)}); + + createDuckDbTable({rowVector}); + + facebook::velox::cudf_velox::registerCudfAggregators( + true /* withCompanionFunctions */, false /* overwrite */); + + auto op = + PlanBuilder() + .values({rowVector}) + .singleAggregation({"c2", "c0"}, {"count_partial(c1)"}) + .localPartition({"c2", "c0"}) + .singleAggregation({"c0"}, {"count_merge(a0)", "count_partial(c2)"}) + .localPartition({"c0"}) + .singleAggregation({"c0"}, {"count_merge(a0)", "count_merge(a1)"}) + .planNode(); + assertQuery( + op, "SELECT c0, count(c1), count(distinct c2) FROM tmp GROUP BY c0"); +} + TEST_F(AggregationTest, partialAggregationMemoryLimit) { auto vectors = { makeRowVector({makeFlatVector( diff --git a/velox/experimental/cudf/tests/CMakeLists.txt b/velox/experimental/cudf/tests/CMakeLists.txt index 52653addd69..2d855abfd3f 100644 --- a/velox/experimental/cudf/tests/CMakeLists.txt +++ b/velox/experimental/cudf/tests/CMakeLists.txt @@ -12,8 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. +add_executable(velox_cudf_hash_join_test HashJoinTest.cpp Main.cpp) add_executable(velox_cudf_order_by_test Main.cpp OrderByTest.cpp) add_executable(velox_cudf_aggregation_test Main.cpp AggregationTest.cpp) +add_executable(velox_cudf_table_scan_test Main.cpp TableScanTest.cpp) +add_executable(velox_cudf_table_write_test Main.cpp TableWriteTest.cpp) +add_executable(velox_cudf_local_partition_test Main.cpp LocalPartitionTest.cpp) +add_executable(velox_cudf_filter_project_test Main.cpp FilterProjectTest.cpp) +add_executable(velox_cudf_limit_test Main.cpp LimitTest.cpp) + +add_test( + NAME velox_cudf_hash_join_test + COMMAND velox_cudf_hash_join_test + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}) add_test( NAME velox_cudf_order_by_test @@ -25,10 +36,59 @@ add_test( COMMAND velox_cudf_aggregation_test WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}) +add_test( + NAME velox_cudf_local_partition_test + COMMAND velox_cudf_local_partition_test + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}) + +add_test( + NAME velox_cudf_table_scan_test + COMMAND velox_cudf_table_scan_test + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}) + +add_test( + NAME velox_cudf_table_write_test + COMMAND velox_cudf_table_write_test + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}) + +add_test( + NAME velox_cudf_filter_project_test + COMMAND velox_cudf_filter_project_test + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}) + +add_test( + NAME velox_cudf_limit_test + COMMAND velox_cudf_limit_test + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}) + +set_tests_properties(velox_cudf_hash_join_test PROPERTIES LABELS cuda_driver + TIMEOUT 3000) set_tests_properties(velox_cudf_order_by_test PROPERTIES LABELS cuda_driver TIMEOUT 3000) set_tests_properties(velox_cudf_aggregation_test PROPERTIES LABELS cuda_driver TIMEOUT 3000) +set_tests_properties(velox_cudf_local_partition_test + PROPERTIES LABELS cuda_driver TIMEOUT 3000) +set_tests_properties(velox_cudf_table_scan_test PROPERTIES LABELS cuda_driver + TIMEOUT 3000) +set_tests_properties(velox_cudf_table_write_test PROPERTIES LABELS cuda_driver + TIMEOUT 3000) +set_tests_properties(velox_cudf_filter_project_test + PROPERTIES LABELS cuda_driver TIMEOUT 3000) +set_tests_properties(velox_cudf_limit_test PROPERTIES LABELS cuda_driver + TIMEOUT 3000) + +target_link_libraries( + velox_cudf_hash_join_test + velox_cudf_exec + velox_exec + velox_exec_test_lib + velox_test_util + velox_vector_fuzzer + gtest + gtest_main + Folly::folly + fmt::fmt) target_link_libraries( velox_cudf_order_by_test @@ -50,3 +110,55 @@ target_link_libraries( gtest gtest_main fmt::fmt) + +target_link_libraries( + velox_cudf_local_partition_test + velox_cudf_exec + velox_exec + velox_exec_test_lib + velox_test_util + gtest + gtest_main + fmt::fmt) + +target_link_libraries( + velox_cudf_table_scan_test + velox_cudf_exec_test_lib + velox_cudf_parquet_connector + velox_exec + velox_exec_test_lib + velox_test_util + gtest + gtest_main + fmt::fmt) + +target_link_libraries( + velox_cudf_table_write_test + velox_cudf_exec_test_lib + velox_cudf_parquet_connector + velox_exec + velox_exec_test_lib + velox_test_util + gtest + gtest_main + fmt::fmt) + +target_link_libraries( + velox_cudf_filter_project_test + velox_cudf_exec + velox_exec + velox_exec_test_lib + velox_test_util + gtest + gtest_main) + +target_link_libraries( + velox_cudf_limit_test + velox_cudf_exec + velox_exec + velox_exec_test_lib + velox_test_util + gtest + gtest_main) + +add_subdirectory(utils) diff --git a/velox/experimental/cudf/tests/FilterProjectTest.cpp b/velox/experimental/cudf/tests/FilterProjectTest.cpp new file mode 100644 index 00000000000..8c83ba4f2b6 --- /dev/null +++ b/velox/experimental/cudf/tests/FilterProjectTest.cpp @@ -0,0 +1,777 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ +#include "velox/experimental/cudf/exec/CudfFilterProject.h" +#include "velox/experimental/cudf/exec/ToCudf.h" + +#include "velox/common/base/tests/GTestUtils.h" +#include "velox/dwio/common/tests/utils/BatchMaker.h" +#include "velox/exec/tests/utils/OperatorTestBase.h" +#include "velox/exec/tests/utils/PlanBuilder.h" + +using namespace facebook::velox; +using namespace facebook::velox::exec; +using namespace facebook::velox::exec::test; +using namespace facebook::velox::common::testutil; + +namespace { + +template +T getColValue(const std::vector& input, int col, int32_t index) { + return input[0]->as()->childAt(col)->as>()->valueAt( + index); +} + +class CudfFilterProjectTest : public OperatorTestBase { + protected: + void SetUp() override { + OperatorTestBase::SetUp(); + filesystems::registerLocalFileSystem(); + cudf_velox::registerCudf(); + rng_.seed(123); + + rowType_ = ROW({{"c0", INTEGER()}, {"c1", DOUBLE()}, {"c2", VARCHAR()}}); + } + + void TearDown() override { + cudf_velox::unregisterCudf(); + OperatorTestBase::TearDown(); + } + + void testMultiplyOperation(const std::vector& input) { + // Create a plan with a multiply operation + auto plan = + PlanBuilder().values(input).project({"1.0 * c1 AS result"}).planNode(); + + // Run the test + runTest(plan, "SELECT 1.0 * c1 AS result FROM tmp"); + } + + void testDivideOperation(const std::vector& input) { + // Create a plan with a divide operation + auto plan = + PlanBuilder().values(input).project({"c0 / c1 AS result"}).planNode(); + + // Run the test + runTest(plan, "SELECT c0 / c1 AS result FROM tmp"); + } + + void testMultiplyAndMinusOperation(const std::vector& input) { + // Create a plan with a multiply and minus operation + auto plan = PlanBuilder() + .values(input) + .project({"c0 * (1.0 - c1) AS result"}) + .planNode(); + + // Run the test + runTest(plan, "SELECT c0 * (1.0 - c1) AS result FROM tmp"); + } + + void testStringEqualOperation(const std::vector& input) { + // Create a plan with a string equal operation + auto c2Value = input[0] + ->as() + ->childAt(2) + ->as>() + ->valueAt(1) + .str(); + auto plan = PlanBuilder() + .values(input) + .project({"c2 = '" + c2Value + "' AS result"}) + .planNode(); + + // Run the test + runTest(plan, "SELECT c2 = '" + c2Value + "' AS result FROM tmp"); + } + + void testStringNotEqualOperation(const std::vector& input) { + // Create a plan with a string not equal operation + auto c2Value = input[0] + ->as() + ->childAt(2) + ->as>() + ->valueAt(1) + .str(); + auto plan = PlanBuilder() + .values(input) + .project({"c2 <> '" + c2Value + "' AS result"}) + .planNode(); + + // Run the test + runTest(plan, "SELECT c2 <> '" + c2Value + "' AS result FROM tmp"); + } + + void testAndOperation(const std::vector& input) { + // Create a plan with AND operation + auto plan = PlanBuilder() + .values(input) + .project({"c0 = 1 AND c1 = 2.0 AS result"}) + .planNode(); + + // Run the test + runTest(plan, "SELECT c0 = 1 AND c1 = 2.0 AS result FROM tmp"); + } + + void testOrOperation(const std::vector& input) { + // Create a plan with OR operation + auto plan = PlanBuilder() + .values(input) + .project({"c0 = 1 OR c1 = 2.0 AS result"}) + .planNode(); + + // Run the test + runTest(plan, "SELECT c0 = 1 OR c1 = 2.0 AS result FROM tmp"); + } + + void testYearFunction(const std::vector& input) { + // Create a plan with YEAR function + auto plan = + PlanBuilder().values(input).project({"YEAR(c2) AS result"}).planNode(); + + // Run the test + runTest(plan, "SELECT YEAR(c2) AS result FROM tmp"); + } + + void testLengthFunction(const std::vector& input) { + // Create a plan with LENGTH function + auto plan = PlanBuilder() + .values(input) + .project({"LENGTH(c2) AS result"}) + .planNode(); + + // Run the test + runTest(plan, "SELECT LENGTH(c2) AS result FROM tmp"); + } + + void testCaseWhenOperation(const std::vector& input) { + // Create a plan with a CASE WHEN operation + auto plan = + PlanBuilder() + .values(input) + .project({"CASE WHEN c0 = 0 THEN 1.0 ELSE 0.0 END AS result"}) + .planNode(); + + // Run the test + runTest( + plan, + "SELECT CASE WHEN c0 = 0 THEN 1.0 ELSE 0.0 END AS result FROM tmp"); + } + + void testSubstrOperation(const std::vector& input) { + // Create a plan with a substr operation + auto plan = PlanBuilder() + .values(input) + .project({"substr(c2, 1, 3) AS result"}) + .planNode(); + + // Run the test + runTest(plan, "SELECT substr(c2, 1, 3) AS result FROM tmp"); + } + + void testLikeOperation(const std::vector& input) { + // Create a plan with a like operation + auto plan = PlanBuilder() + .values(input) + .project({"c2 LIKE '%test%' AS result"}) + .planNode(); + + // Run the test + runTest(plan, "SELECT c2 LIKE '%test%' AS result FROM tmp"); + } + + void testLessThanOperation(const std::vector& input) { + // Create a plan with a less than operation + auto plan = + PlanBuilder().values(input).project({"c0 < c1 AS result"}).planNode(); + + // Run the test + runTest(plan, "SELECT c0 < c1 AS result FROM tmp"); + + // compare against literals + plan = PlanBuilder().values(input).project({"c0 < 1 AS result"}).planNode(); + + // Run the test + runTest(plan, "SELECT c0 < 1 AS result FROM tmp"); + } + + void testGreaterThanOperation(const std::vector& input) { + // Create a plan with a greater than operation + auto plan = + PlanBuilder().values(input).project({"c0 > c1 AS result"}).planNode(); + + // Run the test + runTest(plan, "SELECT c0 > c1 AS result FROM tmp"); + + // compare against literals + plan = PlanBuilder().values(input).project({"c0 > 1 AS result"}).planNode(); + + // Run the test + runTest(plan, "SELECT c0 > 1 AS result FROM tmp"); + } + + void testLessThanEqualOperation(const std::vector& input) { + // Create a plan with a less than equal operation + auto plan = + PlanBuilder().values(input).project({"c0 <= c1 AS result"}).planNode(); + + // Run the test + runTest(plan, "SELECT c0 <= c1 AS result FROM tmp"); + } + + void testGreaterThanEqualOperation(const std::vector& input) { + // Create a plan with a greater than equal operation + auto plan = + PlanBuilder().values(input).project({"c0 >= c1 AS result"}).planNode(); + + // Run the test + runTest(plan, "SELECT c0 >= c1 AS result FROM tmp"); + } + + void testNotOperation(const std::vector& input) { + // Create a plan with a NOT operation + auto plan = PlanBuilder() + .values(input) + .project({"NOT (c0 = 1) AS result"}) + .planNode(); + + // Run the test + runTest(plan, "SELECT NOT (c0 = 1) AS result FROM tmp"); + } + + void testBetweenOperation(const std::vector& input) { + // Create a plan with a BETWEEN operation + auto plan = PlanBuilder() + .values(input) + .project({"c0 BETWEEN 1 AND 100 AS result"}) + .planNode(); + + // Run the test + runTest(plan, "SELECT c0 BETWEEN 1 AND 100 AS result FROM tmp"); + } + + void testMultiInputAndOperation(const std::vector& input) { + // Create a plan with multiple AND operations + auto c2Value = getColValue(input, 2, 1).str(); + auto plan = PlanBuilder() + .values(input) + .project( + {"c0 > 1000 AND c0 < 20000 AND c2 = '" + c2Value + + "' AS result"}) + .planNode(); + + // Run the test + runTest( + plan, + "SELECT c0 > 1000 AND c0 < 20000 AND c2 = '" + c2Value + + "' AS result FROM tmp"); + } + + void testMultiInputOrOperation(const std::vector& input) { + // Create a plan with multiple OR operations + auto c2Value = getColValue(input, 2, 1).str(); + auto plan = PlanBuilder() + .values(input) + .project( + {"c0 > 16000 OR c0 < 8000 OR c1 = 2.0 OR c2 = '" + + c2Value + "' AS result"}) + .planNode(); + + // Run the test + runTest( + plan, + "SELECT c0 > 16000 OR c0 < 8000 OR c1 = 2.0 OR c2 = '" + c2Value + + "' AS result FROM tmp"); + } + + void testIntegerInOperation(const std::vector& input) { + // Create a plan with an IN operation for integers + std::vector c0Values; + for (int32_t i = 0; i < 5; i++) { + c0Values.push_back(getColValue(input, 0, i)); + } + std::string c0ValuesStr; + for (size_t i = 0; i < c0Values.size(); ++i) { + c0ValuesStr += std::to_string(c0Values[i]) + ","; + } + c0ValuesStr.pop_back(); + auto plan = PlanBuilder(pool_.get()) + .values(input) + .project({"c0 IN (" + c0ValuesStr + ") AS result"}) + .planNode(); + + // Run the test + runTest(plan, "SELECT c0 IN (" + c0ValuesStr + ") AS result FROM tmp"); + } + + void testDoubleInOperation(const std::vector& input) { + // Create a plan with an IN operation for doubles + std::vector c1Values; + for (int32_t i = 0; i < 4; i++) { + c1Values.push_back(getColValue(input, 1, i)); + } + std::string c1ValuesStr; + for (size_t i = 0; i < c1Values.size(); ++i) { + c1ValuesStr += std::to_string(c1Values[i]) + ","; + } + c1ValuesStr.pop_back(); + auto plan = PlanBuilder(pool_.get()) + .values(input) + .project({"c1 IN (" + c1ValuesStr + ") AS result"}) + .planNode(); + + // Run the test + runTest(plan, "SELECT c1 IN (" + c1ValuesStr + ") AS result FROM tmp"); + } + + void testStringInOperation(const std::vector& input) { + // Create a plan with an IN operation for strings + std::vector c2Values; + for (int32_t i = 0; i < 3; i++) { + c2Values.push_back(getColValue(input, 2, i)); + } + std::string c2ValuesStr; + for (size_t i = 0; i < c2Values.size(); ++i) { + c2ValuesStr += "'" + c2Values[i].str() + "',"; + } + c2ValuesStr.pop_back(); + auto plan = PlanBuilder(pool_.get()) + .values(input) + .project({"c2 IN (" + c2ValuesStr + ") AS result"}) + .planNode(); + + // Run the test + runTest(plan, "SELECT c2 IN (" + c2ValuesStr + ") AS result FROM tmp"); + } + + void testMixedInOperation(const std::vector& input) { + // Create a plan that combines multiple IN operations + auto plan = + PlanBuilder(pool_.get()) + .values(input) + .project( + {"c0 IN (1, 2, 3) OR c1 IN (1.5, 2.5) OR c2 IN ('test1', 'test2') AS result"}) + .planNode(); + + // Run the test + runTest( + plan, + "SELECT c0 IN (1, 2, 3) OR c1 IN (1.5, 2.5) OR c2 IN ('test1', 'test2') AS result FROM tmp"); + } + + void runTest(core::PlanNodePtr planNode, const std::string& duckDbSql) { + SCOPED_TRACE("run without spilling"); + assertQuery(planNode, duckDbSql); + } + + std::vector makeVectors( + const RowTypePtr& rowType, + int32_t numVectors, + int32_t rowsPerVector) { + std::vector vectors; + for (int32_t i = 0; i < numVectors; ++i) { + auto vector = std::dynamic_pointer_cast( + facebook::velox::test::BatchMaker::createBatch( + rowType, rowsPerVector, *pool_)); + vectors.push_back(vector); + } + return vectors; + } + + folly::Random::DefaultGenerator rng_; + RowTypePtr rowType_; +}; + +TEST_F(CudfFilterProjectTest, multiplyOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testMultiplyOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, divideOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testDivideOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, multiplyAndMinusOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testMultiplyAndMinusOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, stringEqualOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testStringEqualOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, stringNotEqualOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testStringNotEqualOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, andOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testAndOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, orOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testOrOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, lengthFunction) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testLengthFunction(vectors); +} + +TEST_F(CudfFilterProjectTest, yearFunction) { + // Update row type to use TIMESTAMP directly + auto rowType = + ROW({{"c0", INTEGER()}, {"c1", DOUBLE()}, {"c2", TIMESTAMP()}}); + + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType, 2, batchSize); + + // Set timestamp values directly + for (auto& vector : vectors) { + auto timestampVector = vector->childAt(2)->asFlatVector(); + for (vector_size_t i = 0; i < batchSize; ++i) { + // Set to 2024-03-14 12:34:56 + Timestamp ts(1710415496, 0); // seconds, nanos + timestampVector->set(i, ts); + } + } + + createDuckDbTable(vectors); + testYearFunction(vectors); +} + +TEST_F(CudfFilterProjectTest, DISABLED_caseWhenOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + // failing because switch copies nulls too. + createDuckDbTable(vectors); + + testCaseWhenOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, substrOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testSubstrOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, likeOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testLikeOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, lessThanOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testLessThanOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, greaterThanOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testGreaterThanOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, lessThanEqualOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testLessThanEqualOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, greaterThanEqualOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testGreaterThanEqualOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, notOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testNotOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, betweenOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testBetweenOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, multiInputAndOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testMultiInputAndOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, multiInputOrOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testMultiInputOrOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, integerInOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testIntegerInOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, doubleInOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testDoubleInOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, stringInOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testStringInOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, mixedInOperation) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + testMixedInOperation(vectors); +} + +TEST_F(CudfFilterProjectTest, simpleFilter) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + // Create a plan with a simple filter + auto plan = PlanBuilder() + .values(vectors) + .filter("c0 > 500") + .project({"c0", "c1", "c2"}) + .planNode(); + + // Run the test + assertQuery(plan, "SELECT c0, c1, c2 FROM tmp WHERE c0 > 500"); +} + +TEST_F(CudfFilterProjectTest, filterWithProject) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + // Create a plan with filter and project + auto plan = + PlanBuilder() + .values(vectors) + .filter("c0 > 500") + .project({"c0 + 2 as doubled", "c1 + 1.0 as incremented", "c2"}) + .planNode(); + + // Run the test + assertQuery( + plan, + "SELECT c0 + 2 as doubled, c1 + 1.0 as incremented, c2 FROM tmp WHERE c0 > 500"); +} + +TEST_F(CudfFilterProjectTest, complexFilter) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + // Create a plan with a complex filter condition + auto plan = PlanBuilder() + .values(vectors) + .filter("c0 > 500 AND c1 < 0.5 AND c2 LIKE '%test%'") + .project({"c0", "c1", "c2"}) + .planNode(); + + // Run the test + assertQuery( + plan, + "SELECT c0, c1, c2 FROM tmp WHERE c0 > 500 AND c1 < 0.5 AND c2 LIKE '%test%'"); +} + +TEST_F(CudfFilterProjectTest, filterWithNullValues) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + + // Add some null values to the vectors + for (auto& vector : vectors) { + auto c0Vector = vector->childAt(0)->asFlatVector(); + auto c1Vector = vector->childAt(1)->asFlatVector(); + for (vector_size_t i = 0; i < batchSize; i += 10) { + c0Vector->setNull(i, true); + c1Vector->setNull(i, true); + } + } + + createDuckDbTable(vectors); + + // Create a plan with filter that handles null values + auto plan = PlanBuilder() + .values(vectors) + .filter("c0 IS NOT NULL AND c1 IS NOT NULL") + .project({"c0", "c1", "c2"}) + .planNode(); + + // Run the test + assertQuery( + plan, + "SELECT c0, c1, c2 FROM tmp WHERE c0 IS NOT NULL AND c1 IS NOT NULL"); +} + +TEST_F(CudfFilterProjectTest, filterWithOrCondition) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + // Create a plan with OR condition in filter + auto plan = PlanBuilder() + .values(vectors) + .filter("c0 > 500 OR c1 < 0.5") + .project({"c0", "c1", "c2"}) + .planNode(); + + // Run the test + assertQuery(plan, "SELECT c0, c1, c2 FROM tmp WHERE c0 > 500 OR c1 < 0.5"); +} + +TEST_F(CudfFilterProjectTest, filterWithInCondition) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + // Create a plan with IN condition in filter + auto plan = PlanBuilder(pool_.get()) + .values(vectors) + .filter("c0 IN (100, 200, 300, 400, 500)") + .project({"c0", "c1", "c2"}) + .planNode(); + + // Run the test + assertQuery( + plan, "SELECT c0, c1, c2 FROM tmp WHERE c0 IN (100, 200, 300, 400, 500)"); +} + +TEST_F(CudfFilterProjectTest, filterWithBetweenCondition) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + // Create a plan with BETWEEN condition in filter + auto plan = PlanBuilder() + .values(vectors) + .filter("c0 BETWEEN 100 AND 500") + .project({"c0", "c1", "c2"}) + .planNode(); + + // Run the test + assertQuery(plan, "SELECT c0, c1, c2 FROM tmp WHERE c0 BETWEEN 100 AND 500"); +} + +TEST_F(CudfFilterProjectTest, filterWithStringOperations) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + // Create a plan with string operations in filter + auto plan = PlanBuilder() + .values(vectors) + .filter("LENGTH(c2) > 5") + .project({"c0", "c1", "c2"}) + .planNode(); + + // Run the test + assertQuery(plan, "SELECT c0, c1, c2 FROM tmp WHERE LENGTH(c2) > 5"); +} + +TEST_F(CudfFilterProjectTest, filterWithoutProject) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + // Create a plan with only filter (no projection) + auto plan = + PlanBuilder().values(vectors).filter("c0 > 500 AND c1 < 0.5").planNode(); + + // Run the test - should return all columns without modification + assertQuery(plan, "SELECT c0, c1, c2 FROM tmp WHERE c0 > 500 AND c1 < 0.5"); +} + +TEST_F(CudfFilterProjectTest, filterWithEmptyResult) { + vector_size_t batchSize = 1000; + auto vectors = makeVectors(rowType_, 2, batchSize); + createDuckDbTable(vectors); + + // Create a plan with a filter that should return no rows + auto plan = PlanBuilder() + .values(vectors) + .filter("c0 < 0 AND c0 > 1000") // Impossible condition + .planNode(); + + // Run the test - should return empty result + assertQuery(plan, "SELECT c0, c1, c2 FROM tmp WHERE c0 < 0 AND c0 > 1000"); +} +} // namespace diff --git a/velox/experimental/cudf/tests/HashJoinTest.cpp b/velox/experimental/cudf/tests/HashJoinTest.cpp new file mode 100644 index 00000000000..520087e62d2 --- /dev/null +++ b/velox/experimental/cudf/tests/HashJoinTest.cpp @@ -0,0 +1,8252 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#include "velox/experimental/cudf/exec/ToCudf.h" + +#include "folly/experimental/EventCount.h" +#include "velox/common/base/tests/GTestUtils.h" +#include "velox/common/testutil/TestValue.h" +#include "velox/dwio/common/tests/utils/BatchMaker.h" +#include "velox/exec/Cursor.h" +#include "velox/exec/HashBuild.h" +#include "velox/exec/HashJoinBridge.h" +#include "velox/exec/OperatorUtils.h" +#include "velox/exec/PlanNodeStats.h" +#include "velox/exec/tests/utils/ArbitratorTestUtil.h" +#include "velox/exec/tests/utils/AssertQueryBuilder.h" +#include "velox/exec/tests/utils/HashJoinTestBase.h" +#include "velox/exec/tests/utils/HiveConnectorTestBase.h" +#include "velox/exec/tests/utils/PlanBuilder.h" +#include "velox/exec/tests/utils/TempDirectoryPath.h" +#include "velox/exec/tests/utils/VectorTestUtil.h" +#include "velox/vector/fuzzer/VectorFuzzer.h" + +#include +#include + +using namespace facebook::velox; +using namespace facebook::velox::exec; +using namespace facebook::velox::exec::test; +using namespace facebook::velox::common::testutil; + +using facebook::velox::test::BatchMaker; + +namespace { + +class HashJoinTest : public HashJoinTestBase { + public: + HashJoinTest() : HashJoinTestBase(TestParam(1)) {} + + explicit HashJoinTest(const TestParam& param) : HashJoinTestBase(param) {} + + void SetUp() override { + HashJoinTestBase::SetUp(); + cudf_velox::registerCudf(); + } + + void TearDown() override { + cudf_velox::unregisterCudf(); + HashJoinTestBase::TearDown(); + } +}; + +class MultiThreadedHashJoinTest + : public HashJoinTest, + public testing::WithParamInterface { + public: + MultiThreadedHashJoinTest() : HashJoinTest(GetParam()) {} + + static std::vector getTestParams() { + return std::vector({TestParam{1}, TestParam{3}}); + } +}; + +TEST_P(MultiThreadedHashJoinTest, bigintArray) { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .numDrivers(numDrivers_) + .keyTypes({BIGINT()}) + .probeVectors(1600, 5) + .buildVectors(1500, 5) + .referenceQuery( + "SELECT t_k0, t_data, u_k0, u_data FROM t, u WHERE t.t_k0 = u.u_k0") + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, outOfJoinKeyColumnOrder) { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .numDrivers(numDrivers_) + .probeType(probeType_) + .probeKeys({"t_k2"}) + .probeVectors(5, 10) + .buildType(buildType_) + .buildKeys({"u_k2"}) + .buildVectors(64, 15) + .joinOutputLayout({"t_k1", "t_k2", "u_k1", "u_k2", "u_v1"}) + .referenceQuery( + "SELECT t_k1, t_k2, u_k1, u_k2, u_v1 FROM t, u WHERE t_k2 = u_k2") + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, joinWithCancellation) { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .numDrivers(numDrivers_) + .keyTypes({BIGINT()}) + .probeVectors(1600, 5) + .buildVectors(1500, 5) + .injectTaskCancellation(true) + .referenceQuery( + "SELECT t_k0, t_data, u_k0, u_data FROM t, u WHERE t.t_k0 = u.u_k0") + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + auto stats = task->taskStats(); + EXPECT_GT(stats.terminationTimeMs, 0); + }) + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, testJoinWithSpillenabledCancellation) { + auto spillDirectory = exec::test::TempDirectoryPath::create(); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .keyTypes({BIGINT()}) + .probeVectors(1600, 5) + .buildVectors(1500, 5) + .injectTaskCancellation(true) + .injectSpill(false) + // Need spill directory so that canSpill() is true for HashProbe + .spillDirectory(spillDirectory->getPath()) + .referenceQuery( + "SELECT t_k0, t_data, u_k0, u_data FROM t, u WHERE t.t_k0 = u.u_k0") + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, emptyBuild) { + const std::vector finishOnEmptys = {false, true}; + for (const auto finishOnEmpty : finishOnEmptys) { + SCOPED_TRACE(fmt::format("finishOnEmpty: {}", finishOnEmpty)); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .hashProbeFinishEarlyOnEmptyBuild(finishOnEmpty) + .numDrivers(numDrivers_) + .keyTypes({BIGINT()}) + .probeVectors(1600, 5) + .buildVectors(0, 5) + .referenceQuery( + "SELECT t_k0, t_data, u_k0, u_data FROM t, u WHERE t_k0 = u_k0") + .checkSpillStats(false) + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + const auto statsPair = taskSpilledStats(*task); + ASSERT_EQ(statsPair.first.spilledRows, 0); + ASSERT_EQ(statsPair.first.spilledBytes, 0); + ASSERT_EQ(statsPair.first.spilledPartitions, 0); + ASSERT_EQ(statsPair.first.spilledFiles, 0); + ASSERT_EQ(statsPair.second.spilledRows, 0); + ASSERT_EQ(statsPair.second.spilledBytes, 0); + ASSERT_EQ(statsPair.second.spilledPartitions, 0); + ASSERT_EQ(statsPair.second.spilledFiles, 0); + verifyTaskSpilledRuntimeStats(*task, false); + // Check the hash probe has processed probe input rows. + if (finishOnEmpty) { + ASSERT_EQ(getInputPositions(task, 1), 0); + } else { + ASSERT_GT(getInputPositions(task, 1), 0); + } + }) + .run(); + } +} + +TEST_P(MultiThreadedHashJoinTest, emptyProbe) { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .numDrivers(numDrivers_) + .keyTypes({BIGINT()}) + .probeVectors(0, 5) + .buildVectors(1500, 5) + .checkSpillStats(false) + .referenceQuery( + "SELECT t_k0, t_data, u_k0, u_data FROM t, u WHERE t_k0 = u_k0") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + const auto statsPair = taskSpilledStats(*task); + if (hasSpill) { + ASSERT_GT(statsPair.first.spilledRows, 0); + ASSERT_GT(statsPair.first.spilledBytes, 0); + ASSERT_GT(statsPair.first.spilledPartitions, 0); + ASSERT_GT(statsPair.first.spilledFiles, 0); + // There is no spilling at empty probe side. + ASSERT_EQ(statsPair.second.spilledRows, 0); + ASSERT_EQ(statsPair.second.spilledBytes, 0); + ASSERT_GT(statsPair.second.spilledPartitions, 0); + ASSERT_EQ(statsPair.second.spilledFiles, 0); + } else { + ASSERT_EQ(statsPair.first.spilledRows, 0); + ASSERT_EQ(statsPair.first.spilledBytes, 0); + ASSERT_EQ(statsPair.first.spilledPartitions, 0); + ASSERT_EQ(statsPair.first.spilledFiles, 0); + ASSERT_EQ(statsPair.second.spilledRows, 0); + ASSERT_EQ(statsPair.second.spilledBytes, 0); + ASSERT_EQ(statsPair.second.spilledPartitions, 0); + ASSERT_EQ(statsPair.second.spilledFiles, 0); + verifyTaskSpilledRuntimeStats(*task, false); + } + }) + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, normalizedKey) { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .numDrivers(numDrivers_) + .keyTypes({BIGINT(), VARCHAR()}) + .probeVectors(1600, 5) + .buildVectors(1500, 5) + .referenceQuery( + "SELECT t_k0, t_k1, t_data, u_k0, u_k1, u_data FROM t, u WHERE t_k0 = u_k0 AND t_k1 = u_k1") + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, normalizedKeyOverflow) { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .keyTypes({BIGINT(), VARCHAR(), BIGINT(), BIGINT(), BIGINT(), BIGINT()}) + .probeVectors(1600, 5) + .buildVectors(1500, 5) + .referenceQuery( + "SELECT t_k0, t_k1, t_k2, t_k3, t_k4, t_k5, t_data, u_k0, u_k1, u_k2, u_k3, u_k4, u_k5, u_data FROM t, u WHERE t_k0 = u_k0 AND t_k1 = u_k1 AND t_k2 = u_k2 AND t_k3 = u_k3 AND t_k4 = u_k4 AND t_k5 = u_k5") + .run(); +} + +DEBUG_ONLY_TEST_P(MultiThreadedHashJoinTest, parallelJoinBuildCheck) { + std::atomic isParallelBuild{false}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::HashTable::parallelJoinBuild", + std::function([&](void*) { isParallelBuild = true; })); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .keyTypes({BIGINT(), VARCHAR()}) + .probeVectors(1600, 5) + .buildVectors(1500, 5) + .referenceQuery( + "SELECT t_k0, t_k1, t_data, u_k0, u_k1, u_data FROM t, u WHERE t_k0 = u_k0 AND t_k1 = u_k1") + .injectSpill(false) + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + auto joinStats = task->taskStats() + .pipelineStats.back() + .operatorStats.back() + .runtimeStats; + ASSERT_GT(joinStats["hashtable.buildWallNanos"].sum, 0); + ASSERT_GE(joinStats["hashtable.buildWallNanos"].count, 1); + }) + .run(); + ASSERT_EQ(numDrivers_ == 1, !isParallelBuild); +} + +DEBUG_ONLY_TEST_P( + MultiThreadedHashJoinTest, + raceBetweenTaskTerminateAndTableBuild) { + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::HashBuild::finishHashBuild", + std::function([&](Operator* op) { + auto task = op->operatorCtx()->task(); + task->requestAbort(); + })); + VELOX_ASSERT_THROW( + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .keyTypes({BIGINT(), VARCHAR()}) + .probeVectors(1600, 5) + .buildVectors(1500, 5) + .referenceQuery( + "SELECT t_k0, t_k1, t_data, u_k0, u_k1, u_data FROM t, u WHERE t_k0 = u_k0 AND t_k1 = u_k1") + .injectSpill(false) + .run(), + "Aborted for external error"); +} + +TEST_P(MultiThreadedHashJoinTest, allTypes) { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .keyTypes( + {BIGINT(), + VARCHAR(), + REAL(), + DOUBLE(), + INTEGER(), + SMALLINT(), + TINYINT()}) + .probeVectors(1600, 5) + .buildVectors(1500, 5) + .referenceQuery( + "SELECT t_k0, t_k1, t_k2, t_k3, t_k4, t_k5, t_k6, t_data, u_k0, u_k1, u_k2, u_k3, u_k4, u_k5, u_k6, u_data FROM t, u WHERE t_k0 = u_k0 AND t_k1 = u_k1 AND t_k2 = u_k2 AND t_k3 = u_k3 AND t_k4 = u_k4 AND t_k5 = u_k5 AND t_k6 = u_k6") + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, filter) { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .numDrivers(numDrivers_) + .keyTypes({BIGINT()}) + .probeVectors(1600, 5) + .buildVectors(1500, 5) + .joinFilter("((t_k0 % 100) + (u_k0 % 100)) % 40 < 20") + .referenceQuery( + "SELECT t_k0, t_data, u_k0, u_data FROM t, u WHERE t_k0 = u_k0 AND ((t_k0 % 100) + (u_k0 % 100)) % 40 < 20") + .run(); +} + +DEBUG_ONLY_TEST_P(MultiThreadedHashJoinTest, filterSpillOnFirstProbeInput) { + auto spillDirectory = exec::test::TempDirectoryPath::create(); + std::atomic_bool injectProbeSpillOnce{true}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::getOutput", + std::function([&](Operator* op) { + if (!isHashProbeMemoryPool(*op->pool())) { + return; + } + HashProbe* probeOp = static_cast(op); + if (!probeOp->testingHasPendingInput()) { + return; + } + if (!injectProbeSpillOnce.exchange(false)) { + return; + } + testingRunArbitration(op->pool()); + ASSERT_EQ(op->pool()->usedBytes(), 40960); + ASSERT_EQ(op->pool()->reservedBytes(), 1048576); + })); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .keyTypes({BIGINT()}) + .numDrivers(1) + .probeVectors(1600, 5) + .buildVectors(1500, 5) + .injectSpill(false) + .spillDirectory(spillDirectory->getPath()) + .joinFilter("((t_k0 % 100) + (u_k0 % 100)) % 40 < 20") + .referenceQuery( + "SELECT t_k0, t_data, u_k0, u_data FROM t, u WHERE t_k0 = u_k0 AND ((t_k0 % 100) + (u_k0 % 100)) % 40 < 20") + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + const auto statsPair = taskSpilledStats(*task); + ASSERT_EQ(statsPair.first.spilledRows, 0); + ASSERT_EQ(statsPair.first.spilledBytes, 0); + ASSERT_EQ(statsPair.first.spilledPartitions, 0); + ASSERT_EQ(statsPair.first.spilledFiles, 0); + ASSERT_GT(statsPair.second.spilledRows, 0); + ASSERT_GT(statsPair.second.spilledBytes, 0); + ASSERT_GT(statsPair.second.spilledPartitions, 0); + ASSERT_GT(statsPair.second.spilledFiles, 0); + }) + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, nullAwareAntiJoinWithNull) { + struct { + double probeNullRatio; + double buildNullRatio; + + std::string debugString() const { + return fmt::format( + "probeNullRatio: {}, buildNullRatio: {}", + probeNullRatio, + buildNullRatio); + } + } testSettings[] = { + {0.0, 1.0}, {0.0, 0.1}, {0.1, 1.0}, {0.1, 0.1}, {1.0, 1.0}, {1.0, 0.1}}; + for (const auto& testData : testSettings) { + SCOPED_TRACE(testData.debugString()); + + std::vector probeVectors = + makeBatches(5, 3, probeType_, pool_.get(), testData.probeNullRatio); + + // The first half number of build batches having no nulls to trigger it + // later during the processing. + std::vector buildVectors = mergeBatches( + makeBatches(5, 6, buildType_, pool_.get(), 0.0), + makeBatches(5, 6, buildType_, pool_.get(), testData.buildNullRatio)); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .probeType(probeType_) + .probeKeys({"t_k2"}) + .probeVectors(std::move(probeVectors)) + .buildType(buildType_) + .buildKeys({"u_k2"}) + .buildVectors(std::move(buildVectors)) + .joinType(core::JoinType::kAnti) + .nullAware(true) + .joinOutputLayout({"t_k1", "t_k2"}) + .referenceQuery( + "SELECT t_k1, t_k2 FROM t WHERE t.t_k2 NOT IN (SELECT u_k2 FROM u)") + // NOTE: we might not trigger spilling at build side if we detect the + // null join key in the build rows early. + .checkSpillStats(false) + .run(); + } +} + +TEST_P(MultiThreadedHashJoinTest, rightSemiJoinFilterWithLargeOutput) { + // Build the identical left and right vectors to generate large join + // outputs. + std::vector probeVectors = + makeBatches(4, [&](uint32_t /*unused*/) { + return makeRowVector( + {"t0", "t1"}, + {makeFlatVector(2048, [](auto row) { return row; }), + makeFlatVector(2048, [](auto row) { return row; })}); + }); + + std::vector buildVectors = + makeBatches(4, [&](uint32_t /*unused*/) { + return makeRowVector( + {"u0", "u1"}, + {makeFlatVector(2048, [](auto row) { return row; }), + makeFlatVector(2048, [](auto row) { return row; })}); + }); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .probeKeys({"t0"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u0"}) + .buildVectors(std::move(buildVectors)) + .joinType(core::JoinType::kRightSemiFilter) + .joinOutputLayout({"u1"}) + .referenceQuery("SELECT u.u1 FROM u WHERE u.u0 IN (SELECT t0 FROM t)") + .run(); +} + +/// Test hash join where build-side keys come from a small range and allow for +/// array-based lookup instead of a hash table. +// Array-based lookup is not implemented for cudfHashJoin. +TEST_P(MultiThreadedHashJoinTest, DISABLED_arrayBasedLookup) { + auto oddIndices = makeIndices(500, [](auto i) { return 2 * i + 1; }); + + std::vector probeVectors = { + // Join key vector is flat. + makeRowVector({ + makeFlatVector(1'000, [](auto row) { return row; }), + makeFlatVector(1'000, [](auto row) { return row; }), + }), + // Join key vector is constant. There is a match in the build side. + makeRowVector({ + makeConstant(4, 2'000), + makeFlatVector(2'000, [](auto row) { return row; }), + }), + // Join key vector is constant. There is no match. + makeRowVector({ + makeConstant(5, 2'000), + makeFlatVector(2'000, [](auto row) { return row; }), + }), + // Join key vector is a dictionary. + makeRowVector({ + wrapInDictionary( + oddIndices, + 500, + makeFlatVector(1'000, [](auto row) { return row * 4; })), + makeFlatVector(1'000, [](auto row) { return row; }), + })}; + + // 100 key values in [0, 198] range. + std::vector buildVectors = { + makeRowVector( + {makeFlatVector(100, [](auto row) { return row / 2; })}), + makeRowVector( + {makeFlatVector(100, [](auto row) { return row * 2; })}), + makeRowVector( + {makeFlatVector(100, [](auto row) { return row; })})}; + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .numDrivers(numDrivers_) + .probeKeys({"c0"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"c0"}) + .buildVectors(std::move(buildVectors)) + .joinOutputLayout({"c1"}) + .outputProjections({"c1 + 1"}) + .referenceQuery("SELECT t.c1 + 1 FROM t, u WHERE t.c0 = u.c0") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + if (hasSpill) { + return; + } + auto joinStats = task->taskStats() + .pipelineStats.back() + .operatorStats.back() + .runtimeStats; + ASSERT_EQ(151, joinStats["distinctKey0"].sum); + ASSERT_EQ(200, joinStats["rangeKey0"].sum); + }) + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, joinSidesDifferentSchema) { + // In this join, the tables have different schema. LHS table t has schema + // {INTEGER, VARCHAR, INTEGER}. RHS table u has schema {INTEGER, REAL, + // INTEGER}. The filter predicate uses + // a column from the right table before the left and the corresponding + // columns at the same channel number(1) have different types. This has been + // a source of crashes in the join logic. + size_t batchSize = 100; + + std::vector stringVector = {"aaa", "bbb", "ccc", "ddd", "eee"}; + std::vector probeVectors = + makeBatches(5, [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector(batchSize, [](auto row) { return row; }), + makeFlatVector( + batchSize, + [&](auto row) { + return StringView(stringVector[row % stringVector.size()]); + }), + makeFlatVector(batchSize, [](auto row) { return row; }), + }); + }); + std::vector buildVectors = + makeBatches(5, [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector(batchSize, [](auto row) { return row; }), + makeFlatVector( + batchSize, [](auto row) { return row * 5.0; }), + makeFlatVector(batchSize, [](auto row) { return row; }), + }); + }); + + // In this hash join the 2 tables have a common key which is the + // first channel in both tables. + const std::string referenceQuery = + "SELECT t.c0 * t.c2/2 FROM " + " t, u " + " WHERE t.c0 = u.c0 AND " + // TODO: enable ltrim test after the race condition in expression + // execution gets fixed. + //" u.c2 > 10 AND ltrim(t.c1) = 'aaa'"; + " u.c2 > 10"; + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .numDrivers(numDrivers_) + .probeKeys({"t_c0"}) + .probeVectors(std::move(probeVectors)) + .probeProjections({"c0 AS t_c0", "c1 AS t_c1", "c2 AS t_c2"}) + .buildKeys({"u_c0"}) + .buildVectors(std::move(buildVectors)) + .buildProjections({"c0 AS u_c0", "c1 AS u_c1", "c2 AS u_c2"}) + //.joinFilter("u_c2 > 10 AND ltrim(t_c1) == 'aaa'") + .joinFilter("u_c2 > 10") + .joinOutputLayout({"t_c0", "t_c2"}) + .outputProjections({"t_c0 * t_c2/2"}) + .referenceQuery(referenceQuery) + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, innerJoinWithEmptyBuild) { + const std::vector finishOnEmptys = {false, true}; + for (auto finishOnEmpty : finishOnEmptys) { + SCOPED_TRACE(fmt::format("finishOnEmpty: {}", finishOnEmpty)); + + std::vector probeVectors = makeBatches(5, [&](int32_t batch) { + return makeRowVector({ + makeFlatVector( + 123, + [batch](auto row) { return row * 11 / std::max(batch, 1); }, + nullEvery(13)), + makeFlatVector(1'234, [](auto row) { return row; }), + }); + }); + std::vector buildVectors = + makeBatches(10, [&](int32_t batch) { + return makeRowVector({makeFlatVector( + 123, + [batch](auto row) { return row % std::max(batch, 1); }, + nullEvery(7))}); + }); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .hashProbeFinishEarlyOnEmptyBuild(finishOnEmpty) + .numDrivers(numDrivers_) + .probeKeys({"c0"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"c0"}) + .buildVectors(std::move(buildVectors)) + .buildFilter("c0 < 0") + .joinOutputLayout({"c1"}) + .referenceQuery("SELECT null LIMIT 0") + .checkSpillStats(false) + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + const auto statsPair = taskSpilledStats(*task); + ASSERT_EQ(statsPair.first.spilledRows, 0); + ASSERT_EQ(statsPair.first.spilledBytes, 0); + ASSERT_EQ(statsPair.first.spilledPartitions, 0); + ASSERT_EQ(statsPair.first.spilledFiles, 0); + ASSERT_EQ(statsPair.second.spilledRows, 0); + ASSERT_EQ(statsPair.second.spilledBytes, 0); + ASSERT_EQ(statsPair.second.spilledPartitions, 0); + ASSERT_EQ(statsPair.second.spilledFiles, 0); + verifyTaskSpilledRuntimeStats(*task, false); + ASSERT_EQ(maxHashBuildSpillLevel(*task), -1); + // Check the hash probe has processed probe input rows. + if (finishOnEmpty) { + ASSERT_EQ(getInputPositions(task, 1), 0); + } else { + ASSERT_GT(getInputPositions(task, 1), 0); + } + }) + .run(); + } +} + +TEST_P(MultiThreadedHashJoinTest, leftSemiJoinFilter) { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .numDrivers(numDrivers_) + .probeType(probeType_) + .probeVectors(174, 5) + .probeKeys({"t_k1"}) + .buildType(buildType_) + .buildVectors(133, 4) + .buildKeys({"u_k1"}) + .joinType(core::JoinType::kLeftSemiFilter) + .joinOutputLayout({"t_k2"}) + .referenceQuery("SELECT t_k2 FROM t WHERE t_k1 IN (SELECT u_k1 FROM u)") + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, leftSemiJoinFilterWithEmptyBuild) { + const std::vector finishOnEmptys = {false, true}; + for (const auto finishOnEmpty : finishOnEmptys) { + SCOPED_TRACE(fmt::format("finishOnEmpty: {}", finishOnEmpty)); + + std::vector probeVectors = + makeBatches(10, [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 1'234, [](auto row) { return row % 11; }, nullEvery(13)), + makeFlatVector(1'234, [](auto row) { return row; }), + }); + }); + std::vector buildVectors = + makeBatches(10, [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 123, [](auto row) { return row % 5; }, nullEvery(7)), + }); + }); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .hashProbeFinishEarlyOnEmptyBuild(finishOnEmpty) + .numDrivers(numDrivers_) + .probeKeys({"c0"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"c0"}) + .buildVectors(std::move(buildVectors)) + .joinType(core::JoinType::kLeftSemiFilter) + .joinFilter("c0 < 0") + .joinOutputLayout({"c1"}) + .referenceQuery( + "SELECT t.c1 FROM t WHERE t.c0 IN (SELECT c0 FROM u WHERE c0 < 0)") + .run(); + } +} + +TEST_P(MultiThreadedHashJoinTest, leftSemiJoinFilterWithExtraFilter) { + std::vector probeVectors = makeBatches(5, [&](int32_t batch) { + return makeRowVector( + {"t0", "t1"}, + { + makeFlatVector( + 250, [batch](auto row) { return row % (11 + batch); }), + makeFlatVector( + 250, [batch](auto row) { return row * batch; }), + }); + }); + + std::vector buildVectors = makeBatches(5, [&](int32_t batch) { + return makeRowVector( + {"u0", "u1"}, + { + makeFlatVector( + 123, [batch](auto row) { return row % (5 + batch); }), + makeFlatVector( + 123, [batch](auto row) { return row * batch; }), + }); + }); + + { + auto testProbeVectors = probeVectors; + auto testBuildVectors = buildVectors; + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .numDrivers(numDrivers_) + .probeKeys({"t0"}) + .probeVectors(std::move(testProbeVectors)) + .buildKeys({"u0"}) + .buildVectors(std::move(testBuildVectors)) + .joinType(core::JoinType::kLeftSemiFilter) + .joinOutputLayout({"t0", "t1"}) + .referenceQuery( + "SELECT t.* FROM t WHERE EXISTS (SELECT u0 FROM u WHERE t0 = u0)") + .run(); + } + + { + auto testProbeVectors = probeVectors; + auto testBuildVectors = buildVectors; + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .numDrivers(numDrivers_) + .probeKeys({"t0"}) + .probeVectors(std::move(testProbeVectors)) + .buildKeys({"u0"}) + .buildVectors(std::move(testBuildVectors)) + .joinType(core::JoinType::kLeftSemiFilter) + .joinFilter("t1 != u1") + .joinOutputLayout({"t0", "t1"}) + .referenceQuery( + "SELECT t.* FROM t WHERE EXISTS (SELECT u0, u1 FROM u WHERE t0 = u0 AND t1 <> u1)") + .run(); + } +} + +TEST_P(MultiThreadedHashJoinTest, rightSemiJoinFilter) { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .probeType(probeType_) + .probeVectors(133, 3) + .probeKeys({"t_k1"}) + .buildType(buildType_) + .buildVectors(174, 4) + .buildKeys({"u_k1"}) + .joinType(core::JoinType::kRightSemiFilter) + .joinOutputLayout({"u_k2"}) + .referenceQuery("SELECT u_k2 FROM u WHERE u_k1 IN (SELECT t_k1 FROM t)") + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, rightSemiJoinFilterWithEmptyBuild) { + const std::vector finishOnEmptys = {false, true}; + for (const auto finishOnEmpty : finishOnEmptys) { + SCOPED_TRACE(fmt::format("finishOnEmpty: {}", finishOnEmpty)); + + // probeVectors size is greater than buildVector size. + std::vector probeVectors = + makeBatches(5, [&](uint32_t /*unused*/) { + return makeRowVector( + {"t0", "t1"}, + {makeFlatVector( + 431, [](auto row) { return row % 11; }, nullEvery(13)), + makeFlatVector(431, [](auto row) { return row; })}); + }); + + std::vector buildVectors = + makeBatches(5, [&](uint32_t /*unused*/) { + return makeRowVector( + {"u0", "u1"}, + { + makeFlatVector( + 434, [](auto row) { return row % 5; }, nullEvery(7)), + makeFlatVector(434, [](auto row) { return row; }), + }); + }); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .hashProbeFinishEarlyOnEmptyBuild(finishOnEmpty) + .numDrivers(numDrivers_) + .probeKeys({"t0"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u0"}) + .buildVectors(std::move(buildVectors)) + .buildFilter("u0 < 0") + .joinType(core::JoinType::kRightSemiFilter) + .joinOutputLayout({"u1"}) + .referenceQuery( + "SELECT u.u1 FROM u WHERE u.u0 IN (SELECT t0 FROM t) AND u.u0 < 0") + .checkSpillStats(false) + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + const auto statsPair = taskSpilledStats(*task); + ASSERT_EQ(statsPair.first.spilledRows, 0); + ASSERT_EQ(statsPair.first.spilledBytes, 0); + ASSERT_EQ(statsPair.first.spilledPartitions, 0); + ASSERT_EQ(statsPair.first.spilledFiles, 0); + ASSERT_EQ(statsPair.second.spilledRows, 0); + ASSERT_EQ(statsPair.second.spilledBytes, 0); + ASSERT_EQ(statsPair.second.spilledPartitions, 0); + ASSERT_EQ(statsPair.second.spilledFiles, 0); + verifyTaskSpilledRuntimeStats(*task, false); + ASSERT_EQ(maxHashBuildSpillLevel(*task), -1); + // Check the hash probe has processed probe input rows. + if (finishOnEmpty) { + ASSERT_EQ(getInputPositions(task, 1), 0); + } else { + ASSERT_GT(getInputPositions(task, 1), 0); + } + }) + .run(); + } +} + +TEST_P(MultiThreadedHashJoinTest, rightSemiJoinFilterWithAllMatches) { + // Make build side larger to test all rows are returned. + std::vector probeVectors = + makeBatches(3, [&](uint32_t /*unused*/) { + return makeRowVector( + {"t0", "t1"}, + { + makeFlatVector( + 123, [](auto row) { return row % 5; }, nullEvery(7)), + makeFlatVector(123, [](auto row) { return row; }), + }); + }); + + std::vector buildVectors = + makeBatches(5, [&](uint32_t /*unused*/) { + return makeRowVector( + {"u0", "u1"}, + {makeFlatVector( + 314, [](auto row) { return row % 11; }, nullEvery(13)), + makeFlatVector(314, [](auto row) { return row; })}); + }); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .probeKeys({"t0"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u0"}) + .buildVectors(std::move(buildVectors)) + .joinType(core::JoinType::kRightSemiFilter) + .joinOutputLayout({"u1"}) + .referenceQuery("SELECT u.u1 FROM u WHERE u.u0 IN (SELECT t0 FROM t)") + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, rightSemiJoinFilterWithExtraFilter) { + auto probeVectors = makeBatches(4, [&](int32_t /*unused*/) { + return makeRowVector( + {"t0", "t1"}, + { + makeFlatVector(345, [](auto row) { return row; }), + makeFlatVector(345, [](auto row) { return row; }), + }); + }); + + auto buildVectors = makeBatches(4, [&](int32_t /*unused*/) { + return makeRowVector( + {"u0", "u1"}, + { + makeFlatVector(250, [](auto row) { return row; }), + makeFlatVector(250, [](auto row) { return row; }), + }); + }); + + // Always true filter. + { + auto testProbeVectors = probeVectors; + auto testBuildVectors = buildVectors; + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .probeKeys({"t0"}) + .probeVectors(std::move(testProbeVectors)) + .buildKeys({"u0"}) + .buildVectors(std::move(testBuildVectors)) + .joinType(core::JoinType::kRightSemiFilter) + .joinFilter("t1 > -1") + .joinOutputLayout({"u0", "u1"}) + .referenceQuery( + "SELECT u.* FROM u WHERE EXISTS (SELECT t0 FROM t WHERE u0 = t0 AND t1 > -1)") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + ASSERT_EQ( + getOutputPositions(task, "HashProbe"), 200 * 5 * numDrivers_); + }) + .run(); + } + + // Always false filter. + { + auto testProbeVectors = probeVectors; + auto testBuildVectors = buildVectors; + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .probeKeys({"t0"}) + .probeVectors(std::move(testProbeVectors)) + .buildKeys({"u0"}) + .buildVectors(std::move(testBuildVectors)) + .joinType(core::JoinType::kRightSemiFilter) + .joinFilter("t1 > 100000") + .joinOutputLayout({"u0", "u1"}) + .referenceQuery( + "SELECT u.* FROM u WHERE EXISTS (SELECT t0 FROM t WHERE u0 = t0 AND t1 > 100000)") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + ASSERT_EQ(getOutputPositions(task, "HashProbe"), 0); + }) + .run(); + } + + // Selective filter. + { + auto testProbeVectors = probeVectors; + auto testBuildVectors = buildVectors; + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .probeKeys({"t0"}) + .probeVectors(std::move(testProbeVectors)) + .buildKeys({"u0"}) + .buildVectors(std::move(testBuildVectors)) + .joinType(core::JoinType::kRightSemiFilter) + .joinFilter("t1 % 5 = 0") + .joinOutputLayout({"u0", "u1"}) + .referenceQuery( + "SELECT u.* FROM u WHERE EXISTS (SELECT t0 FROM t WHERE u0 = t0 AND t1 % 5 = 0)") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + ASSERT_EQ( + getOutputPositions(task, "HashProbe"), 200 / 5 * 5 * numDrivers_); + }) + .run(); + } +} + +TEST_P(MultiThreadedHashJoinTest, semiFilterOverLazyVectors) { + auto probeVectors = makeBatches(1, [&](auto /*unused*/) { + return makeRowVector( + {"t0", "t1"}, + { + makeFlatVector(1'000, [](auto row) { return row; }), + makeFlatVector(1'000, [](auto row) { return row * 10; }), + }); + }); + + auto buildVectors = makeBatches(3, [&](auto /*unused*/) { + return makeRowVector( + {"u0", "u1"}, + { + makeFlatVector( + 1'000, [](auto row) { return -100 + (row / 5); }), + makeFlatVector( + 1'000, [](auto row) { return -1000 + (row / 5) * 10; }), + }); + }); + + std::shared_ptr probeFile = TempFilePath::create(); + writeToFile(probeFile->getPath(), probeVectors); + + std::shared_ptr buildFile = TempFilePath::create(); + writeToFile(buildFile->getPath(), buildVectors); + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + core::PlanNodeId probeScanId; + core::PlanNodeId buildScanId; + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .tableScan(asRowType(probeVectors[0]->type())) + .capturePlanNodeId(probeScanId) + .hashJoin( + {"t0"}, + {"u0"}, + PlanBuilder(planNodeIdGenerator) + .tableScan(asRowType(buildVectors[0]->type())) + .capturePlanNodeId(buildScanId) + .planNode(), + "", + {"t0", "t1"}, + core::JoinType::kLeftSemiFilter) + .planNode(); + + SplitInput splitInput = { + {probeScanId, + {exec::Split(makeHiveConnectorSplit(probeFile->getPath()))}}, + {buildScanId, + {exec::Split(makeHiveConnectorSplit(buildFile->getPath()))}}, + }; + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(plan) + .inputSplits(splitInput) + .checkSpillStats(false) + .referenceQuery("SELECT t0, t1 FROM t WHERE t0 IN (SELECT u0 FROM u)") + .run(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(flipJoinSides(plan)) + .inputSplits(splitInput) + .checkSpillStats(false) + .referenceQuery("SELECT t0, t1 FROM t WHERE t0 IN (SELECT u0 FROM u)") + .run(); + + // With extra filter. + planNodeIdGenerator = std::make_shared(); + plan = PlanBuilder(planNodeIdGenerator) + .tableScan(asRowType(probeVectors[0]->type())) + .capturePlanNodeId(probeScanId) + .hashJoin( + {"t0"}, + {"u0"}, + PlanBuilder(planNodeIdGenerator) + .tableScan(asRowType(buildVectors[0]->type())) + .capturePlanNodeId(buildScanId) + .planNode(), + "(t1 + u1) % 3 = 0", + {"t0", "t1"}, + core::JoinType::kLeftSemiFilter) + .planNode(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(plan) + .inputSplits(splitInput) + .checkSpillStats(false) + .referenceQuery( + "SELECT t0, t1 FROM t WHERE t0 IN (SELECT u0 FROM u WHERE (t1 + u1) % 3 = 0)") + .run(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(flipJoinSides(plan)) + .inputSplits(splitInput) + .checkSpillStats(false) + .referenceQuery( + "SELECT t0, t1 FROM t WHERE t0 IN (SELECT u0 FROM u WHERE (t1 + u1) % 3 = 0)") + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, nullAwareAntiJoin) { + std::vector probeVectors = + makeBatches(5, [&](uint32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 1'000, [](auto row) { return row % 11; }, nullEvery(13)), + makeFlatVector(1'000, [](auto row) { return row; }), + }); + }); + + std::vector buildVectors = + makeBatches(5, [&](uint32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 1'234, [](auto row) { return row % 5; }, nullEvery(7)), + }); + }); + + { + auto testProbeVectors = probeVectors; + auto testBuildVectors = buildVectors; + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .probeKeys({"c0"}) + .probeVectors(std::move(testProbeVectors)) + .buildKeys({"c0"}) + .buildVectors(std::move(testBuildVectors)) + .buildFilter("c0 IS NOT NULL") + .joinType(core::JoinType::kAnti) + .nullAware(true) + .joinOutputLayout({"c1"}) + .referenceQuery( + "SELECT t.c1 FROM t WHERE t.c0 NOT IN (SELECT c0 FROM u WHERE c0 IS NOT NULL)") + .checkSpillStats(false) + .run(); + } + + // Empty build side. + { + auto testProbeVectors = probeVectors; + auto testBuildVectors = buildVectors; + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .probeKeys({"c0"}) + .probeVectors(std::move(testProbeVectors)) + .buildKeys({"c0"}) + .buildVectors(std::move(testBuildVectors)) + .buildFilter("c0 < 0") + .joinType(core::JoinType::kAnti) + .nullAware(true) + .joinOutputLayout({"c1"}) + .referenceQuery( + "SELECT t.c1 FROM t WHERE t.c0 NOT IN (SELECT c0 FROM u WHERE c0 < 0)") + .checkSpillStats(false) + .run(); + } + + // Build side with nulls. Null-aware Anti join always returns nothing. + { + auto testProbeVectors = probeVectors; + auto testBuildVectors = buildVectors; + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .probeKeys({"c0"}) + .probeVectors(std::move(testProbeVectors)) + .buildKeys({"c0"}) + .buildVectors(std::move(testBuildVectors)) + .joinType(core::JoinType::kAnti) + .nullAware(true) + .joinOutputLayout({"c1"}) + .referenceQuery( + "SELECT t.c1 FROM t WHERE t.c0 NOT IN (SELECT c0 FROM u)") + .checkSpillStats(false) + .run(); + } +} + +TEST_P(MultiThreadedHashJoinTest, nullAwareAntiJoinWithFilter) { + std::vector probeVectors = + makeBatches(5, [&](int32_t /*unused*/) { + return makeRowVector( + {"t0", "t1"}, + { + makeFlatVector(128, [](auto row) { return row % 11; }), + makeFlatVector(128, [](auto row) { return row; }), + }); + }); + + std::vector buildVectors = + makeBatches(5, [&](int32_t /*unused*/) { + return makeRowVector( + {"u0", "u1"}, + { + makeFlatVector(123, [](auto row) { return row % 5; }), + makeFlatVector(123, [](auto row) { return row; }), + }); + }); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .probeKeys({"t0"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u0"}) + .buildVectors(std::move(buildVectors)) + .joinType(core::JoinType::kAnti) + .nullAware(true) + .joinFilter("t1 != u1") + .joinOutputLayout({"t0", "t1"}) + .referenceQuery( + "SELECT t.* FROM t WHERE NOT EXISTS (SELECT * FROM u WHERE t0 = u0 AND t1 <> u1)") + .checkSpillStats(false) + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + // Verify spilling is not triggered in case of null-aware anti-join + // with filter. + const auto statsPair = taskSpilledStats(*task); + ASSERT_EQ(statsPair.first.spilledRows, 0); + ASSERT_EQ(statsPair.first.spilledBytes, 0); + ASSERT_EQ(statsPair.first.spilledPartitions, 0); + ASSERT_EQ(statsPair.first.spilledFiles, 0); + ASSERT_EQ(statsPair.second.spilledRows, 0); + ASSERT_EQ(statsPair.second.spilledBytes, 0); + ASSERT_EQ(statsPair.second.spilledPartitions, 0); + ASSERT_EQ(statsPair.second.spilledFiles, 0); + verifyTaskSpilledRuntimeStats(*task, false); + ASSERT_EQ(maxHashBuildSpillLevel(*task), -1); + }) + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, nullAwareAntiJoinWithFilterAndEmptyBuild) { + const std::vector finishOnEmptys = {false, true}; + for (const auto finishOnEmpty : finishOnEmptys) { + SCOPED_TRACE(fmt::format("finishOnEmpty: {}", finishOnEmpty)); + + auto probeVectors = makeBatches(4, [&](int32_t /*unused*/) { + return makeRowVector( + {"t0", "t1"}, + { + makeNullableFlatVector({std::nullopt, 1, 2}), + makeFlatVector({0, 1, 2}), + }); + }); + auto buildVectors = makeBatches(4, [&](int32_t /*unused*/) { + return makeRowVector( + {"u0", "u1"}, + { + makeNullableFlatVector({3, 2, 3}), + makeFlatVector({0, 2, 3}), + }); + }); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .hashProbeFinishEarlyOnEmptyBuild(finishOnEmpty) + .numDrivers(numDrivers_) + .probeKeys({"t0"}) + .probeVectors(std::vector(probeVectors)) + .buildKeys({"u0"}) + .buildVectors(std::vector(buildVectors)) + .buildFilter("u0 < 0") + .joinType(core::JoinType::kAnti) + .nullAware(true) + .joinFilter("u1 > t1") + .joinOutputLayout({"t0", "t1"}) + .referenceQuery( + "SELECT t.* FROM t WHERE NOT EXISTS (SELECT * FROM u WHERE u0 < 0 AND u.u0 = t.t0)") + .checkSpillStats(false) + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + // Verify spilling is not triggered in case of null-aware anti-join + // with filter. + const auto statsPair = taskSpilledStats(*task); + ASSERT_EQ(statsPair.first.spilledRows, 0); + ASSERT_EQ(statsPair.first.spilledBytes, 0); + ASSERT_EQ(statsPair.first.spilledPartitions, 0); + ASSERT_EQ(statsPair.first.spilledFiles, 0); + ASSERT_EQ(statsPair.second.spilledRows, 0); + ASSERT_EQ(statsPair.second.spilledBytes, 0); + ASSERT_EQ(statsPair.second.spilledPartitions, 0); + ASSERT_EQ(statsPair.second.spilledFiles, 0); + verifyTaskSpilledRuntimeStats(*task, false); + ASSERT_EQ(maxHashBuildSpillLevel(*task), -1); + }) + .run(); + } +} + +TEST_P(MultiThreadedHashJoinTest, nullAwareAntiJoinWithFilterAndNullKey) { + auto probeVectors = makeBatches(4, [&](int32_t /*unused*/) { + return makeRowVector( + {"t0", "t1"}, + { + makeNullableFlatVector({std::nullopt, 1, 2}), + makeFlatVector({0, 1, 2}), + }); + }); + auto buildVectors = makeBatches(4, [&](int32_t /*unused*/) { + return makeRowVector( + {"u0", "u1"}, + { + makeNullableFlatVector({std::nullopt, 2, 3}), + makeFlatVector({0, 2, 3}), + }); + }); + + std::vector filters({"u1 > t1", "u1 * t1 > 0"}); + for (const std::string& filter : filters) { + const auto referenceSql = fmt::format( + "SELECT t.* FROM t WHERE t0 NOT IN (SELECT u0 FROM u WHERE {})", + filter); + + auto testProbeVectors = probeVectors; + auto testBuildVectors = buildVectors; + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .probeKeys({"t0"}) + .probeVectors(std::move(testProbeVectors)) + .buildKeys({"u0"}) + .buildVectors(std::move(testBuildVectors)) + .joinType(core::JoinType::kAnti) + .nullAware(true) + .joinFilter(filter) + .joinOutputLayout({"t0", "t1"}) + .referenceQuery(referenceSql) + .checkSpillStats(false) + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + // Verify spilling is not triggered in case of null-aware anti-join + // with filter. + const auto statsPair = taskSpilledStats(*task); + ASSERT_EQ(statsPair.first.spilledRows, 0); + ASSERT_EQ(statsPair.first.spilledBytes, 0); + ASSERT_EQ(statsPair.first.spilledPartitions, 0); + ASSERT_EQ(statsPair.first.spilledFiles, 0); + ASSERT_EQ(statsPair.second.spilledRows, 0); + ASSERT_EQ(statsPair.second.spilledBytes, 0); + ASSERT_EQ(statsPair.second.spilledPartitions, 0); + ASSERT_EQ(statsPair.second.spilledFiles, 0); + verifyTaskSpilledRuntimeStats(*task, false); + ASSERT_EQ(maxHashBuildSpillLevel(*task), -1); + }) + .run(); + } +} + +TEST_P( + MultiThreadedHashJoinTest, + hashModeNullAwareAntiJoinWithFilterAndNullKey) { + // Use float type keys to trigger hash mode table. + auto probeVectors = makeBatches(50, [&](int32_t /*unused*/) { + return makeRowVector( + {"t0", "t1"}, + { + makeNullableFlatVector({std::nullopt, 1, 2}), + makeFlatVector({1, 1, 2}), + }); + }); + auto buildVectors = makeBatches(5, [&](int32_t /*unused*/) { + return makeRowVector( + {"u0", "u1"}, + { + makeNullableFlatVector({std::nullopt, 2, 3}), + makeFlatVector({0, 2, 3}), + }); + }); + + std::vector filters({"u1 < t1", "u1 + t1 = 0"}); + for (const std::string& filter : filters) { + const auto referenceSql = fmt::format( + "SELECT t.* FROM t WHERE t0 NOT IN (SELECT u0 FROM u WHERE {})", + filter); + + auto testProbeVectors = probeVectors; + auto testBuildVectors = buildVectors; + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .probeKeys({"t0"}) + .probeVectors(std::move(testProbeVectors)) + .buildKeys({"u0"}) + .buildVectors(std::move(testBuildVectors)) + .joinType(core::JoinType::kAnti) + .nullAware(true) + .joinFilter(filter) + .joinOutputLayout({"t0", "t1"}) + .referenceQuery(referenceSql) + .checkSpillStats(false) + .run(); + } +} + +TEST_P(MultiThreadedHashJoinTest, nullAwareAntiJoinWithFilterOnNullableColumn) { + const std::string referenceSql = + "SELECT t.* FROM t WHERE t0 NOT IN (SELECT u0 FROM u WHERE t1 <> u1)"; + const std::string joinFilter = "t1 <> u1"; + { + SCOPED_TRACE("null filter column"); + auto probeVectors = makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector( + {"t0", "t1"}, + { + makeFlatVector(200, [](auto row) { return row % 11; }), + makeFlatVector(200, folly::identity, nullEvery(97)), + }); + }); + auto buildVectors = makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector( + {"u0", "u1"}, + { + makeFlatVector(234, [](auto row) { return row % 5; }), + makeFlatVector(234, folly::identity, nullEvery(91)), + }); + }); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .probeKeys({"t0"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u0"}) + .buildVectors(std::move(buildVectors)) + .joinType(core::JoinType::kAnti) + .nullAware(true) + .joinFilter(joinFilter) + .joinOutputLayout({"t0", "t1"}) + .referenceQuery(referenceSql) + .checkSpillStats(false) + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + // Verify spilling is not triggered in case of null-aware anti-join + // with filter. + const auto statsPair = taskSpilledStats(*task); + ASSERT_EQ(statsPair.first.spilledRows, 0); + ASSERT_EQ(statsPair.first.spilledBytes, 0); + ASSERT_EQ(statsPair.first.spilledPartitions, 0); + ASSERT_EQ(statsPair.first.spilledFiles, 0); + ASSERT_EQ(statsPair.second.spilledRows, 0); + ASSERT_EQ(statsPair.second.spilledBytes, 0); + ASSERT_EQ(statsPair.second.spilledPartitions, 0); + ASSERT_EQ(statsPair.second.spilledFiles, 0); + verifyTaskSpilledRuntimeStats(*task, false); + ASSERT_EQ(maxHashBuildSpillLevel(*task), -1); + }) + .run(); + } + + { + SCOPED_TRACE("null filter and key column"); + auto probeVectors = makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector( + {"t0", "t1"}, + { + makeFlatVector( + 200, [](auto row) { return row % 11; }, nullEvery(23)), + makeFlatVector(200, folly::identity, nullEvery(29)), + }); + }); + auto buildVectors = makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector( + {"u0", "u1"}, + { + makeFlatVector( + 234, [](auto row) { return row % 5; }, nullEvery(31)), + makeFlatVector(234, folly::identity, nullEvery(37)), + }); + }); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .probeKeys({"t0"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u0"}) + .buildVectors(std::move(buildVectors)) + .joinType(core::JoinType::kAnti) + .nullAware(true) + .joinFilter(joinFilter) + .joinOutputLayout({"t0", "t1"}) + .referenceQuery(referenceSql) + .checkSpillStats(false) + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + // Verify spilling is not triggered in case of null-aware anti-join + // with filter. + const auto statsPair = taskSpilledStats(*task); + ASSERT_EQ(statsPair.first.spilledRows, 0); + ASSERT_EQ(statsPair.first.spilledBytes, 0); + ASSERT_EQ(statsPair.first.spilledPartitions, 0); + ASSERT_EQ(statsPair.first.spilledFiles, 0); + ASSERT_EQ(statsPair.second.spilledRows, 0); + ASSERT_EQ(statsPair.second.spilledBytes, 0); + ASSERT_EQ(statsPair.second.spilledPartitions, 0); + ASSERT_EQ(statsPair.second.spilledFiles, 0); + verifyTaskSpilledRuntimeStats(*task, false); + ASSERT_EQ(maxHashBuildSpillLevel(*task), -1); + }) + .run(); + } +} + +TEST_P(MultiThreadedHashJoinTest, antiJoin) { + auto probeVectors = makeBatches(64, [&](int32_t /*unused*/) { + return makeRowVector( + {"t0", "t1"}, + { + makeNullableFlatVector({std::nullopt, 1, 2}), + makeFlatVector({0, 1, 2}), + }); + }); + auto buildVectors = makeBatches(64, [&](int32_t /*unused*/) { + return makeRowVector( + {"u0", "u1"}, + { + makeNullableFlatVector({std::nullopt, 2, 3}), + makeFlatVector({0, 2, 3}), + }); + }); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .numDrivers(numDrivers_) + .probeKeys({"t0"}) + .probeVectors(std::vector(probeVectors)) + .buildKeys({"u0"}) + .buildVectors(std::vector(buildVectors)) + .joinType(core::JoinType::kAnti) + .joinOutputLayout({"t0", "t1"}) + .referenceQuery( + "SELECT t.* FROM t WHERE NOT EXISTS (SELECT * FROM u WHERE u.u0 = t.t0)") + .run(); + + std::vector filters({ + "u1 > t1", "u1 * t1 > 0", + // This filter is true on rows without a match. It should not prevent + // the row from being returned. + // Disabling this because coalesce is not supported in cudf. + // "coalesce(u1, t1, 0::integer) is not null", + // This filter throws if evaluated on rows without a match. The join + // should not evaluate filter on those rows and therefore should not + // fail. + // Disabling this because coalesce is not supported in cudf. + // "t1 / coalesce(u1, 0::integer) is not null", + // This filter triggers memory pool allocation at + // HashBuild::setupFilterForAntiJoins, which should not be invoked in + // operator's constructor. + // Disabling this because contains is not supported in cudf. + // "contains(array[1, 2, NULL], 1)", + }); + for (const std::string& filter : filters) { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .numDrivers(numDrivers_) + .probeKeys({"t0"}) + .probeVectors(std::vector(probeVectors)) + .buildKeys({"u0"}) + .buildVectors(std::vector(buildVectors)) + .joinType(core::JoinType::kAnti) + .joinFilter(filter) + .joinOutputLayout({"t0", "t1"}) + .referenceQuery(fmt::format( + "SELECT t.* FROM t WHERE NOT EXISTS (SELECT * FROM u WHERE u.u0 = t.t0 AND {})", + filter)) + .run(); + } +} + +TEST_P(MultiThreadedHashJoinTest, antiJoinWithFilterAndEmptyBuild) { + const std::vector finishOnEmptys = {false, true}; + for (const auto finishOnEmpty : finishOnEmptys) { + SCOPED_TRACE(fmt::format("finishOnEmpty: {}", finishOnEmpty)); + + auto probeVectors = makeBatches(4, [&](int32_t /*unused*/) { + return makeRowVector( + {"t0", "t1"}, + { + makeNullableFlatVector({std::nullopt, 1, 2}), + makeFlatVector({0, 1, 2}), + }); + }); + auto buildVectors = makeBatches(4, [&](int32_t /*unused*/) { + return makeRowVector( + {"u0", "u1"}, + { + makeNullableFlatVector({3, 2, 3}), + makeFlatVector({0, 2, 3}), + }); + }); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .hashProbeFinishEarlyOnEmptyBuild(finishOnEmpty) + .numDrivers(numDrivers_) + .probeKeys({"t0"}) + .probeVectors(std::vector(probeVectors)) + .buildKeys({"u0"}) + .buildVectors(std::vector(buildVectors)) + .buildFilter("u0 < 0") + .joinType(core::JoinType::kAnti) + .joinFilter("u1 > t1") + .joinOutputLayout({"t0", "t1"}) + .referenceQuery( + "SELECT t.* FROM t WHERE NOT EXISTS (SELECT * FROM u WHERE u0 < 0 AND u.u0 = t.t0)") + .checkSpillStats(false) + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + const auto statsPair = taskSpilledStats(*task); + ASSERT_EQ(statsPair.first.spilledRows, 0); + ASSERT_EQ(statsPair.first.spilledBytes, 0); + ASSERT_EQ(statsPair.first.spilledPartitions, 0); + ASSERT_EQ(statsPair.first.spilledFiles, 0); + ASSERT_EQ(statsPair.second.spilledRows, 0); + ASSERT_EQ(statsPair.second.spilledBytes, 0); + ASSERT_EQ(statsPair.second.spilledPartitions, 0); + ASSERT_EQ(statsPair.second.spilledFiles, 0); + verifyTaskSpilledRuntimeStats(*task, false); + ASSERT_EQ(maxHashBuildSpillLevel(*task), -1); + }) + .run(); + } +} + +TEST_P(MultiThreadedHashJoinTest, leftJoin) { + // Left side keys are [0, 1, 2,..20]. + // Use 3-rd column as row number to allow for asserting the order of + // results. + std::vector probeVectors = mergeBatches( + makeBatches( + 3, + [&](int32_t /*unused*/) { + return makeRowVector( + {"c0", "c1", "row_number"}, + { + makeFlatVector( + 77, [](auto row) { return row % 21; }, nullEvery(13)), + makeFlatVector(77, [](auto row) { return row; }), + makeFlatVector(77, [](auto row) { return row; }), + }); + }), + makeBatches( + 2, + [&](int32_t /*unused*/) { + return makeRowVector( + {"c0", "c1", "row_number"}, + { + makeFlatVector( + 97, + [](auto row) { return (row + 3) % 21; }, + nullEvery(13)), + makeFlatVector(97, [](auto row) { return row; }), + makeFlatVector( + 97, [](auto row) { return 97 + row; }), + }); + }), + true); + + std::vector buildVectors = + makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 73, [](auto row) { return row % 5; }, nullEvery(7)), + makeFlatVector( + 73, [](auto row) { return -111 + row * 2; }, nullEvery(7)), + }); + }); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .numDrivers(numDrivers_) + .probeKeys({"c0"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u_c0"}) + .buildVectors(std::move(buildVectors)) + .buildProjections({"c0 AS u_c0", "c1 AS u_c1"}) + .joinType(core::JoinType::kLeft) + .joinOutputLayout({"row_number", "c0", "c1", "u_c0"}) + .referenceQuery( + "SELECT t.row_number, t.c0, t.c1, u.c0 FROM t LEFT JOIN u ON t.c0 = u.c0") + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + int nullJoinBuildKeyCount = 0; + int nullJoinProbeKeyCount = 0; + + for (auto& pipeline : task->taskStats().pipelineStats) { + for (auto op : pipeline.operatorStats) { + if (op.operatorType == "CudfHashJoinBuild") { + nullJoinBuildKeyCount += op.numNullKeys; + } + if (op.operatorType == "CudfHashJoinProbe") { + nullJoinProbeKeyCount += op.numNullKeys; + } + } + } + ASSERT_EQ(nullJoinBuildKeyCount, 33 * GetParam().numDrivers); + ASSERT_EQ(nullJoinProbeKeyCount, 34 * GetParam().numDrivers); + }) + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, nullStatsWithEmptyBuild) { + std::vector probeVectors = + makeBatches(1, [&](int32_t /*unused*/) { + return makeRowVector( + {"c0", "c1", "row_number"}, + { + makeFlatVector( + 77, [](auto row) { return row % 21; }, nullEvery(13)), + makeFlatVector(77, [](auto row) { return row; }), + makeFlatVector(77, [](auto row) { return row; }), + }); + }); + + // All null keys on build side. + std::vector buildVectors = + makeBatches(1, [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 1, [](auto row) { return row % 5; }, nullEvery(1)), + makeFlatVector( + 1, [](auto row) { return -111 + row * 2; }, nullEvery(1)), + }); + }); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .probeKeys({"c0"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u_c0"}) + .buildVectors(std::move(buildVectors)) + .buildProjections({"c0 AS u_c0", "c1 AS u_c1"}) + .joinType(core::JoinType::kLeft) + .joinOutputLayout({"row_number", "c0", "c1", "u_c0"}) + .referenceQuery( + "SELECT t.row_number, t.c0, t.c1, u.c0 FROM t LEFT JOIN u ON t.c0 = u.c0") + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + int nullJoinBuildKeyCount = 0; + int nullJoinProbeKeyCount = 0; + + for (auto& pipeline : task->taskStats().pipelineStats) { + for (auto op : pipeline.operatorStats) { + if (op.operatorType == "CudfHashJoinBuild") { + nullJoinBuildKeyCount += op.numNullKeys; + } + if (op.operatorType == "CudfHashJoinProbe") { + nullJoinProbeKeyCount += op.numNullKeys; + } + } + } + // Due to inaccurate stats tracking in case of empty build side, + // we will report 0 null keys on probe side. + // CudfHashJoinProbe will reports correct null keys + // since early exit is not implemented. + ASSERT_EQ(nullJoinProbeKeyCount, 6 * GetParam().numDrivers); + ASSERT_EQ(nullJoinBuildKeyCount, 1 * GetParam().numDrivers); + }) + .checkSpillStats(false) + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, leftJoinWithEmptyBuild) { + const std::vector finishOnEmptys = {false, true}; + for (const auto finishOnEmpty : finishOnEmptys) { + SCOPED_TRACE(fmt::format("finishOnEmpty: {}", finishOnEmpty)); + + // Left side keys are [0, 1, 2,..10]. + // Use 3-rd column as row number to allow for asserting the order of + // results. + std::vector probeVectors = mergeBatches( + makeBatches( + 3, + [&](int32_t /*unused*/) { + return makeRowVector( + {"c0", "c1", "row_number"}, + { + makeFlatVector( + 77, [](auto row) { return row % 11; }, nullEvery(13)), + makeFlatVector(77, [](auto row) { return row; }), + makeFlatVector(77, [](auto row) { return row; }), + }); + }), + makeBatches( + 2, + [&](int32_t /*unused*/) { + return makeRowVector( + {"c0", "c1", "row_number"}, + { + makeFlatVector( + 97, + [](auto row) { return (row + 3) % 11; }, + nullEvery(13)), + makeFlatVector(97, [](auto row) { return row; }), + makeFlatVector( + 97, [](auto row) { return 97 + row; }), + }); + }), + true); + + std::vector buildVectors = + makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 73, [](auto row) { return row % 5; }, nullEvery(7)), + makeFlatVector( + 73, [](auto row) { return -111 + row * 2; }, nullEvery(7)), + }); + }); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .hashProbeFinishEarlyOnEmptyBuild(finishOnEmpty) + .numDrivers(numDrivers_) + .probeKeys({"c0"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u_c0"}) + .buildVectors(std::move(buildVectors)) + .buildProjections({"c0 AS u_c0", "c1 AS u_c1"}) + .buildFilter("c0 < 0") + .joinType(core::JoinType::kLeft) + .joinOutputLayout({"row_number", "c1"}) + .referenceQuery( + "SELECT t.row_number, t.c1 FROM t LEFT JOIN (SELECT c0 FROM u WHERE c0 < 0) u ON t.c0 = u.c0") + .checkSpillStats(false) + .run(); + } +} + +TEST_P(MultiThreadedHashJoinTest, leftJoinWithNoJoin) { + // Left side keys are [0, 1, 2,..10]. + // Use 3-rd column as row number to allow for asserting the order of + // results. + std::vector probeVectors = mergeBatches( + makeBatches( + 3, + [&](int32_t /*unused*/) { + return makeRowVector( + {"c0", "c1", "row_number"}, + { + makeFlatVector( + 77, [](auto row) { return row % 11; }, nullEvery(13)), + makeFlatVector(77, [](auto row) { return row; }), + makeFlatVector(77, [](auto row) { return row; }), + }); + }), + makeBatches( + 2, + [&](int32_t /*unused*/) { + return makeRowVector( + {"c0", "c1", "row_number"}, + { + makeFlatVector( + 97, + [](auto row) { return (row + 3) % 11; }, + nullEvery(13)), + makeFlatVector(97, [](auto row) { return row; }), + makeFlatVector( + 97, [](auto row) { return 97 + row; }), + }); + }), + true); + + std::vector buildVectors = + makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 73, [](auto row) { return row % 5; }, nullEvery(7)), + makeFlatVector( + 73, [](auto row) { return -111 + row * 2; }, nullEvery(7)), + }); + }); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .numDrivers(numDrivers_) + .probeKeys({"c0"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u_c0"}) + .buildVectors(std::move(buildVectors)) + .buildProjections({"c0 - 123::INTEGER AS u_c0", "c1 AS u_c1"}) + .joinType(core::JoinType::kLeft) + .joinOutputLayout({"row_number", "c0", "u_c1"}) + .referenceQuery( + "SELECT t.row_number, t.c0, u.c1 FROM t LEFT JOIN (SELECT c0 - 123::INTEGER AS u_c0, c1 FROM u) u ON t.c0 = u.u_c0") + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, leftJoinWithAllMatch) { + // Left side keys are [0, 1, 2,..10]. + // Use 3-rd column as row number to allow for asserting the order of + // results. + std::vector probeVectors = mergeBatches( + makeBatches( + 3, + [&](int32_t /*unused*/) { + return makeRowVector( + {"c0", "c1", "row_number"}, + { + makeFlatVector( + 77, [](auto row) { return row % 11; }, nullEvery(13)), + makeFlatVector(77, [](auto row) { return row; }), + makeFlatVector(77, [](auto row) { return row; }), + }); + }), + makeBatches( + 2, + [&](int32_t /*unused*/) { + return makeRowVector( + {"c0", "c1", "row_number"}, + { + makeFlatVector( + 97, + [](auto row) { return (row + 3) % 11; }, + nullEvery(13)), + makeFlatVector(97, [](auto row) { return row; }), + makeFlatVector( + 97, [](auto row) { return 97 + row; }), + }); + }), + true); + + std::vector buildVectors = + makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 73, [](auto row) { return row % 5; }, nullEvery(7)), + makeFlatVector( + 73, [](auto row) { return -111 + row * 2; }, nullEvery(7)), + }); + }); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .numDrivers(numDrivers_) + .probeKeys({"c0"}) + .probeVectors(std::move(probeVectors)) + .probeFilter("c0 < 5") + .buildKeys({"u_c0"}) + .buildVectors(std::move(buildVectors)) + .buildProjections({"c0 AS u_c0", "c1 AS u_c1"}) + .joinType(core::JoinType::kLeft) + .joinOutputLayout({"row_number", "c0", "c1", "u_c1"}) + .referenceQuery( + "SELECT t.row_number, t.c0, t.c1, u.c1 FROM (SELECT * FROM t WHERE c0 < 5) t LEFT JOIN u ON t.c0 = u.c0") + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, leftJoinWithFilter) { + // Left side keys are [0, 1, 2,..10]. + // Use 3-rd column as row number to allow for asserting the order of + // results. + std::vector probeVectors = mergeBatches( + makeBatches( + 3, + [&](int32_t /*unused*/) { + return makeRowVector( + {"c0", "c1", "row_number"}, + { + makeFlatVector( + 77, [](auto row) { return row % 11; }, nullEvery(13)), + makeFlatVector(77, [](auto row) { return row; }), + makeFlatVector(77, [](auto row) { return row; }), + }); + }), + makeBatches( + 2, + [&](int32_t /*unused*/) { + return makeRowVector( + {"c0", "c1", "row_number"}, + { + makeFlatVector( + 97, + [](auto row) { return (row + 3) % 11; }, + nullEvery(13)), + makeFlatVector(97, [](auto row) { return row; }), + makeFlatVector( + 97, [](auto row) { return 97 + row; }), + }); + }), + true); + + std::vector buildVectors = + makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 73, [](auto row) { return row % 5; }, nullEvery(7)), + makeFlatVector( + 73, [](auto row) { return -111 + row * 2; }, nullEvery(7)), + }); + }); + + // Additional filter. + { + auto testProbeVectors = probeVectors; + auto testBuildVectors = buildVectors; + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .numDrivers(numDrivers_) + .probeKeys({"c0"}) + .probeVectors(std::move(testProbeVectors)) + .buildKeys({"u_c0"}) + .buildVectors(std::move(testBuildVectors)) + .buildProjections({"c0 AS u_c0", "c1 AS u_c1"}) + .joinType(core::JoinType::kLeft) + .joinFilter("(c1 + u_c1) % 2 = 1") + .joinOutputLayout({"row_number", "c0", "c1", "u_c1"}) + .referenceQuery( + "SELECT t.row_number, t.c0, t.c1, u.c1 FROM t LEFT JOIN u ON t.c0 = u.c0 AND (t.c1 + u.c1) % 2 = 1") + .run(); + } + + // No rows pass the additional filter. + { + auto testProbeVectors = probeVectors; + auto testBuildVectors = buildVectors; + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .numDrivers(numDrivers_) + .probeKeys({"c0"}) + .probeVectors(std::move(testProbeVectors)) + .buildKeys({"u_c0"}) + .buildVectors(std::move(testBuildVectors)) + .buildProjections({"c0 AS u_c0", "c1 AS u_c1"}) + .joinType(core::JoinType::kLeft) + .joinFilter("(c1 + u_c1) % 2 = 3") + .joinOutputLayout({"row_number", "c0", "c1", "u_c1"}) + .referenceQuery( + "SELECT t.row_number, t.c0, t.c1, u.c1 FROM t LEFT JOIN u ON t.c0 = u.c0 AND (t.c1 + u.c1) % 2 = 3") + .run(); + } +} + +/// Tests left join with a filter that may evaluate to true, false or null. +/// Makes sure that null filter results are handled correctly, e.g. as if the +/// filter returned false. +TEST_P(MultiThreadedHashJoinTest, leftJoinWithNullableFilter) { + std::vector probeVectors = mergeBatches( + makeBatches( + 5, + [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector({1, 2, 3, 4, 5}), + makeNullableFlatVector( + {10, std::nullopt, 30, std::nullopt, 50}), + }); + }), + makeBatches( + 5, + [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector({1, 2, 3, 4, 5}), + makeNullableFlatVector( + {std::nullopt, 20, 30, std::nullopt, 50}), + }); + }), + true); + + std::vector buildVectors = + makeBatches(5, [&](int32_t /*unused*/) { + return makeRowVector( + {makeFlatVector(128, [](vector_size_t row) { + if (row < 3) { + return row; + } + return row + 10; + })}); + }); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .numDrivers(numDrivers_) + .probeKeys({"c0"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u_c0"}) + .buildVectors(std::move(buildVectors)) + .buildProjections({"c0 AS u_c0"}) + .joinType(core::JoinType::kLeft) + .joinFilter("c1 + u_c0 > 0") + .joinOutputLayout({"c0", "c1", "u_c0"}) + .referenceQuery( + "SELECT * FROM t LEFT JOIN u ON (t.c0 = u.c0 AND t.c1 + u.c0 > 0)") + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, rightJoin) { + // Left side keys are [0, 1, 2,..20]. + std::vector probeVectors = mergeBatches( + makeBatches( + 3, + [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 137, [](auto row) { return row % 21; }, nullEvery(13)), + makeFlatVector(137, [](auto row) { return row; }), + }); + }), + makeBatches( + 3, + [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 234, + [](auto row) { return (row + 3) % 21; }, + nullEvery(13)), + makeFlatVector(234, [](auto row) { return row; }), + }); + }), + true); + + // Right side keys are [-3, -2, -1, 0, 1, 2, 3]. + std::vector buildVectors = + makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 123, [](auto row) { return -3 + row % 7; }, nullEvery(11)), + makeFlatVector( + 123, [](auto row) { return -111 + row * 2; }, nullEvery(13)), + }); + }); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .injectSpill(false) + .probeKeys({"c0"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u_c0"}) + .buildVectors(std::move(buildVectors)) + .buildProjections({"c0 AS u_c0", "c1 AS u_c1"}) + .joinType(core::JoinType::kRight) + .joinOutputLayout({"c0", "c1", "u_c1"}) + .referenceQuery( + "SELECT t.c0, t.c1, u.c1 FROM t RIGHT JOIN u ON t.c0 = u.c0") + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, rightJoinWithEmptyBuild) { + const std::vector finishOnEmptys = {false, true}; + for (const auto finishOnEmpty : finishOnEmptys) { + SCOPED_TRACE(fmt::format("finishOnEmpty: {}", finishOnEmpty)); + + // Left side keys are [0, 1, 2,..10]. + std::vector probeVectors = mergeBatches( + makeBatches( + 3, + [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 137, [](auto row) { return row % 11; }, nullEvery(13)), + makeFlatVector(137, [](auto row) { return row; }), + }); + }), + makeBatches( + 3, + [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 234, + [](auto row) { return (row + 3) % 11; }, + nullEvery(13)), + makeFlatVector(234, [](auto row) { return row; }), + }); + }), + true); + + // Right side keys are [-3, -2, -1, 0, 1, 2, 3]. + std::vector buildVectors = + makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 123, [](auto row) { return -3 + row % 7; }, nullEvery(11)), + makeFlatVector( + 123, [](auto row) { return -111 + row * 2; }, nullEvery(13)), + }); + }); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .hashProbeFinishEarlyOnEmptyBuild(finishOnEmpty) + .numDrivers(numDrivers_) + .injectSpill(false) + .probeKeys({"c0"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u_c0"}) + .buildVectors(std::move(buildVectors)) + .buildFilter("c0 > 100") + .buildProjections({"c0 AS u_c0", "c1 AS u_c1"}) + .joinType(core::JoinType::kRight) + .joinOutputLayout({"c1"}) + .referenceQuery("SELECT null LIMIT 0") + .checkSpillStats(false) + .run(); + } +} + +TEST_P(MultiThreadedHashJoinTest, rightJoinWithAllMatch) { + // Left side keys are [0, 1, 2,..20]. + std::vector probeVectors = mergeBatches( + makeBatches( + 3, + [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 137, [](auto row) { return row % 21; }, nullEvery(13)), + makeFlatVector(137, [](auto row) { return row; }), + }); + }), + makeBatches( + 3, + [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 234, + [](auto row) { return (row + 3) % 21; }, + nullEvery(13)), + makeFlatVector(234, [](auto row) { return row; }), + }); + }), + true); + + // Right side keys are [-3, -2, -1, 0, 1, 2, 3]. + std::vector buildVectors = + makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 123, [](auto row) { return -3 + row % 7; }, nullEvery(11)), + makeFlatVector( + 123, [](auto row) { return -111 + row * 2; }, nullEvery(13)), + }); + }); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .injectSpill(false) + .probeKeys({"c0"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u_c0"}) + .buildVectors(std::move(buildVectors)) + .buildFilter("c0 >= 0") + .buildProjections({"c0 AS u_c0", "c1 AS u_c1"}) + .joinType(core::JoinType::kRight) + .joinOutputLayout({"c0", "c1", "u_c1"}) + .referenceQuery( + "SELECT t.c0, t.c1, u.c1 FROM t RIGHT JOIN (SELECT * FROM u WHERE c0 >= 0) u ON t.c0 = u.c0") + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, rightJoinWithFilter) { + // Left side keys are [0, 1, 2,..20]. + std::vector probeVectors = mergeBatches( + makeBatches( + 3, + [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 137, [](auto row) { return row % 21; }, nullEvery(13)), + makeFlatVector(137, [](auto row) { return row; }), + }); + }), + makeBatches( + 3, + [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 234, + [](auto row) { return (row + 3) % 21; }, + nullEvery(13)), + makeFlatVector(234, [](auto row) { return row; }), + }); + }), + true); + + // Right side keys are [-3, -2, -1, 0, 1, 2, 3]. + std::vector buildVectors = + makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 123, [](auto row) { return -3 + row % 7; }, nullEvery(11)), + makeFlatVector( + 123, [](auto row) { return -111 + row * 2; }, nullEvery(13)), + }); + }); + + // Filter with passed rows. + { + auto testProbeVectors = probeVectors; + auto testBuildVectors = buildVectors; + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .injectSpill(false) + .probeKeys({"c0"}) + .probeVectors(std::move(testProbeVectors)) + .buildKeys({"u_c0"}) + .buildVectors(std::move(testBuildVectors)) + .buildProjections({"c0 AS u_c0", "c1 AS u_c1"}) + .joinType(core::JoinType::kRight) + .joinFilter("(c1 + u_c1) % 2 = 1") + .joinOutputLayout({"c0", "c1", "u_c1"}) + .referenceQuery( + "SELECT t.c0, t.c1, u.c1 FROM t RIGHT JOIN u ON t.c0 = u.c0 AND (t.c1 + u.c1) % 2 = 1") + .run(); + } + + // Filter without passed rows. + { + auto testProbeVectors = probeVectors; + auto testBuildVectors = buildVectors; + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .injectSpill(false) + .probeKeys({"c0"}) + .probeVectors(std::move(testProbeVectors)) + .buildKeys({"u_c0"}) + .buildVectors(std::move(testBuildVectors)) + .buildProjections({"c0 AS u_c0", "c1 AS u_c1"}) + .joinType(core::JoinType::kRight) + .joinFilter("(c1 + u_c1) % 2 = 3") + .joinOutputLayout({"c0", "c1", "u_c1"}) + .referenceQuery( + "SELECT t.c0, t.c1, u.c1 FROM t RIGHT JOIN u ON t.c0 = u.c0 AND (t.c1 + u.c1) % 2 = 3") + .run(); + } +} + +TEST_P(MultiThreadedHashJoinTest, fullJoin) { + // Left side keys are [0, 1, 2,..20]. + std::vector probeVectors = mergeBatches( + makeBatches( + 3, + [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 213, [](auto row) { return row % 21; }, nullEvery(13)), + makeFlatVector(213, [](auto row) { return row; }), + }); + }), + makeBatches( + 2, + [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 137, + [](auto row) { return (row + 3) % 21; }, + nullEvery(13)), + makeFlatVector(137, [](auto row) { return row; }), + }); + }), + true); + + // Right side keys are [-3, -2, -1, + // 0, 1, 2, 3]. + std::vector buildVectors = + makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 123, [](auto row) { return -3 + row % 7; }, nullEvery(11)), + makeFlatVector( + 123, [](auto row) { return -111 + row * 2; }, nullEvery(13)), + }); + }); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .injectSpill(false) + .probeKeys({"c0"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u_c0"}) + .buildVectors(std::move(buildVectors)) + .buildProjections({"c0 AS u_c0", "c1 AS u_c1"}) + .joinType(core::JoinType::kFull) + .joinOutputLayout({"c0", "c1", "u_c1"}) + .referenceQuery( + "SELECT t.c0, t.c1, u.c1 FROM t FULL OUTER JOIN u ON t.c0 = u.c0") + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, fullJoinWithEmptyBuild) { + const std::vector finishOnEmptys = {false, true}; + for (const auto finishOnEmpty : finishOnEmptys) { + SCOPED_TRACE(fmt::format("finishOnEmpty: {}", finishOnEmpty)); + + // Left side keys are [0, 1, 2,..10]. + std::vector probeVectors = mergeBatches( + makeBatches( + 3, + [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 213, [](auto row) { return row % 11; }, nullEvery(13)), + makeFlatVector(213, [](auto row) { return row; }), + }); + }), + makeBatches( + 2, + [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 137, + [](auto row) { return (row + 3) % 11; }, + nullEvery(13)), + makeFlatVector(137, [](auto row) { return row; }), + }); + }), + true); + + // Right side keys are [-3, -2, -1, 0, 1, 2, 3]. + std::vector buildVectors = + makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 123, [](auto row) { return -3 + row % 7; }, nullEvery(11)), + makeFlatVector( + 123, [](auto row) { return -111 + row * 2; }, nullEvery(13)), + }); + }); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .hashProbeFinishEarlyOnEmptyBuild(finishOnEmpty) + .numDrivers(numDrivers_) + .injectSpill(false) + .probeKeys({"c0"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u_c0"}) + .buildVectors(std::move(buildVectors)) + .buildFilter("c0 > 100") + .buildProjections({"c0 AS u_c0", "c1 AS u_c1"}) + .joinType(core::JoinType::kFull) + .joinOutputLayout({"c1"}) + .referenceQuery( + "SELECT t.c1 FROM t FULL OUTER JOIN (SELECT * FROM u WHERE c0 > 100) u ON t.c0 = u.c0") + .checkSpillStats(false) + .run(); + } +} + +TEST_P(MultiThreadedHashJoinTest, fullJoinWithNoMatch) { + // Left side keys are [0, 1, 2,..10]. + std::vector probeVectors = mergeBatches( + makeBatches( + 3, + [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 213, [](auto row) { return row % 11; }, nullEvery(13)), + makeFlatVector(213, [](auto row) { return row; }), + }); + }), + makeBatches( + 2, + [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 137, + [](auto row) { return (row + 3) % 11; }, + nullEvery(13)), + makeFlatVector(137, [](auto row) { return row; }), + }); + }), + true); + + // Right side keys are [-3, -2, -1, 0, 1, 2, 3]. + std::vector buildVectors = + makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 123, [](auto row) { return -3 + row % 7; }, nullEvery(11)), + makeFlatVector( + 123, [](auto row) { return -111 + row * 2; }, nullEvery(13)), + }); + }); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .injectSpill(false) + .probeKeys({"c0"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u_c0"}) + .buildVectors(std::move(buildVectors)) + .buildFilter("c0 < 0") + .buildProjections({"c0 AS u_c0", "c1 AS u_c1"}) + .joinType(core::JoinType::kFull) + .joinOutputLayout({"c1"}) + .referenceQuery( + "SELECT t.c1 FROM t FULL OUTER JOIN (SELECT * FROM u WHERE c0 < 0) u ON t.c0 = u.c0") + .run(); +} + +TEST_P(MultiThreadedHashJoinTest, fullJoinWithFilters) { + // Left side keys are [0, 1, 2,..10]. + std::vector probeVectors = mergeBatches( + makeBatches( + 3, + [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 213, [](auto row) { return row % 11; }, nullEvery(13)), + makeFlatVector(213, [](auto row) { return row; }), + }); + }), + makeBatches( + 2, + [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 137, + [](auto row) { return (row + 3) % 11; }, + nullEvery(13)), + makeFlatVector(137, [](auto row) { return row; }), + }); + }), + true); + + // Right side keys are [-3, -2, -1, 0, 1, 2, 3]. + std::vector buildVectors = + makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector( + 123, [](auto row) { return -3 + row % 7; }, nullEvery(11)), + makeFlatVector( + 123, [](auto row) { return -111 + row * 2; }, nullEvery(13)), + }); + }); + + // Filter with passed rows. + { + auto testProbeVectors = probeVectors; + auto testBuildVectors = buildVectors; + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .injectSpill(false) + .probeKeys({"c0"}) + .probeVectors(std::move(testProbeVectors)) + .buildKeys({"u_c0"}) + .buildVectors(std::move(testBuildVectors)) + .buildProjections({"c0 AS u_c0", "c1 AS u_c1"}) + .joinType(core::JoinType::kFull) + .joinFilter("(c1 + u_c1) % 2 = 1") + .joinOutputLayout({"c0", "c1", "u_c1"}) + .referenceQuery( + "SELECT t.c0, t.c1, u.c1 FROM t FULL OUTER JOIN u ON t.c0 = u.c0 AND (t.c1 + u.c1) % 2 = 1") + .run(); + } + + // Filter without passed rows. + { + auto testProbeVectors = probeVectors; + auto testBuildVectors = buildVectors; + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .injectSpill(false) + .probeKeys({"c0"}) + .probeVectors(std::move(testProbeVectors)) + .buildKeys({"u_c0"}) + .buildVectors(std::move(testBuildVectors)) + .buildProjections({"c0 AS u_c0", "c1 AS u_c1"}) + .joinType(core::JoinType::kFull) + .joinFilter("(c1 + u_c1) % 2 = 3") + .joinOutputLayout({"c0", "c1", "u_c1"}) + .referenceQuery( + "SELECT t.c0, t.c1, u.c1 FROM t FULL OUTER JOIN u ON t.c0 = u.c0 AND (t.c1 + u.c1) % 2 = 3") + .run(); + } +} + +TEST_P(MultiThreadedHashJoinTest, noSpillLevelLimit) { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .injectSpill(false) + .keyTypes({INTEGER()}) + .probeVectors(1600, 5) + .buildVectors(1500, 5) + .referenceQuery( + "SELECT t_k0, t_data, u_k0, u_data FROM t, u WHERE t.t_k0 = u.u_k0") + .maxSpillLevel(-1) + .config(core::QueryConfig::kSpillStartPartitionBit, "51") + .config(core::QueryConfig::kSpillNumPartitionBits, "3") + .checkSpillStats(false) + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + if (!hasSpill) { + return; + } + ASSERT_EQ(maxHashBuildSpillLevel(*task), 3); + }) + .run(); +} + +// Verify that dynamic filter pushed down is turned off for null-aware right +// semi project join. +TEST_F(HashJoinTest, nullAwareRightSemiProjectOverScan) { + std::vector probes; + std::vector builds; + // Matches present: + probes.push_back(makeRowVector( + {"t0"}, + { + makeNullableFlatVector({1, std::nullopt, 2}), + })); + builds.push_back(makeRowVector( + {"u0"}, + { + makeNullableFlatVector({1, 2, 3, std::nullopt}), + })); + + // No matches present: + probes.push_back(makeRowVector( + {"t0"}, + { + makeFlatVector({5, 6}), + })); + builds.push_back(makeRowVector( + {"u0"}, + { + makeNullableFlatVector({1, 2, 3, std::nullopt}), + })); + + for (int i = 0; i < probes.size(); i++) { + RowVectorPtr& probe = probes[i]; + RowVectorPtr& build = builds[i]; + std::shared_ptr probeFile = TempFilePath::create(); + writeToFile(probeFile->getPath(), {probe}); + + std::shared_ptr buildFile = TempFilePath::create(); + writeToFile(buildFile->getPath(), {build}); + + createDuckDbTable("t", {probe}); + createDuckDbTable("u", {build}); + + core::PlanNodeId probeScanId; + core::PlanNodeId buildScanId; + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .tableScan(asRowType(probe->type())) + .capturePlanNodeId(probeScanId) + .hashJoin( + {"t0"}, + {"u0"}, + PlanBuilder(planNodeIdGenerator) + .tableScan(asRowType(build->type())) + .capturePlanNodeId(buildScanId) + .planNode(), + "", + {"u0", "match"}, + core::JoinType::kRightSemiProject, + true /*nullAware*/) + .planNode(); + + SplitInput splitInput = { + {probeScanId, + {exec::Split(makeHiveConnectorSplit(probeFile->getPath()))}}, + {buildScanId, + {exec::Split(makeHiveConnectorSplit(buildFile->getPath()))}}, + }; + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(plan) + .inputSplits(splitInput) + .checkSpillStats(false) + .referenceQuery("SELECT u0, u0 IN (SELECT t0 FROM t) FROM u") + .run(); + } +} + +TEST_F(HashJoinTest, duplicateJoinKeys) { + auto leftVectors = makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector({ + makeNullableFlatVector( + {1, 2, 2, 3, 3, std::nullopt, 4, 5, 5, 6, 7}), + makeNullableFlatVector( + {1, 2, 2, std::nullopt, 3, 3, 4, 5, 5, 6, 8}), + }); + }); + + auto rightVectors = makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector({ + makeNullableFlatVector({1, 1, 3, 4, std::nullopt, 5, 7, 8}), + makeNullableFlatVector({1, 1, 3, 4, 5, std::nullopt, 7, 8}), + }); + }); + + createDuckDbTable("t", leftVectors); + createDuckDbTable("u", rightVectors); + + auto planNodeIdGenerator = std::make_shared(); + + auto assertPlan = [&](const std::vector& leftProject, + const std::vector& leftKeys, + const std::vector& rightProject, + const std::vector& rightKeys, + const std::vector& outputLayout, + core::JoinType joinType, + const std::string& query) { + auto plan = PlanBuilder(planNodeIdGenerator) + .values(leftVectors) + .project(leftProject) + .hashJoin( + leftKeys, + rightKeys, + PlanBuilder(planNodeIdGenerator) + .values(rightVectors) + .project(rightProject) + .planNode(), + "", + outputLayout, + joinType) + .planNode(); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(plan) + .injectSpill(false) + .checkSpillStats(false) + .referenceQuery(query) + .run(); + }; + + std::vector> joins = { + {core::JoinType::kInner, "INNER JOIN"}, + {core::JoinType::kLeft, "LEFT JOIN"}, + {core::JoinType::kRight, "RIGHT JOIN"}, + {core::JoinType::kFull, "FULL OUTER JOIN"}}; + + for (const auto& [joinType, joinTypeSql] : joins) { + // Duplicate keys on the build side. + assertPlan( + {"c0 AS t0", "c1 as t1"}, // leftProject + {"t0", "t1"}, // leftKeys + {"c0 AS u0"}, // rightProject + {"u0", "u0"}, // rightKeys + {"t0", "t1", "u0"}, // outputLayout + joinType, + "SELECT t.c0, t.c1, u.c0 FROM t " + joinTypeSql + + " u ON t.c0 = u.c0 and t.c1 = u.c0"); + } + + for (const auto& [joinType, joinTypeSql] : joins) { + // Duplicated keys on the probe side. + assertPlan( + {"c0 AS t0"}, // leftProject + {"t0", "t0"}, // leftKeys + {"c0 AS u0", "c1 AS u1"}, // rightProject + {"u0", "u1"}, // rightKeys + {"t0", "u0", "u1"}, // outputLayout + joinType, + "SELECT t.c0, u.c0, u.c1 FROM t " + joinTypeSql + + " u ON t.c0 = u.c0 and t.c0 = u.c1"); + } +} + +TEST_F(HashJoinTest, semiProject) { + // Some keys have multiple rows: 2, 3, 5. + auto probeVectors = makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector({1, 2, 2, 3, 3, 3, 4, 5, 5, 6, 7}), + makeFlatVector({10, 20, 21, 30, 31, 32, 40, 50, 51, 60, 70}), + }); + }); + + // Some keys are missing: 2, 6. + // Some have multiple rows: 1, 5. + // Some keys are not present on probe side: 8. + auto buildVectors = makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector({ + makeFlatVector({1, 1, 3, 4, 5, 5, 7, 8}), + makeFlatVector({100, 101, 300, 400, 500, 501, 700, 800}), + }); + }); + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors) + .project({"c0 AS t0", "c1 AS t1"}) + .hashJoin( + {"t0"}, + {"u0"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors) + .project({"c0 AS u0", "c1 AS u1"}) + .planNode(), + "", + {"t0", "t1", "match"}, + core::JoinType::kLeftSemiProject) + .planNode(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .planNode(plan) + .referenceQuery( + "SELECT t.c0, t.c1, EXISTS (SELECT * FROM u WHERE t.c0 = u.c0) FROM t") + .run(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .planNode(flipJoinSides(plan)) + .referenceQuery( + "SELECT t.c0, t.c1, EXISTS (SELECT * FROM u WHERE t.c0 = u.c0) FROM t") + .run(); + + // With extra filter. + planNodeIdGenerator = std::make_shared(); + plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors) + .project({"c0 AS t0", "c1 AS t1"}) + .hashJoin( + {"t0"}, + {"u0"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors) + .project({"c0 AS u0", "c1 AS u1"}) + .planNode(), + "t1 * 10 <> u1", + {"t0", "t1", "match"}, + core::JoinType::kLeftSemiProject) + .planNode(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .planNode(plan) + .referenceQuery( + "SELECT t.c0, t.c1, EXISTS (SELECT * FROM u WHERE t.c0 = u.c0 AND t.c1 * 10 <> u.c1) FROM t") + .run(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .planNode(flipJoinSides(plan)) + .referenceQuery( + "SELECT t.c0, t.c1, EXISTS (SELECT * FROM u WHERE t.c0 = u.c0 AND t.c1 * 10 <> u.c1) FROM t") + .run(); + + // Empty build side. + planNodeIdGenerator = std::make_shared(); + plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors) + .project({"c0 AS t0", "c1 AS t1"}) + .hashJoin( + {"t0"}, + {"u0"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors) + .project({"c0 AS u0", "c1 AS u1"}) + .filter("u0 < 0") + .planNode(), + "", + {"t0", "t1", "match"}, + core::JoinType::kLeftSemiProject) + .planNode(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .planNode(plan) + .referenceQuery( + "SELECT t.c0, t.c1, EXISTS (SELECT * FROM u WHERE u.c0 < 0 AND t.c0 = u.c0) FROM t") + // NOTE: there is no spilling in empty build test case as all the + // build-side rows have been filtered out. + .run(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .planNode(flipJoinSides(plan)) + .referenceQuery( + "SELECT t.c0, t.c1, EXISTS (SELECT * FROM u WHERE u.c0 < 0 AND t.c0 = u.c0) FROM t") + // NOTE: there is no spilling in empty build test case as all the + // build-side rows have been filtered out. + .checkSpillStats(false) + .run(); +} + +TEST_F(HashJoinTest, semiProjectWithNullKeys) { + // Some keys have multiple rows: 2, 3, 5. + auto probeVectors = makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector( + {"t0", "t1"}, + { + makeNullableFlatVector( + {1, 2, 2, 3, 3, 3, 4, std::nullopt, 5, 5, 6, 7}), + makeFlatVector( + {10, 20, 21, 30, 31, 32, 40, -1, 50, 51, 60, 70}), + }); + }); + + // Some keys are missing: 2, 6. + // Some have multiple rows: 1, 5. + // Some keys are not present on probe side: 8. + auto buildVectors = makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector( + {"u0", "u1"}, + { + makeNullableFlatVector( + {1, 1, 3, 4, std::nullopt, 5, 5, 7, 8}), + makeFlatVector( + {100, 101, 300, 400, -100, 500, 501, 700, 800}), + }); + }); + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + auto makePlan = [&](bool nullAware, + const std::string& probeFilter = "", + const std::string& buildFilter = "") { + auto planNodeIdGenerator = std::make_shared(); + return PlanBuilder(planNodeIdGenerator) + .values(probeVectors) + .optionalFilter(probeFilter) + .hashJoin( + {"t0"}, + {"u0"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors) + .optionalFilter(buildFilter) + .planNode(), + "", + {"t0", "t1", "match"}, + core::JoinType::kLeftSemiProject, + nullAware) + .planNode(); + }; + + // Null join keys on both sides. + auto plan = makePlan(false /*nullAware*/); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .planNode(plan) + .referenceQuery( + "SELECT t0, t1, EXISTS (SELECT * FROM u WHERE u0 = t0) FROM t") + .run(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .planNode(flipJoinSides(plan)) + .referenceQuery( + "SELECT t0, t1, EXISTS (SELECT * FROM u WHERE u0 = t0) FROM t") + .run(); + + plan = makePlan(true /*nullAware*/); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .planNode(plan) + .referenceQuery("SELECT t0, t1, t0 IN (SELECT u0 FROM u) FROM t") + .run(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .planNode(flipJoinSides(plan)) + .referenceQuery("SELECT t0, t1, t0 IN (SELECT u0 FROM u) FROM t") + .run(); + + // Null join keys on build side-only. + plan = makePlan(false /*nullAware*/, "t0 IS NOT NULL"); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .planNode(plan) + .referenceQuery( + "SELECT t0, t1, EXISTS (SELECT * FROM u WHERE u0 = t0) FROM t WHERE t0 IS NOT NULL") + .run(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .planNode(flipJoinSides(plan)) + .referenceQuery( + "SELECT t0, t1, EXISTS (SELECT * FROM u WHERE u0 = t0) FROM t WHERE t0 IS NOT NULL") + .run(); + + plan = makePlan(true /*nullAware*/, "t0 IS NOT NULL"); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .planNode(plan) + .referenceQuery( + "SELECT t0, t1, t0 IN (SELECT u0 FROM u) FROM t WHERE t0 IS NOT NULL") + .run(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .planNode(flipJoinSides(plan)) + .referenceQuery( + "SELECT t0, t1, t0 IN (SELECT u0 FROM u) FROM t WHERE t0 IS NOT NULL") + .run(); + + // Null join keys on probe side-only. + plan = makePlan(false /*nullAware*/, "", "u0 IS NOT NULL"); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .planNode(plan) + .referenceQuery( + "SELECT t0, t1, EXISTS (SELECT * FROM u WHERE u0 = t0 AND u0 IS NOT NULL) FROM t") + .run(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .planNode(flipJoinSides(plan)) + .referenceQuery( + "SELECT t0, t1, EXISTS (SELECT * FROM u WHERE u0 = t0 AND u0 IS NOT NULL) FROM t") + .run(); + + plan = makePlan(true /*nullAware*/, "", "u0 IS NOT NULL"); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .planNode(plan) + .referenceQuery( + "SELECT t0, t1, t0 IN (SELECT u0 FROM u WHERE u0 IS NOT NULL) FROM t") + .run(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .planNode(flipJoinSides(plan)) + .referenceQuery( + "SELECT t0, t1, t0 IN (SELECT u0 FROM u WHERE u0 IS NOT NULL) FROM t") + .run(); + + // Empty build side. + plan = makePlan(false /*nullAware*/, "", "u0 < 0"); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, executor_.get()) + .planNode(plan) + .injectSpill(false) + .checkSpillStats(false) + .referenceQuery( + "SELECT t0, t1, EXISTS (SELECT * FROM u WHERE u0 = t0 AND u0 < 0) FROM t") + .run(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, executor_.get()) + .planNode(flipJoinSides(plan)) + .injectSpill(false) + .checkSpillStats(false) + .referenceQuery( + "SELECT t0, t1, EXISTS (SELECT * FROM u WHERE u0 = t0 AND u0 < 0) FROM t") + .run(); + + plan = makePlan(true /*nullAware*/, "", "u0 < 0"); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, executor_.get()) + .planNode(plan) + .injectSpill(false) + .checkSpillStats(false) + .referenceQuery( + "SELECT t0, t1, t0 IN (SELECT u0 FROM u WHERE u0 < 0) FROM t") + .run(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, executor_.get()) + .planNode(flipJoinSides(plan)) + .injectSpill(false) + .checkSpillStats(false) + .referenceQuery( + "SELECT t0, t1, t0 IN (SELECT u0 FROM u WHERE u0 < 0) FROM t") + .run(); + + // Build side with all rows having null join keys. + plan = makePlan(false /*nullAware*/, "", "u0 IS NULL"); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, executor_.get()) + .planNode(plan) + .injectSpill(false) + .checkSpillStats(false) + .referenceQuery( + "SELECT t0, t1, EXISTS (SELECT * FROM u WHERE u0 = t0 AND u0 IS NULL) FROM t") + .run(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, executor_.get()) + .planNode(flipJoinSides(plan)) + .injectSpill(false) + .checkSpillStats(false) + .referenceQuery( + "SELECT t0, t1, EXISTS (SELECT * FROM u WHERE u0 = t0 AND u0 IS NULL) FROM t") + .run(); + + plan = makePlan(true /*nullAware*/, "", "u0 IS NULL"); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, executor_.get()) + .planNode(plan) + .injectSpill(false) + .checkSpillStats(false) + .referenceQuery( + "SELECT t0, t1, t0 IN (SELECT u0 FROM u WHERE u0 IS NULL) FROM t") + .run(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, executor_.get()) + .planNode(flipJoinSides(plan)) + .injectSpill(false) + .checkSpillStats(false) + .referenceQuery( + "SELECT t0, t1, t0 IN (SELECT u0 FROM u WHERE u0 IS NULL) FROM t") + .run(); +} + +TEST_F(HashJoinTest, semiProjectWithFilter) { + auto probeVectors = makeBatches(3, [&](auto /*unused*/) { + return makeRowVector( + {"t0", "t1"}, + { + makeNullableFlatVector({1, 2, 3, std::nullopt, 5}), + makeFlatVector({10, 20, 30, 40, 50}), + }); + }); + + auto buildVectors = makeBatches(3, [&](auto /*unused*/) { + return makeRowVector( + {"u0", "u1"}, + { + makeNullableFlatVector({1, 2, 3, std::nullopt}), + makeFlatVector({11, 22, 33, 44}), + }); + }); + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + auto makePlan = [&](bool nullAware, const std::string& filter) { + auto planNodeIdGenerator = std::make_shared(); + return PlanBuilder(planNodeIdGenerator) + .values(probeVectors) + .hashJoin( + {"t0"}, + {"u0"}, + PlanBuilder(planNodeIdGenerator).values(buildVectors).planNode(), + filter, + {"t0", "t1", "match"}, + core::JoinType::kLeftSemiProject, + nullAware) + .planNode(); + }; + + std::vector filters = { + "t1 <> u1", + "t1 < u1", + "t1 > u1", + "t1 is not null AND u1 is not null", + "t1 is null OR u1 is null", + }; + for (const auto& filter : filters) { + auto plan = makePlan(true /*nullAware*/, filter); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(plan) + .referenceQuery(fmt::format( + "SELECT t0, t1, t0 IN (SELECT u0 FROM u WHERE {}) FROM t", filter)) + .injectSpill(false) + .run(); + + plan = makePlan(false /*nullAware*/, filter); + + // DuckDB Exists operator returns NULL when u0 or t0 is NULL. We exclude + // these values. + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(plan) + .referenceQuery(fmt::format( + "SELECT t0, t1, EXISTS (SELECT * FROM u WHERE (u0 is not null OR t0 is not null) AND u0 = t0 AND {}) FROM t", + filter)) + .injectSpill(false) + .run(); + } +} + +TEST_F(HashJoinTest, nullAwareRightSemiProjectWithFilterNotAllowed) { + auto probe = makeRowVector(ROW({"t0", "t1"}, {INTEGER(), BIGINT()}), 10); + auto build = makeRowVector(ROW({"u0", "u1"}, {INTEGER(), BIGINT()}), 10); + + auto planNodeIdGenerator = std::make_shared(); + VELOX_ASSERT_THROW( + PlanBuilder(planNodeIdGenerator) + .values({probe}) + .hashJoin( + {"t0"}, + {"u0"}, + PlanBuilder(planNodeIdGenerator).values({build}).planNode(), + "t1 > u1", + {"u0", "u1", "match"}, + core::JoinType::kRightSemiProject, + true /* nullAware */), + "Null-aware right semi project join doesn't support extra filter"); +} + +TEST_F(HashJoinTest, leftSemiJoinWithExtraOutputCapacity) { + std::vector probeVectors; + std::vector buildVectors; + probeVectors.push_back(makeRowVector( + {"t0", "t1"}, + { + makeFlatVector({1, 2, 3, 4, 5, 6}), + makeFlatVector({10, 10, 10, 10, 10, 10}), + })); + + buildVectors.push_back(makeRowVector( + {"u0", "u1"}, + { + makeFlatVector({1, 1, 1, 1, 1}), + makeFlatVector({10, 10, 10, 10, 10}), + })); + buildVectors.push_back(makeRowVector( + {"u0", "u1"}, + { + makeFlatVector({2, 3, 4, 5, 6}), + makeFlatVector({10, 10, 10, 10, 10}), + })); + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + auto runQuery = [&](const std::string& query, + const std::string& filter, + core::JoinType joinType) { + auto planNodeIdGenerator = std::make_shared(); + std::vector outputLayout = {"t0", "t1"}; + if (joinType == core::JoinType::kLeftSemiProject) { + outputLayout.push_back("match"); + } + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors) + .hashJoin( + {"t0"}, + {"u0"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors) + .planNode(), + filter, + outputLayout, + joinType, + false) + .planNode(); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(plan) + .config(core::QueryConfig::kPreferredOutputBatchRows, "5") + .referenceQuery(query) + .injectSpill(false) + .run(); + }; + { + SCOPED_TRACE("left semi filter join"); + std::string filter = "t1 = u1"; + runQuery( + fmt::format( + "SELECT t0, t1 FROM t WHERE EXISTS (SELECT u0 FROM u WHERE t0 = u0 AND {})", + filter), + filter, + core::JoinType::kLeftSemiFilter); + } + + { + SCOPED_TRACE("left semi project join"); + std::string filter = "t1 <> u1"; + runQuery( + fmt::format( + "SELECT t0, t1, t0 IN (SELECT u0 FROM u WHERE {}) FROM t", filter), + filter, + core::JoinType::kLeftSemiProject); + } +} + +TEST_F(HashJoinTest, nullAwareMultiKeyNotAllowed) { + auto probe = makeRowVector( + ROW({"t0", "t1", "t2"}, {INTEGER(), BIGINT(), VARCHAR()}), 10); + auto build = makeRowVector( + ROW({"u0", "u1", "u2"}, {INTEGER(), BIGINT(), VARCHAR()}), 10); + + // Null-aware left semi project join. + auto planNodeIdGenerator = std::make_shared(); + VELOX_ASSERT_THROW( + PlanBuilder(planNodeIdGenerator) + .values({probe}) + .hashJoin( + {"t0", "t1"}, + {"u0", "u1"}, + PlanBuilder(planNodeIdGenerator).values({build}).planNode(), + "", + {"t0", "t1", "match"}, + core::JoinType::kLeftSemiProject, + true /* nullAware */), + "Null-aware joins allow only one join key"); + + // Null-aware right semi project join. + VELOX_ASSERT_THROW( + PlanBuilder(planNodeIdGenerator) + .values({probe}) + .hashJoin( + {"t0", "t1"}, + {"u0", "u1"}, + PlanBuilder(planNodeIdGenerator).values({build}).planNode(), + "", + {"u0", "u1", "match"}, + core::JoinType::kRightSemiProject, + true /* nullAware */), + "Null-aware joins allow only one join key"); + + // Null-aware anti join. + VELOX_ASSERT_THROW( + PlanBuilder(planNodeIdGenerator) + .values({probe}) + .hashJoin( + {"t0", "t1"}, + {"u0", "u1"}, + PlanBuilder(planNodeIdGenerator).values({build}).planNode(), + "", + {"t0", "t1"}, + core::JoinType::kAnti, + true /* nullAware */), + "Null-aware joins allow only one join key"); +} + +TEST_F(HashJoinTest, semiProjectOverLazyVectors) { + auto probeVectors = makeBatches(1, [&](auto /*unused*/) { + return makeRowVector( + {"t0", "t1"}, + { + makeFlatVector(1'000, [](auto row) { return row; }), + makeFlatVector(1'000, [](auto row) { return row * 10; }), + }); + }); + + auto buildVectors = makeBatches(3, [&](auto /*unused*/) { + return makeRowVector( + {"u0", "u1"}, + { + makeFlatVector( + 1'000, [](auto row) { return -100 + (row / 5); }), + makeFlatVector( + 1'000, [](auto row) { return -1000 + (row / 5) * 10; }), + }); + }); + + std::shared_ptr probeFile = TempFilePath::create(); + writeToFile(probeFile->getPath(), probeVectors); + + std::shared_ptr buildFile = TempFilePath::create(); + writeToFile(buildFile->getPath(), buildVectors); + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + core::PlanNodeId probeScanId; + core::PlanNodeId buildScanId; + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .tableScan(asRowType(probeVectors[0]->type())) + .capturePlanNodeId(probeScanId) + .hashJoin( + {"t0"}, + {"u0"}, + PlanBuilder(planNodeIdGenerator) + .tableScan(asRowType(buildVectors[0]->type())) + .capturePlanNodeId(buildScanId) + .planNode(), + "", + {"t0", "t1", "match"}, + core::JoinType::kLeftSemiProject) + .planNode(); + + SplitInput splitInput = { + {probeScanId, + {exec::Split(makeHiveConnectorSplit(probeFile->getPath()))}}, + {buildScanId, + {exec::Split(makeHiveConnectorSplit(buildFile->getPath()))}}, + }; + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(plan) + .inputSplits(splitInput) + .checkSpillStats(false) + .referenceQuery("SELECT t0, t1, t0 IN (SELECT u0 FROM u) FROM t") + .run(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(flipJoinSides(plan)) + .inputSplits(splitInput) + .checkSpillStats(false) + .referenceQuery("SELECT t0, t1, t0 IN (SELECT u0 FROM u) FROM t") + .run(); + + // With extra filter. + planNodeIdGenerator = std::make_shared(); + plan = PlanBuilder(planNodeIdGenerator) + .tableScan(asRowType(probeVectors[0]->type())) + .capturePlanNodeId(probeScanId) + .hashJoin( + {"t0"}, + {"u0"}, + PlanBuilder(planNodeIdGenerator) + .tableScan(asRowType(buildVectors[0]->type())) + .capturePlanNodeId(buildScanId) + .planNode(), + "(t1 + u1) % 3 = 0", + {"t0", "t1", "match"}, + core::JoinType::kLeftSemiProject) + .planNode(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(plan) + .inputSplits(splitInput) + .checkSpillStats(false) + .referenceQuery( + "SELECT t0, t1, t0 IN (SELECT u0 FROM u WHERE (t1 + u1) % 3 = 0) FROM t") + .run(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(flipJoinSides(plan)) + .inputSplits(splitInput) + .checkSpillStats(false) + .referenceQuery( + "SELECT t0, t1, t0 IN (SELECT u0 FROM u WHERE (t1 + u1) % 3 = 0) FROM t") + .run(); +} + +VELOX_INSTANTIATE_TEST_SUITE_P( + HashJoinTest, + MultiThreadedHashJoinTest, + testing::ValuesIn(MultiThreadedHashJoinTest::getTestParams())); + +// TODO: try to parallelize the following test cases if possible. +TEST_F(HashJoinTest, memory) { + // Measures memory allocation in a 1:n hash join followed by + // projection and aggregation. We expect vectors to be mostly + // reused, except for t_k0 + 1, which is a dictionary after the + // join. + std::vector probeVectors = + makeBatches(10, [&](int32_t /*unused*/) { + return std::dynamic_pointer_cast( + BatchMaker::createBatch(probeType_, 1000, *pool_)); + }); + + // auto buildType = makeRowType(keyTypes, "u_"); + std::vector buildVectors = + makeBatches(10, [&](int32_t /*unused*/) { + return std::dynamic_pointer_cast( + BatchMaker::createBatch(buildType_, 1000, *pool_)); + }); + + auto planNodeIdGenerator = std::make_shared(); + CursorParameters params; + params.planNode = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, true) + .hashJoin( + {"t_k1"}, + {"u_k1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, true) + .planNode(), + "", + concat(probeType_->names(), buildType_->names())) + .project({"t_k1 % 1000 AS k1", "u_k1 % 1000 AS k2"}) + .singleAggregation({}, {"sum(k1)", "sum(k2)"}) + .planNode(); + params.queryCtx = core::QueryCtx::create(driverExecutor_.get()); + auto [taskCursor, rows] = readCursor(params); + EXPECT_GT(3'500, params.queryCtx->pool()->stats().numAllocs); + EXPECT_GT(40'000'000, params.queryCtx->pool()->stats().cumulativeBytes); +} + +TEST_F(HashJoinTest, lazyVectors) { + // a dataset of multiple row groups with multiple columns. We create + // different dictionary wrappings for different columns and load the + // rows in scope at different times. + auto probeVectors = makeBatches(3, [&](int32_t /*unused*/) { + return makeRowVector( + {makeFlatVector(3'000, [](auto row) { return row; }), + makeFlatVector(30'000, [](auto row) { return row % 23; }), + makeFlatVector(30'000, [](auto row) { return row % 31; }), + makeFlatVector(30'000, [](auto row) { + return StringView::makeInline(fmt::format("{} string", row % 43)); + })}); + }); + + std::vector buildVectors = + makeBatches(4, [&](int32_t /*unused*/) { + return makeRowVector( + {makeFlatVector(1'000, [](auto row) { return row * 3; }), + makeFlatVector( + 10'000, [](auto row) { return row % 31; })}); + }); + + std::vector> tempFiles; + + for (const auto& probeVector : probeVectors) { + tempFiles.push_back(TempFilePath::create()); + writeToFile(tempFiles.back()->getPath(), probeVector); + } + createDuckDbTable("t", probeVectors); + + for (const auto& buildVector : buildVectors) { + tempFiles.push_back(TempFilePath::create()); + writeToFile(tempFiles.back()->getPath(), buildVector); + } + createDuckDbTable("u", buildVectors); + + auto makeInputSplits = [&](const core::PlanNodeId& probeScanId, + const core::PlanNodeId& buildScanId) { + return [&] { + std::vector probeSplits; + for (int i = 0; i < probeVectors.size(); ++i) { + probeSplits.push_back( + exec::Split(makeHiveConnectorSplit(tempFiles[i]->getPath()))); + } + std::vector buildSplits; + for (int i = 0; i < buildVectors.size(); ++i) { + buildSplits.push_back(exec::Split(makeHiveConnectorSplit( + tempFiles[probeSplits.size() + i]->getPath()))); + } + SplitInput splits; + splits.emplace(probeScanId, probeSplits); + splits.emplace(buildScanId, buildSplits); + return splits; + }; + }; + + { + auto planNodeIdGenerator = std::make_shared(); + core::PlanNodeId probeScanId; + core::PlanNodeId buildScanId; + auto op = PlanBuilder(planNodeIdGenerator) + .tableScan(ROW({"c0", "c1"}, {INTEGER(), BIGINT()})) + .capturePlanNodeId(probeScanId) + .hashJoin( + {"c0"}, + {"c0"}, + PlanBuilder(planNodeIdGenerator) + .tableScan(ROW({"c0"}, {INTEGER()})) + .capturePlanNodeId(buildScanId) + .planNode(), + "", + {"c1"}) + .project({"c1 + 1"}) + .planNode(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .planNode(std::move(op)) + .makeInputSplits(makeInputSplits(probeScanId, buildScanId)) + .referenceQuery("SELECT t.c1 + 1 FROM t, u WHERE t.c0 = u.c0") + .run(); + } + + { + auto planNodeIdGenerator = std::make_shared(); + core::PlanNodeId probeScanId; + core::PlanNodeId buildScanId; + auto op = PlanBuilder(planNodeIdGenerator) + .tableScan( + ROW({"c0", "c1", "c2", "c3"}, + {INTEGER(), BIGINT(), INTEGER(), VARCHAR()})) + .capturePlanNodeId(probeScanId) + .filter("c2 < 29") + .hashJoin( + {"c0"}, + {"bc0"}, + PlanBuilder(planNodeIdGenerator) + .tableScan(ROW({"c0", "c1"}, {INTEGER(), BIGINT()})) + .capturePlanNodeId(buildScanId) + .project({"c0 as bc0", "c1 as bc1"}) + .planNode(), + "(c1 + bc1) % 33 < 27", + {"c1", "bc1", "c3"}) + .project({"c1 + 1", "bc1", "length(c3)"}) + .planNode(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .checkSpillStats(false) + .planNode(std::move(op)) + .makeInputSplits(makeInputSplits(probeScanId, buildScanId)) + .referenceQuery( + "SELECT t.c1 + 1, U.c1, length(t.c3) FROM t, u WHERE t.c0 = u.c0 and t.c2 < 29 and (t.c1 + u.c1) % 33 < 27") + .run(); + } +} + +TEST_F(HashJoinTest, lazyVectorNotLoadedInFilter) { + // Ensure that if lazy vectors are temporarily wrapped during a filter's + // execution and remain unloaded, the temporary wrap is promptly + // discarded. This precaution prevents the generation of the probe's output + // from wrapping an unloaded vector while the temporary wrap is + // still alive. + // This is done by generating a sufficiently small batch to allow the lazy + // vector to remain unloaded, as it doesn't need to be split between batches. + // Then we use a filter that skips the execution of the expression containing + // the lazy vector, thereby avoiding its loading. + + testLazyVectorsWithFilter( + core::JoinType::kInner, + "c1 >= 0 OR c2 > 0", + {"c1", "c2"}, + "SELECT t.c1, t.c2 FROM t, u WHERE t.c0 = u.c0"); +} + +TEST_F(HashJoinTest, lazyVectorPartiallyLoadedInFilterLeftJoin) { + // Test the case where a filter loads a subset of the rows that will be output + // from a column on the probe side. + + testLazyVectorsWithFilter( + core::JoinType::kLeft, + "c1 > 0 AND c2 > 0", + {"c1", "c2"}, + "SELECT t.c1, t.c2 FROM t LEFT JOIN u ON t.c0 = u.c0 AND (c1 > 0 AND c2 > 0)"); +} + +TEST_F(HashJoinTest, lazyVectorPartiallyLoadedInFilterFullJoin) { + // Test the case where a filter loads a subset of the rows that will be output + // from a column on the probe side. + + testLazyVectorsWithFilter( + core::JoinType::kFull, + "c1 > 0 AND c2 > 0", + {"c1", "c2"}, + "SELECT t.c1, t.c2 FROM t FULL OUTER JOIN u ON t.c0 = u.c0 AND (c1 > 0 AND c2 > 0)"); +} + +TEST_F(HashJoinTest, lazyVectorPartiallyLoadedInFilterLeftSemiProject) { + // Test the case where a filter loads a subset of the rows that will be output + // from a column on the probe side. + + testLazyVectorsWithFilter( + core::JoinType::kLeftSemiProject, + "c1 > 0 AND c2 > 0", + {"c1", "c2", "match"}, + "SELECT t.c1, t.c2, EXISTS (SELECT * FROM u WHERE t.c0 = u.c0 AND (t.c1 > 0 AND t.c2 > 0)) FROM t"); +} + +TEST_F(HashJoinTest, lazyVectorPartiallyLoadedInFilterAntiJoin) { + // Test the case where a filter loads a subset of the rows that will be output + // from a column on the probe side. + + testLazyVectorsWithFilter( + core::JoinType::kAnti, + "c1 > 0 AND c2 > 0", + {"c1", "c2"}, + "SELECT t.c1, t.c2 FROM t WHERE NOT EXISTS (SELECT * FROM u WHERE t.c0 = u.c0 AND (t.c1 > 0 AND t.c2 > 0))"); +} + +TEST_F(HashJoinTest, lazyVectorPartiallyLoadedInFilterInnerJoin) { + // Test the case where a filter loads a subset of the rows that will be output + // from a column on the probe side. + + testLazyVectorsWithFilter( + core::JoinType::kInner, + "not (c1 < 15 and c2 >= 0)", + {"c1", "c2"}, + "SELECT t.c1, t.c2 FROM t, u WHERE t.c0 = u.c0 AND NOT (c1 < 15 AND c2 >= 0)"); +} + +TEST_F(HashJoinTest, lazyVectorPartiallyLoadedInFilterLeftSemiFilter) { + // Test the case where a filter loads a subset of the rows that will be output + // from a column on the probe side. + + testLazyVectorsWithFilter( + core::JoinType::kLeftSemiFilter, + "not (c1 < 15 and c2 >= 0)", + {"c1", "c2"}, + "SELECT t.c1, t.c2 FROM t WHERE c0 IN (SELECT u.c0 FROM u WHERE t.c0 = u.c0 AND NOT (t.c1 < 15 AND t.c2 >= 0))"); +} + +// Disabled because dynamic filters are not supported yet in velox-cudf. +TEST_F(HashJoinTest, DISABLED_dynamicFilters) { + const int32_t numSplits = 10; + const int32_t numRowsProbe = 333; + const int32_t numRowsBuild = 100; + + std::vector probeVectors; + probeVectors.reserve(numSplits); + + std::vector> tempFiles; + for (int32_t i = 0; i < numSplits; ++i) { + auto rowVector = makeRowVector({ + makeFlatVector( + numRowsProbe, [&](auto row) { return row - i * 10; }), + makeFlatVector(numRowsProbe, [](auto row) { return row; }), + }); + probeVectors.push_back(rowVector); + tempFiles.push_back(TempFilePath::create()); + writeToFile(tempFiles.back()->getPath(), rowVector); + } + auto makeInputSplits = [&](const core::PlanNodeId& nodeId) { + return [&] { + std::vector probeSplits; + for (auto& file : tempFiles) { + probeSplits.push_back( + exec::Split(makeHiveConnectorSplit(file->getPath()))); + } + SplitInput splits; + splits.emplace(nodeId, probeSplits); + return splits; + }; + }; + + // 100 key values in [35, 233] range. + std::vector buildVectors; + for (int i = 0; i < 5; ++i) { + buildVectors.push_back(makeRowVector({ + makeFlatVector( + numRowsBuild / 5, + [i](auto row) { return 35 + 2 * (row + i * numRowsBuild / 5); }), + makeFlatVector(numRowsBuild / 5, [](auto row) { return row; }), + })); + } + std::vector keyOnlyBuildVectors; + for (int i = 0; i < 5; ++i) { + keyOnlyBuildVectors.push_back( + makeRowVector({makeFlatVector(numRowsBuild / 5, [i](auto row) { + return 35 + 2 * (row + i * numRowsBuild / 5); + })})); + } + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + auto probeType = ROW({"c0", "c1"}, {INTEGER(), BIGINT()}); + + auto planNodeIdGenerator = std::make_shared(); + + auto buildSide = PlanBuilder(planNodeIdGenerator, pool_.get()) + .values(buildVectors) + .project({"c0 AS u_c0", "c1 AS u_c1"}) + .planNode(); + auto keyOnlyBuildSide = PlanBuilder(planNodeIdGenerator, pool_.get()) + .values(keyOnlyBuildVectors) + .project({"c0 AS u_c0"}) + .planNode(); + + // Basic push-down. + { + // Inner join. + core::PlanNodeId probeScanId; + core::PlanNodeId joinId; + auto op = PlanBuilder(planNodeIdGenerator, pool_.get()) + .tableScan(probeType) + .capturePlanNodeId(probeScanId) + .hashJoin( + {"c0"}, + {"u_c0"}, + buildSide, + "", + {"c0", "c1", "u_c1"}, + core::JoinType::kInner) + .capturePlanNodeId(joinId) + .project({"c0", "c1 + 1", "c1 + u_c1"}) + .planNode(); + { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .makeInputSplits(makeInputSplits(probeScanId)) + .referenceQuery( + "SELECT t.c0, t.c1 + 1, t.c1 + u.c1 FROM t, u WHERE t.c0 = u.c0") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + SCOPED_TRACE(fmt::format("hasSpill:{}", hasSpill)); + auto planStats = toPlanStats(task->taskStats()); + if (hasSpill) { + // Dynamic filtering should be disabled with spilling triggered. + ASSERT_EQ(0, getFiltersProduced(task, 1).sum); + ASSERT_EQ(0, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_TRUE(planStats.at(probeScanId).dynamicFilterStats.empty()); + } else { + ASSERT_EQ(1, getFiltersProduced(task, 1).sum); + ASSERT_EQ(1, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(0, getReplacedWithFilterRows(task, 1).sum); + ASSERT_LT(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_EQ( + planStats.at(probeScanId).dynamicFilterStats.producerNodeIds, + std::unordered_set({joinId})); + } + }) + .run(); + } + + // Left semi join. + op = PlanBuilder(planNodeIdGenerator, pool_.get()) + .tableScan(probeType) + .capturePlanNodeId(probeScanId) + .hashJoin( + {"c0"}, + {"u_c0"}, + buildSide, + "", + {"c0", "c1"}, + core::JoinType::kLeftSemiFilter) + .capturePlanNodeId(joinId) + .project({"c0", "c1 + 1"}) + .planNode(); + + { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .makeInputSplits(makeInputSplits(probeScanId)) + .referenceQuery( + "SELECT t.c0, t.c1 + 1 FROM t WHERE t.c0 IN (SELECT c0 FROM u)") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + SCOPED_TRACE(fmt::format("hasSpill:{}", hasSpill)); + auto planStats = toPlanStats(task->taskStats()); + if (hasSpill) { + // Dynamic filtering should be disabled with spilling triggered. + ASSERT_EQ(0, getFiltersProduced(task, 1).sum); + ASSERT_EQ(0, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(0, getReplacedWithFilterRows(task, 1).sum); + ASSERT_EQ(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_TRUE(planStats.at(probeScanId).dynamicFilterStats.empty()); + } else { + ASSERT_EQ(1, getFiltersProduced(task, 1).sum); + ASSERT_EQ(1, getFiltersAccepted(task, 0).sum); + ASSERT_GT(getReplacedWithFilterRows(task, 1).sum, 0); + ASSERT_LT(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_EQ( + planStats.at(probeScanId).dynamicFilterStats.producerNodeIds, + std::unordered_set({joinId})); + } + }) + .run(); + } + + // Right semi join. + op = PlanBuilder(planNodeIdGenerator, pool_.get()) + .tableScan(probeType) + .capturePlanNodeId(probeScanId) + .hashJoin( + {"c0"}, + {"u_c0"}, + buildSide, + "", + {"u_c0", "u_c1"}, + core::JoinType::kRightSemiFilter) + .capturePlanNodeId(joinId) + .project({"u_c0", "u_c1 + 1"}) + .planNode(); + + { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .makeInputSplits(makeInputSplits(probeScanId)) + .referenceQuery( + "SELECT u.c0, u.c1 + 1 FROM u WHERE u.c0 IN (SELECT c0 FROM t)") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + SCOPED_TRACE(fmt::format("hasSpill:{}", hasSpill)); + auto planStats = toPlanStats(task->taskStats()); + if (hasSpill) { + // Dynamic filtering should be disabled with spilling triggered. + ASSERT_EQ(0, getFiltersProduced(task, 1).sum); + ASSERT_EQ(0, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(getReplacedWithFilterRows(task, 1).sum, 0); + ASSERT_EQ(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_TRUE(planStats.at(probeScanId).dynamicFilterStats.empty()); + } else { + ASSERT_EQ(1, getFiltersProduced(task, 1).sum); + ASSERT_EQ(1, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(getReplacedWithFilterRows(task, 1).sum, 0); + ASSERT_LT(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_EQ( + planStats.at(probeScanId).dynamicFilterStats.producerNodeIds, + std::unordered_set({joinId})); + } + }) + .run(); + } + + // Right join. + op = PlanBuilder(planNodeIdGenerator, pool_.get()) + .tableScan(probeType) + .capturePlanNodeId(probeScanId) + .hashJoin( + {"c0"}, + {"u_c0"}, + buildSide, + "", + {"c0", "c1", "u_c1"}, + core::JoinType::kRight) + .capturePlanNodeId(joinId) + .project({"c0", "c1 + 1", "c1 + u_c1"}) + .planNode(); + { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .makeInputSplits(makeInputSplits(probeScanId)) + .referenceQuery( + "SELECT t.c0, t.c1 + 1, t.c1 + u.c1 FROM t RIGHT JOIN u ON t.c0 = u.c0") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + SCOPED_TRACE(fmt::format("hasSpill:{}", hasSpill)); + auto planStats = toPlanStats(task->taskStats()); + if (hasSpill) { + // Dynamic filtering should be disabled with spilling triggered. + ASSERT_EQ(0, getFiltersProduced(task, 1).sum); + ASSERT_EQ(0, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_TRUE(planStats.at(probeScanId).dynamicFilterStats.empty()); + } else { + ASSERT_EQ(1, getFiltersProduced(task, 1).sum); + ASSERT_EQ(1, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(0, getReplacedWithFilterRows(task, 1).sum); + ASSERT_LT(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_EQ( + planStats.at(probeScanId).dynamicFilterStats.producerNodeIds, + std::unordered_set({joinId})); + } + }) + .run(); + } + } + + // Basic push-down with column names projected out of the table scan + // having different names than column names in the files. + { + auto scanOutputType = ROW({"a", "b"}, {INTEGER(), BIGINT()}); + ColumnHandleMap assignments; + assignments["a"] = regularColumn("c0", INTEGER()); + assignments["b"] = regularColumn("c1", BIGINT()); + + core::PlanNodeId probeScanId; + core::PlanNodeId joinId; + auto op = PlanBuilder(planNodeIdGenerator, pool_.get()) + .startTableScan() + .outputType(scanOutputType) + .assignments(assignments) + .endTableScan() + .capturePlanNodeId(probeScanId) + .hashJoin({"a"}, {"u_c0"}, buildSide, "", {"a", "b", "u_c1"}) + .capturePlanNodeId(joinId) + .project({"a", "b + 1", "b + u_c1"}) + .planNode(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .makeInputSplits(makeInputSplits(probeScanId)) + .referenceQuery( + "SELECT t.c0, t.c1 + 1, t.c1 + u.c1 FROM t, u WHERE t.c0 = u.c0") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + SCOPED_TRACE(fmt::format("hasSpill:{}", hasSpill)); + auto planStats = toPlanStats(task->taskStats()); + if (hasSpill) { + // Dynamic filtering should be disabled with spilling triggered. + ASSERT_EQ(0, getFiltersProduced(task, 1).sum); + ASSERT_EQ(0, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(0, getReplacedWithFilterRows(task, 1).sum); + ASSERT_EQ(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_TRUE(planStats.at(probeScanId).dynamicFilterStats.empty()); + } else { + ASSERT_EQ(1, getFiltersProduced(task, 1).sum); + ASSERT_EQ(1, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(0, getReplacedWithFilterRows(task, 1).sum); + ASSERT_LT(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_EQ( + planStats.at(probeScanId).dynamicFilterStats.producerNodeIds, + std::unordered_set({joinId})); + } + }) + .run(); + } + + // Push-down that requires merging filters. + { + core::PlanNodeId probeScanId; + core::PlanNodeId joinId; + auto op = PlanBuilder(planNodeIdGenerator, pool_.get()) + .tableScan(probeType, {"c0 < 500::INTEGER"}) + .capturePlanNodeId(probeScanId) + .hashJoin({"c0"}, {"u_c0"}, buildSide, "", {"c1", "u_c1"}) + .capturePlanNodeId(joinId) + .project({"c1 + u_c1"}) + .planNode(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .makeInputSplits(makeInputSplits(probeScanId)) + .referenceQuery( + "SELECT t.c1 + u.c1 FROM t, u WHERE t.c0 = u.c0 AND t.c0 < 500") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + SCOPED_TRACE(fmt::format("hasSpill:{}", hasSpill)); + auto planStats = toPlanStats(task->taskStats()); + if (hasSpill) { + // Dynamic filtering should be disabled with spilling triggered. + ASSERT_EQ(0, getFiltersProduced(task, 1).sum); + ASSERT_EQ(0, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(0, getReplacedWithFilterRows(task, 1).sum); + ASSERT_EQ(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_TRUE(planStats.at(probeScanId).dynamicFilterStats.empty()); + } else { + ASSERT_EQ(1, getFiltersProduced(task, 1).sum); + ASSERT_EQ(1, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(0, getReplacedWithFilterRows(task, 1).sum); + ASSERT_LT(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_EQ( + planStats.at(probeScanId).dynamicFilterStats.producerNodeIds, + std::unordered_set({joinId})); + } + }) + .run(); + } + + // Push-down that turns join into a no-op. + { + core::PlanNodeId probeScanId; + core::PlanNodeId joinId; + auto op = + PlanBuilder(planNodeIdGenerator, pool_.get()) + .tableScan(probeType) + .capturePlanNodeId(probeScanId) + .hashJoin({"c0"}, {"u_c0"}, keyOnlyBuildSide, "", {"c0", "c1"}) + .capturePlanNodeId(joinId) + .project({"c0", "c1 + 1"}) + .planNode(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .makeInputSplits(makeInputSplits(probeScanId)) + .referenceQuery("SELECT t.c0, t.c1 + 1 FROM t, u WHERE t.c0 = u.c0") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + SCOPED_TRACE(fmt::format("hasSpill:{}", hasSpill)); + auto planStats = toPlanStats(task->taskStats()); + if (hasSpill) { + // Dynamic filtering should be disabled with spilling triggered. + ASSERT_EQ(0, getFiltersProduced(task, 1).sum); + ASSERT_EQ(0, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(0, getReplacedWithFilterRows(task, 1).sum); + ASSERT_EQ(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_TRUE(planStats.at(probeScanId).dynamicFilterStats.empty()); + } else { + ASSERT_EQ(1, getFiltersProduced(task, 1).sum); + ASSERT_EQ(1, getFiltersAccepted(task, 0).sum); + ASSERT_EQ( + getReplacedWithFilterRows(task, 1).sum, + numRowsBuild * numSplits); + ASSERT_LT(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_EQ( + planStats.at(probeScanId).dynamicFilterStats.producerNodeIds, + std::unordered_set({joinId})); + } + }) + .run(); + } + + // Push-down that turns join into a no-op with output having a different + // number of columns than the input. + { + core::PlanNodeId probeScanId; + core::PlanNodeId joinId; + auto op = PlanBuilder(planNodeIdGenerator, pool_.get()) + .tableScan(probeType) + .capturePlanNodeId(probeScanId) + .hashJoin({"c0"}, {"u_c0"}, keyOnlyBuildSide, "", {"c0"}) + .capturePlanNodeId(joinId) + .planNode(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .makeInputSplits(makeInputSplits(probeScanId)) + .referenceQuery("SELECT t.c0 FROM t JOIN u ON (t.c0 = u.c0)") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + SCOPED_TRACE(fmt::format("hasSpill:{}", hasSpill)); + auto planStats = toPlanStats(task->taskStats()); + if (hasSpill) { + // Dynamic filtering should be disabled with spilling triggered. + ASSERT_EQ(0, getFiltersProduced(task, 1).sum); + ASSERT_EQ(0, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(0, getReplacedWithFilterRows(task, 1).sum); + ASSERT_EQ(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_TRUE(planStats.at(probeScanId).dynamicFilterStats.empty()); + } else { + ASSERT_EQ(1, getFiltersProduced(task, 1).sum); + ASSERT_EQ(1, getFiltersAccepted(task, 0).sum); + ASSERT_EQ( + getReplacedWithFilterRows(task, 1).sum, + numRowsBuild * numSplits); + ASSERT_LT(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_EQ( + planStats.at(probeScanId).dynamicFilterStats.producerNodeIds, + std::unordered_set({joinId})); + } + }) + .run(); + } + + // Push-down that requires merging filters and turns join into a no-op. + { + core::PlanNodeId probeScanId; + core::PlanNodeId joinId; + auto op = PlanBuilder(planNodeIdGenerator, pool_.get()) + .tableScan(probeType, {"c0 < 500::INTEGER"}) + .capturePlanNodeId(probeScanId) + .hashJoin({"c0"}, {"u_c0"}, keyOnlyBuildSide, "", {"c1"}) + .capturePlanNodeId(joinId) + .project({"c1 + 1"}) + .planNode(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .makeInputSplits(makeInputSplits(probeScanId)) + .referenceQuery( + "SELECT t.c1 + 1 FROM t, u WHERE t.c0 = u.c0 AND t.c0 < 500") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + SCOPED_TRACE(fmt::format("hasSpill:{}", hasSpill)); + auto planStats = toPlanStats(task->taskStats()); + if (hasSpill) { + // Dynamic filtering should be disabled with spilling triggered. + ASSERT_EQ(0, getFiltersProduced(task, 1).sum); + ASSERT_EQ(0, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(getReplacedWithFilterRows(task, 1).sum, 0); + ASSERT_EQ(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_TRUE(planStats.at(probeScanId).dynamicFilterStats.empty()); + } else { + ASSERT_EQ(1, getFiltersProduced(task, 1).sum); + ASSERT_EQ(1, getFiltersAccepted(task, 0).sum); + ASSERT_GT(getReplacedWithFilterRows(task, 1).sum, 0); + ASSERT_LT(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_EQ( + planStats.at(probeScanId).dynamicFilterStats.producerNodeIds, + std::unordered_set({joinId})); + } + }) + .run(); + } + + // Push-down with highly selective filter in the scan. + { + // Inner join. + core::PlanNodeId probeScanId; + core::PlanNodeId joinId; + auto op = + PlanBuilder(planNodeIdGenerator, pool_.get()) + .tableScan(probeType, {"c0 < 200::INTEGER"}) + .capturePlanNodeId(probeScanId) + .hashJoin( + {"c0"}, {"u_c0"}, buildSide, "", {"c1"}, core::JoinType::kInner) + .capturePlanNodeId(joinId) + .project({"c1 + 1"}) + .planNode(); + + { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .makeInputSplits(makeInputSplits(probeScanId)) + .referenceQuery( + "SELECT t.c1 + 1 FROM t, u WHERE t.c0 = u.c0 AND t.c0 < 200") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + SCOPED_TRACE(fmt::format("hasSpill:{}", hasSpill)); + auto planStats = toPlanStats(task->taskStats()); + if (hasSpill) { + // Dynamic filtering should be disabled with spilling triggered. + ASSERT_EQ(0, getFiltersProduced(task, 1).sum); + ASSERT_EQ(0, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(getReplacedWithFilterRows(task, 1).sum, 0); + ASSERT_LT(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_TRUE(planStats.at(probeScanId).dynamicFilterStats.empty()); + } else { + ASSERT_EQ(1, getFiltersProduced(task, 1).sum); + ASSERT_EQ(1, getFiltersAccepted(task, 0).sum); + ASSERT_GT(getReplacedWithFilterRows(task, 1).sum, 0); + ASSERT_LT(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_EQ( + planStats.at(probeScanId).dynamicFilterStats.producerNodeIds, + std::unordered_set({joinId})); + } + }) + .run(); + } + + // Left semi join. + op = PlanBuilder(planNodeIdGenerator, pool_.get()) + .tableScan(probeType, {"c0 < 200::INTEGER"}) + .capturePlanNodeId(probeScanId) + .hashJoin( + {"c0"}, + {"u_c0"}, + buildSide, + "", + {"c1"}, + core::JoinType::kLeftSemiFilter) + .capturePlanNodeId(joinId) + .project({"c1 + 1"}) + .planNode(); + + { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .makeInputSplits(makeInputSplits(probeScanId)) + .referenceQuery( + "SELECT t.c1 + 1 FROM t WHERE t.c0 IN (SELECT c0 FROM u) AND t.c0 < 200") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + SCOPED_TRACE(fmt::format("hasSpill:{}", hasSpill)); + auto planStats = toPlanStats(task->taskStats()); + if (hasSpill) { + // Dynamic filtering should be disabled with spilling triggered. + ASSERT_EQ(0, getFiltersProduced(task, 1).sum); + ASSERT_EQ(0, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(getReplacedWithFilterRows(task, 1).sum, 0); + ASSERT_LT(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_TRUE(planStats.at(probeScanId).dynamicFilterStats.empty()); + } else { + ASSERT_EQ(1, getFiltersProduced(task, 1).sum); + ASSERT_EQ(1, getFiltersAccepted(task, 0).sum); + ASSERT_GT(getReplacedWithFilterRows(task, 1).sum, 0); + ASSERT_LT(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_EQ( + planStats.at(probeScanId).dynamicFilterStats.producerNodeIds, + std::unordered_set({joinId})); + } + }) + .run(); + } + + // Right semi join. + op = PlanBuilder(planNodeIdGenerator, pool_.get()) + .tableScan(probeType, {"c0 < 200::INTEGER"}) + .capturePlanNodeId(probeScanId) + .hashJoin( + {"c0"}, + {"u_c0"}, + buildSide, + "", + {"u_c1"}, + core::JoinType::kRightSemiFilter) + .capturePlanNodeId(joinId) + .project({"u_c1 + 1"}) + .planNode(); + + { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .makeInputSplits(makeInputSplits(probeScanId)) + .referenceQuery( + "SELECT u.c1 + 1 FROM u WHERE u.c0 IN (SELECT c0 FROM t) AND u.c0 < 200") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + SCOPED_TRACE(fmt::format("hasSpill:{}", hasSpill)); + auto planStats = toPlanStats(task->taskStats()); + if (hasSpill) { + // Dynamic filtering should be disabled with spilling triggered. + ASSERT_EQ(0, getFiltersProduced(task, 1).sum); + ASSERT_EQ(0, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(getReplacedWithFilterRows(task, 1).sum, 0); + ASSERT_LT(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_TRUE(planStats.at(probeScanId).dynamicFilterStats.empty()); + } else { + ASSERT_EQ(1, getFiltersProduced(task, 1).sum); + ASSERT_EQ(1, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(getReplacedWithFilterRows(task, 1).sum, 0); + ASSERT_LT(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_EQ( + planStats.at(probeScanId).dynamicFilterStats.producerNodeIds, + std::unordered_set({joinId})); + } + }) + .run(); + } + + // Right join. + op = + PlanBuilder(planNodeIdGenerator, pool_.get()) + .tableScan(probeType, {"c0 < 200::INTEGER"}) + .capturePlanNodeId(probeScanId) + .hashJoin( + {"c0"}, {"u_c0"}, buildSide, "", {"c1"}, core::JoinType::kRight) + .capturePlanNodeId(joinId) + .project({"c1 + 1"}) + .planNode(); + + { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .makeInputSplits(makeInputSplits(probeScanId)) + .referenceQuery( + "SELECT t.c1 + 1 FROM (SELECT * FROM t WHERE t.c0 < 200) t RIGHT JOIN u ON t.c0 = u.c0") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + SCOPED_TRACE(fmt::format("hasSpill:{}", hasSpill)); + auto planStats = toPlanStats(task->taskStats()); + if (hasSpill) { + // Dynamic filtering should be disabled with spilling triggered. + ASSERT_EQ(0, getFiltersProduced(task, 1).sum); + ASSERT_EQ(0, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(getReplacedWithFilterRows(task, 1).sum, 0); + ASSERT_LT(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_TRUE(planStats.at(probeScanId).dynamicFilterStats.empty()); + } else { + ASSERT_EQ(1, getFiltersProduced(task, 1).sum); + ASSERT_EQ(1, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(getReplacedWithFilterRows(task, 1).sum, 0); + ASSERT_LT(getInputPositions(task, 1), numRowsProbe * numSplits); + ASSERT_EQ( + planStats.at(probeScanId).dynamicFilterStats.producerNodeIds, + std::unordered_set({joinId})); + } + }) + .run(); + } + } + + // Disable filter push-down by using values in place of scan. + { + core::PlanNodeId joinId; + auto op = PlanBuilder(planNodeIdGenerator, pool_.get()) + .values(probeVectors) + .hashJoin({"c0"}, {"u_c0"}, buildSide, "", {"c1"}) + .capturePlanNodeId(joinId) + .project({"c1 + 1"}) + .planNode(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .referenceQuery("SELECT t.c1 + 1 FROM t, u WHERE t.c0 = u.c0") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + auto planStats = toPlanStats(task->taskStats()); + ASSERT_EQ(0, getFiltersProduced(task, 1).sum); + ASSERT_EQ(0, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(numRowsProbe * numSplits, getInputPositions(task, 1)); + }) + .run(); + } + + // Disable filter push-down by using an expression as the join key on the + // probe side. + { + core::PlanNodeId probeScanId; + core::PlanNodeId joinId; + auto op = PlanBuilder(planNodeIdGenerator, pool_.get()) + .tableScan(probeType) + .capturePlanNodeId(probeScanId) + .project({"cast(c0 + 1 as integer) AS t_key", "c1"}) + .hashJoin({"t_key"}, {"u_c0"}, buildSide, "", {"c1"}) + .capturePlanNodeId(joinId) + .project({"c1 + 1"}) + .planNode(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .makeInputSplits(makeInputSplits(probeScanId)) + .referenceQuery("SELECT t.c1 + 1 FROM t, u WHERE (t.c0 + 1) = u.c0") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + auto planStats = toPlanStats(task->taskStats()); + ASSERT_EQ(0, getFiltersProduced(task, 1).sum); + ASSERT_EQ(0, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(numRowsProbe * numSplits, getInputPositions(task, 1)); + ASSERT_TRUE(planStats.at(probeScanId).dynamicFilterStats.empty()); + }) + .run(); + } +} + +TEST_F(HashJoinTest, DISABLED_dynamicFiltersStatsWithChainedJoins) { + const int32_t numSplits = 10; + const int32_t numProbeRows = 333; + const int32_t numBuildRows = 100; + + std::vector probeVectors; + probeVectors.reserve(numSplits); + std::vector> tempFiles; + for (int32_t i = 0; i < numSplits; ++i) { + auto rowVector = makeRowVector({ + makeFlatVector( + numProbeRows, [&](auto row) { return row - i * 10; }), + makeFlatVector(numProbeRows, [](auto row) { return row; }), + }); + probeVectors.push_back(rowVector); + tempFiles.push_back(TempFilePath::create()); + writeToFile(tempFiles.back()->getPath(), rowVector); + } + auto makeInputSplits = [&](const core::PlanNodeId& nodeId) { + return [&] { + std::vector probeSplits; + for (auto& file : tempFiles) { + probeSplits.push_back( + exec::Split(makeHiveConnectorSplit(file->getPath()))); + } + SplitInput splits; + splits.emplace(nodeId, probeSplits); + return splits; + }; + }; + + // 100 key values in [35, 233] range. + std::vector buildVectors; + for (int i = 0; i < 5; ++i) { + buildVectors.push_back(makeRowVector({ + makeFlatVector( + numBuildRows / 5, + [i](auto row) { return 35 + 2 * (row + i * numBuildRows / 5); }), + makeFlatVector(numBuildRows / 5, [](auto row) { return row; }), + })); + } + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + auto probeType = ROW({"c0", "c1"}, {INTEGER(), BIGINT()}); + + auto planNodeIdGenerator = std::make_shared(); + + auto buildSide1 = PlanBuilder(planNodeIdGenerator, pool_.get()) + .values(buildVectors) + .project({"c0 AS u_c0", "c1 AS u_c1"}) + .planNode(); + auto buildSide2 = PlanBuilder(planNodeIdGenerator, pool_.get()) + .values(buildVectors) + .project({"c0 AS u_c0", "c1 AS u_c1"}) + .planNode(); + // Inner join pushdown. + core::PlanNodeId probeScanId; + core::PlanNodeId joinId1; + core::PlanNodeId joinId2; + auto op = PlanBuilder(planNodeIdGenerator, pool_.get()) + .tableScan(probeType) + .capturePlanNodeId(probeScanId) + .hashJoin( + {"c0"}, + {"u_c0"}, + buildSide1, + "", + {"c0", "c1"}, + core::JoinType::kInner) + .capturePlanNodeId(joinId1) + .hashJoin( + {"c0"}, + {"u_c0"}, + buildSide2, + "", + {"c0", "c1", "u_c1"}, + core::JoinType::kInner) + .capturePlanNodeId(joinId2) + .project({"c0", "c1 + 1", "c1 + u_c1"}) + .planNode(); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .makeInputSplits(makeInputSplits(probeScanId)) + .injectSpill(false) + .referenceQuery( + "SELECT t.c0, t.c1 + 1, t.c1 + u.c1 FROM t, u WHERE t.c0 = u.c0") + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + auto planStats = toPlanStats(task->taskStats()); + ASSERT_EQ( + planStats.at(probeScanId).dynamicFilterStats.producerNodeIds, + std::unordered_set({joinId1, joinId2})); + }) + .run(); +} + +TEST_F(HashJoinTest, DISABLED_dynamicFiltersWithSkippedSplits) { + const int32_t numSplits = 20; + const int32_t numNonSkippedSplits = 10; + const int32_t numRowsProbe = 333; + const int32_t numRowsBuild = 100; + + std::vector probeVectors; + probeVectors.reserve(numSplits); + + std::vector> tempFiles; + // Each split has a column containing + // the split number. This is used to filter out whole splits based + // on metadata. We test how using metadata for dropping splits + // interactts with dynamic filters. In specific, if the first split + // is discarded based on metadata, the dynamic filters must not be + // lost even if there is no actual reader for the split. + for (int32_t i = 0; i < numSplits; ++i) { + auto rowVector = makeRowVector({ + makeFlatVector( + numRowsProbe, [&](auto row) { return row - i * 10; }), + makeFlatVector(numRowsProbe, [](auto row) { return row; }), + makeFlatVector( + numRowsProbe, [&](auto /*row*/) { return i % 2 == 0 ? 0 : i; }), + }); + probeVectors.push_back(rowVector); + tempFiles.push_back(TempFilePath::create()); + writeToFile(tempFiles.back()->getPath(), rowVector); + } + + auto makeInputSplits = [&](const core::PlanNodeId& nodeId) { + return [&] { + std::vector probeSplits; + for (auto& file : tempFiles) { + probeSplits.push_back( + exec::Split(makeHiveConnectorSplit(file->getPath()))); + } + // We add splits that have no rows. + auto makeEmpty = [&]() { + return exec::Split( + HiveConnectorSplitBuilder(tempFiles.back()->getPath()) + .start(10000000) + .length(1) + .build()); + }; + std::vector emptyFront = {makeEmpty(), makeEmpty()}; + std::vector emptyMiddle = {makeEmpty(), makeEmpty()}; + probeSplits.insert( + probeSplits.begin(), emptyFront.begin(), emptyFront.end()); + probeSplits.insert( + probeSplits.begin() + 13, emptyMiddle.begin(), emptyMiddle.end()); + SplitInput splits; + splits.emplace(nodeId, probeSplits); + return splits; + }; + }; + + // 100 key values in [35, 233] range. + std::vector buildVectors; + for (int i = 0; i < 5; ++i) { + buildVectors.push_back(makeRowVector({ + makeFlatVector( + numRowsBuild / 5, + [i](auto row) { return 35 + 2 * (row + i * numRowsBuild / 5); }), + makeFlatVector(numRowsBuild / 5, [](auto row) { return row; }), + })); + } + std::vector keyOnlyBuildVectors; + for (int i = 0; i < 5; ++i) { + keyOnlyBuildVectors.push_back( + makeRowVector({makeFlatVector(numRowsBuild / 5, [i](auto row) { + return 35 + 2 * (row + i * numRowsBuild / 5); + })})); + } + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + auto probeType = ROW({"c0", "c1", "c2"}, {INTEGER(), BIGINT(), BIGINT()}); + + auto planNodeIdGenerator = std::make_shared(); + + auto buildSide = PlanBuilder(planNodeIdGenerator, pool_.get()) + .values(buildVectors) + .project({"c0 AS u_c0", "c1 AS u_c1"}) + .planNode(); + auto keyOnlyBuildSide = PlanBuilder(planNodeIdGenerator, pool_.get()) + .values(keyOnlyBuildVectors) + .project({"c0 AS u_c0"}) + .planNode(); + + // Basic push-down. + { + // Inner join. + core::PlanNodeId probeScanId; + auto op = PlanBuilder(planNodeIdGenerator, pool_.get()) + .tableScan(probeType, {"c2 > 0"}) + .capturePlanNodeId(probeScanId) + .hashJoin( + {"c0"}, + {"u_c0"}, + buildSide, + "", + {"c0", "c1", "u_c1"}, + core::JoinType::kInner) + .project({"c0", "c1 + 1", "c1 + u_c1"}) + .planNode(); + { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .numDrivers(1) + .makeInputSplits(makeInputSplits(probeScanId)) + .referenceQuery( + "SELECT t.c0, t.c1 + 1, t.c1 + u.c1 FROM t, u WHERE t.c0 = u.c0 AND t.c2 > 0") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + SCOPED_TRACE(fmt::format("hasSpill:{}", hasSpill)); + if (hasSpill) { + // Dynamic filtering should be disabled with spilling triggered. + ASSERT_EQ(0, getFiltersProduced(task, 1).sum); + ASSERT_EQ(0, getFiltersAccepted(task, 0).sum); + ASSERT_EQ( + getInputPositions(task, 1), + numRowsProbe * numNonSkippedSplits); + } else { + ASSERT_EQ(1, getFiltersProduced(task, 1).sum); + ASSERT_EQ(1, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(0, getReplacedWithFilterRows(task, 1).sum); + ASSERT_LT( + getInputPositions(task, 1), + numRowsProbe * numNonSkippedSplits); + } + }) + .run(); + } + + // Left semi join. + op = PlanBuilder(planNodeIdGenerator, pool_.get()) + .tableScan(probeType, {"c2 > 0"}) + .capturePlanNodeId(probeScanId) + .hashJoin( + {"c0"}, + {"u_c0"}, + buildSide, + "", + {"c0", "c1"}, + core::JoinType::kLeftSemiFilter) + .project({"c0", "c1 + 1"}) + .planNode(); + + { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .numDrivers(1) + .makeInputSplits(makeInputSplits(probeScanId)) + .referenceQuery( + "SELECT t.c0, t.c1 + 1 FROM t WHERE t.c0 IN (SELECT c0 FROM u) AND t.c2 > 0") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + SCOPED_TRACE(fmt::format("hasSpill:{}", hasSpill)); + if (hasSpill) { + // Dynamic filtering should be disabled with spilling triggered. + ASSERT_EQ(0, getFiltersProduced(task, 1).sum); + ASSERT_EQ(0, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(0, getReplacedWithFilterRows(task, 1).sum); + ASSERT_EQ( + getInputPositions(task, 1), + numRowsProbe * numNonSkippedSplits); + } else { + ASSERT_EQ(1, getFiltersProduced(task, 1).sum); + ASSERT_EQ(1, getFiltersAccepted(task, 0).sum); + ASSERT_GT(getReplacedWithFilterRows(task, 1).sum, 0); + ASSERT_LT( + getInputPositions(task, 1), + numRowsProbe * numNonSkippedSplits); + } + }) + .run(); + } + + // Right semi join. + op = PlanBuilder(planNodeIdGenerator, pool_.get()) + .tableScan(probeType, {"c2 > 0"}) + .capturePlanNodeId(probeScanId) + .hashJoin( + {"c0"}, + {"u_c0"}, + buildSide, + "", + {"u_c0", "u_c1"}, + core::JoinType::kRightSemiFilter) + .project({"u_c0", "u_c1 + 1"}) + .planNode(); + + { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .numDrivers(1) + .makeInputSplits(makeInputSplits(probeScanId)) + .referenceQuery( + "SELECT u.c0, u.c1 + 1 FROM u WHERE u.c0 IN (SELECT c0 FROM t WHERE t.c2 > 0)") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + SCOPED_TRACE(fmt::format("hasSpill:{}", hasSpill)); + if (hasSpill) { + // Dynamic filtering should be disabled with spilling triggered. + ASSERT_EQ(0, getFiltersProduced(task, 1).sum); + ASSERT_EQ(0, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(getReplacedWithFilterRows(task, 1).sum, 0); + ASSERT_EQ( + getInputPositions(task, 1), + numRowsProbe * numNonSkippedSplits); + } else { + ASSERT_EQ(1, getFiltersProduced(task, 1).sum); + ASSERT_EQ(1, getFiltersAccepted(task, 0).sum); + ASSERT_EQ(getReplacedWithFilterRows(task, 1).sum, 0); + ASSERT_LT( + getInputPositions(task, 1), + numRowsProbe * numNonSkippedSplits); + } + }) + .run(); + } + } +} + +TEST_F(HashJoinTest, DISABLED_dynamicFiltersAppliedToPreloadedSplits) { + vector_size_t size = 1000; + const int32_t numSplits = 5; + + std::vector probeVectors; + probeVectors.reserve(numSplits); + + // Prepare probe side table. + std::vector> tempFiles; + std::vector probeSplits; + for (int32_t i = 0; i < numSplits; ++i) { + auto rowVector = makeRowVector( + {"p0", "p1"}, + { + makeFlatVector( + size, [&](auto row) { return (row + 1) * (i + 1); }), + makeFlatVector(size, [&](auto /*row*/) { return i; }), + }); + probeVectors.push_back(rowVector); + tempFiles.push_back(TempFilePath::create()); + writeToFile(tempFiles.back()->getPath(), rowVector); + auto split = HiveConnectorSplitBuilder(tempFiles.back()->getPath()) + .partitionKey("p1", std::to_string(i)) + .build(); + probeSplits.push_back(exec::Split(split)); + } + + auto outputType = ROW({"p0", "p1"}, {BIGINT(), BIGINT()}); + ColumnHandleMap assignments = { + {"p0", regularColumn("p0", BIGINT())}, + {"p1", partitionKey("p1", BIGINT())}}; + createDuckDbTable("p", probeVectors); + + // Prepare build side table. + std::vector buildVectors{ + makeRowVector({"b0"}, {makeFlatVector({0, numSplits})})}; + createDuckDbTable("b", buildVectors); + + // Executing the join with p1=b0, we expect a dynamic filter for p1 to prune + // the entire file/split. There are total of five splits, and all except the + // first one are expected to be pruned. The result 'preloadedSplits' > 1 + // confirms the successful push of dynamic filters to the preloading data + // source. + core::PlanNodeId probeScanId; + core::PlanNodeId joinNodeId; + auto planNodeIdGenerator = std::make_shared(); + auto op = + PlanBuilder(planNodeIdGenerator) + .startTableScan() + .outputType(outputType) + .assignments(assignments) + .endTableScan() + .capturePlanNodeId(probeScanId) + .hashJoin( + {"p1"}, + {"b0"}, + PlanBuilder(planNodeIdGenerator).values(buildVectors).planNode(), + "", + {"p0"}, + core::JoinType::kInner) + .capturePlanNodeId(joinNodeId) + .project({"p0"}) + .planNode(); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .config(core::QueryConfig::kMaxSplitPreloadPerDriver, "3") + .injectSpill(false) + .inputSplits({{probeScanId, probeSplits}}) + .referenceQuery("select p.p0 from p, b where b.b0 = p.p1") + .checkSpillStats(false) + .verifier([&](const std::shared_ptr& task, bool /*hasSpill*/) { + auto planStats = toPlanStats(task->taskStats()); + auto getStatSum = [&](const core::PlanNodeId& id, + const std::string& name) { + return planStats.at(id).customStats.at(name).sum; + }; + ASSERT_EQ(1, getStatSum(joinNodeId, "dynamicFiltersProduced")); + ASSERT_EQ(1, getStatSum(probeScanId, "dynamicFiltersAccepted")); + ASSERT_EQ(4, getStatSum(probeScanId, "skippedSplits")); + ASSERT_LT(1, getStatSum(probeScanId, "preloadedSplits")); + }) + .run(); +} + +TEST_F(HashJoinTest, DISABLED_dynamicFiltersPushDownThroughAgg) { + const int32_t numRowsProbe = 300; + const int32_t numRowsBuild = 100; + + // Create probe data + std::vector probeVectors{makeRowVector({ + makeFlatVector(numRowsProbe, [&](auto row) { return row - 10; }), + makeFlatVector(numRowsProbe, folly::identity), + })}; + std::shared_ptr probeFile = TempFilePath::create(); + writeToFile(probeFile->getPath(), probeVectors); + + // Create build data + std::vector buildVectors{makeRowVector( + {"u0"}, {makeFlatVector(numRowsBuild, [&](auto row) { + return 35 + 2 * (row + numRowsBuild / 5); + })})}; + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + auto probeType = ROW({"c0", "c1"}, {INTEGER(), BIGINT()}); + auto planNodeIdGenerator = std::make_shared(); + auto buildSide = + PlanBuilder(planNodeIdGenerator).values(buildVectors).planNode(); + + // Inner join. + core::PlanNodeId scanNodeId; + core::PlanNodeId joinNodeId; + core::PlanNodeId aggNodeId; + auto op = PlanBuilder(planNodeIdGenerator, pool_.get()) + .tableScan(probeType) + .capturePlanNodeId(scanNodeId) + .partialAggregation({"c0"}, {"sum(c1)"}) + .capturePlanNodeId(aggNodeId) + .hashJoin( + {"c0"}, + {"u0"}, + buildSide, + "", + {"c0", "a0"}, + core::JoinType::kInner) + .capturePlanNodeId(joinNodeId) + .planNode(); + + SplitInput splitInput = { + {scanNodeId, {Split(makeHiveConnectorSplit(probeFile->getPath()))}}}; + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .inputSplits(splitInput) + .injectSpill(false) + .checkSpillStats(false) + .referenceQuery("SELECT c0, sum(c1) FROM t, u WHERE c0 = u0 group by c0") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + auto planStats = toPlanStats(task->taskStats()); + auto dynamicFilterStats = planStats.at(scanNodeId).dynamicFilterStats; + ASSERT_EQ( + 1, getFiltersProduced(task, getOperatorIndex(joinNodeId)).sum); + ASSERT_EQ( + 1, getFiltersAccepted(task, getOperatorIndex(scanNodeId)).sum); + ASSERT_LT( + getInputPositions(task, getOperatorIndex(aggNodeId)), numRowsProbe); + ASSERT_EQ( + dynamicFilterStats.producerNodeIds, + std::unordered_set({joinNodeId})); + }) + .run(); +} + +TEST_F(HashJoinTest, noDynamicFiltersPushDownThroughRightJoin) { + std::vector innerBuild = {makeRowVector( + {"a"}, + { + makeFlatVector(5, [](auto i) { return 2 * i; }), + })}; + std::vector rightBuild = {makeRowVector( + {"b"}, + { + makeFlatVector(5, [](auto i) { return 1 + 2 * i; }), + })}; + std::vector rightProbe = {makeRowVector( + {"aa", "bb"}, + { + makeFlatVector(10, folly::identity), + makeFlatVector(10, folly::identity), + })}; + auto file = TempFilePath::create(); + writeToFile(file->getPath(), rightProbe); + auto planNodeIdGenerator = std::make_shared(); + core::PlanNodeId scanNodeId; + auto plan = + PlanBuilder(planNodeIdGenerator) + .tableScan(asRowType(rightProbe[0]->type())) + .capturePlanNodeId(scanNodeId) + .hashJoin( + {"bb"}, + {"b"}, + PlanBuilder(planNodeIdGenerator).values(rightBuild).planNode(), + "", + {"aa", "b"}, + core::JoinType::kRight) + .hashJoin( + {"aa"}, + {"a"}, + PlanBuilder(planNodeIdGenerator).values(innerBuild).planNode(), + "", + {"aa"}) + .planNode(); + AssertQueryBuilder(plan) + .split(scanNodeId, Split(makeHiveConnectorSplit(file->getPath()))) + .assertResults( + BaseVector::create(innerBuild[0]->type(), 0, pool_.get())); +} + +// Verify the size of the join output vectors when projecting build-side +// variable-width column. +TEST_F(HashJoinTest, memoryUsage) { + std::vector probeVectors = + makeBatches(10, [&](int32_t /*unused*/) { + return makeRowVector( + {makeFlatVector(1'000, [](auto row) { return row % 5; })}); + }); + std::vector buildVectors = + makeBatches(5, [&](int32_t /*unused*/) { + return makeRowVector( + {"u_c0", "u_c1"}, + {makeFlatVector({0, 1, 2}), + makeFlatVector({ + std::string(40, 'a'), + std::string(50, 'b'), + std::string(30, 'c'), + })}); + }); + core::PlanNodeId joinNodeId; + + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors) + .hashJoin( + {"c0"}, + {"u_c0"}, + PlanBuilder(planNodeIdGenerator) + .values({buildVectors}) + .planNode(), + "", + {"c0", "u_c1"}) + .capturePlanNodeId(joinNodeId) + .singleAggregation({}, {"count(1)"}) + .planNode(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .injectSpill(false) + .planNode(std::move(plan)) + .referenceQuery("SELECT 30000") + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + if (hasSpill) { + return; + } + auto planStats = toPlanStats(task->taskStats()); + auto outputBytes = planStats.at(joinNodeId).outputBytes; + ASSERT_LT(outputBytes, ((40 + 50 + 30) / 3 + 8) * 1000 * 10 * 5); + // Verify number of memory allocations. Should not be too high if + // hash join is able to re-use output vectors that contain + // build-side data. + ASSERT_GT(40, task->pool()->stats().numAllocs); + }) + .run(); +} + +/// Test an edge case in producing small output batches where the logic to +/// calculate the set of probe-side rows to load lazy vectors for was +/// triggering a crash. +TEST_F(HashJoinTest, smallOutputBatchSize) { + // Setup probe data with 50 non-null matching keys followed by 50 null + // keys: 1, 2, 1, 2,...null, null. + auto probeVectors = makeRowVector({ + makeFlatVector( + 100, + [](auto row) { return 1 + row % 2; }, + [](auto row) { return row > 50; }), + makeFlatVector(100, [](auto row) { return row * 10; }), + }); + + // Setup build side to match non-null probe side keys. + auto buildVectors = makeRowVector( + {"u_c0", "u_c1"}, + { + makeFlatVector({1, 2}), + makeFlatVector({100, 200}), + }); + + createDuckDbTable("t", {probeVectors}); + createDuckDbTable("u", {buildVectors}); + + // Plan hash inner join with a filter. + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values({probeVectors}) + .hashJoin( + {"c0"}, + {"u_c0"}, + PlanBuilder(planNodeIdGenerator) + .values({buildVectors}) + .planNode(), + "c1 < u_c1", + {"c0", "u_c1"}) + .planNode(); + + // Use small output batch size to trigger logic for calculating set of + // probe-side rows to load lazy vectors for. + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(plan)) + .config(core::QueryConfig::kPreferredOutputBatchRows, std::to_string(10)) + .referenceQuery("SELECT c0, u_c1 FROM t, u WHERE c0 = u_c0 AND c1 < u_c1") + .injectSpill(false) + .run(); +} + +TEST_F(HashJoinTest, DISABLED_spillFileSize) { + const std::vector maxSpillFileSizes({0, 1, 1'000'000'000}); + for (const auto spillFileSize : maxSpillFileSizes) { + SCOPED_TRACE(fmt::format("spillFileSize: {}", spillFileSize)); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .keyTypes({BIGINT()}) + .probeVectors(100, 3) + .buildVectors(100, 3) + .referenceQuery( + "SELECT t_k0, t_data, u_k0, u_data FROM t, u WHERE t.t_k0 = u.u_k0") + .config(core::QueryConfig::kSpillStartPartitionBit, "48") + .config(core::QueryConfig::kSpillNumPartitionBits, "3") + .config( + core::QueryConfig::kMaxSpillFileSize, std::to_string(spillFileSize)) + .checkSpillStats(false) + .maxSpillLevel(0) + .verifier([&](const std::shared_ptr& task, bool hasSpill) { + if (!hasSpill) { + return; + } + const auto statsPair = taskSpilledStats(*task); + const int32_t numPartitions = statsPair.first.spilledPartitions; + ASSERT_EQ(statsPair.second.spilledPartitions, numPartitions); + const auto fileSizes = numTaskSpillFiles(*task); + if (spillFileSize != 1) { + ASSERT_EQ(fileSizes.first, numPartitions); + } else { + ASSERT_GT(fileSizes.first, numPartitions); + } + verifyTaskSpilledRuntimeStats(*task, true); + }) + .run(); + } +} + +// Spilling is not supported for cudfHashJoin. +TEST_F(HashJoinTest, DISABLED_spillPartitionBitsOverlap) { + auto builder = + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .keyTypes({BIGINT(), BIGINT()}) + .probeVectors(2'000, 3) + .buildVectors(2'000, 3) + .referenceQuery( + "SELECT t_k0, t_k1, t_data, u_k0, u_k1, u_data FROM t, u WHERE t_k0 = u_k0 and t_k1 = u_k1") + .config(core::QueryConfig::kSpillStartPartitionBit, "8") + .config(core::QueryConfig::kSpillNumPartitionBits, "1") + .checkSpillStats(false) + .maxSpillLevel(0); + VELOX_ASSERT_THROW(builder.run(), "vs. 8"); +} + +// The test is to verify if the hash build reservation has been released on +// task error. +DEBUG_ONLY_TEST_F(HashJoinTest, buildReservationReleaseCheck) { + std::vector probeVectors = + makeBatches(1, [&](int32_t /*unused*/) { + return std::dynamic_pointer_cast( + BatchMaker::createBatch(probeType_, 1000, *pool_)); + }); + std::vector buildVectors = makeBatches(10, [&](int32_t index) { + return std::dynamic_pointer_cast( + BatchMaker::createBatch(buildType_, 5000 * (1 + index), *pool_)); + }); + + auto planNodeIdGenerator = std::make_shared(); + CursorParameters params; + params.planNode = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, true) + .hashJoin( + {"t_k1"}, + {"u_k1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, true) + .planNode(), + "", + concat(probeType_->names(), buildType_->names())) + .planNode(); + params.queryCtx = core::QueryCtx::create(driverExecutor_.get()); + // NOTE: the spilling setup is to trigger memory reservation code path which + // only gets executed when spilling is enabled. We don't care about if + // spilling is really triggered in test or not. + auto spillDirectory = exec::test::TempDirectoryPath::create(); + params.spillDirectory = spillDirectory->getPath(); + params.queryCtx->testingOverrideConfigUnsafe( + {{core::QueryConfig::kSpillEnabled, "true"}, + {core::QueryConfig::kMaxSpillLevel, "0"}}); + params.maxDrivers = 1; + + auto cursor = TaskCursor::create(params); + auto* task = cursor->task().get(); + + // Set up a testvalue to trigger task abort when hash build tries to reserve + // memory. + SCOPED_TESTVALUE_SET( + "facebook::velox::common::memory::MemoryPoolImpl::maybeReserve", + std::function( + [&](memory::MemoryPool* /*unused*/) { task->requestAbort(); })); + auto runTask = [&]() { + while (cursor->moveNext()) { + } + }; + VELOX_ASSERT_THROW(runTask(), ""); + ASSERT_TRUE(waitForTaskAborted(task, 5'000'000)); +} + +TEST_F(HashJoinTest, dynamicFilterOnPartitionKey) { + vector_size_t size = 10; + auto filePaths = makeFilePaths(1); + auto rowVector = makeRowVector( + {makeFlatVector(size, [&](auto row) { return row; })}); + createDuckDbTable("u", {rowVector}); + writeToFile(filePaths[0]->getPath(), rowVector); + std::vector buildVectors{ + makeRowVector({"c0"}, {makeFlatVector({0, 1, 2})})}; + createDuckDbTable("t", buildVectors); + auto split = facebook::velox::exec::test::HiveConnectorSplitBuilder( + filePaths[0]->getPath()) + .partitionKey("k", "0") + .build(); + auto outputType = ROW({"n1_0", "n1_1"}, {BIGINT(), BIGINT()}); + ColumnHandleMap assignments = { + {"n1_0", regularColumn("c0", BIGINT())}, + {"n1_1", partitionKey("k", BIGINT())}}; + + core::PlanNodeId probeScanId; + auto planNodeIdGenerator = std::make_shared(); + auto op = + PlanBuilder(planNodeIdGenerator) + .startTableScan() + .outputType(outputType) + .assignments(assignments) + .endTableScan() + .capturePlanNodeId(probeScanId) + .hashJoin( + {"n1_1"}, + {"c0"}, + PlanBuilder(planNodeIdGenerator).values(buildVectors).planNode(), + "", + {"c0"}, + core::JoinType::kInner) + .project({"c0"}) + .planNode(); + SplitInput splits = {{probeScanId, {exec::Split(split)}}}; + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .inputSplits(splits) + .referenceQuery("select t.c0 from t, u where t.c0 = 0") + .checkSpillStats(false) + .run(); +} + +TEST_F(HashJoinTest, DISABLED_probeMemoryLimitOnBuildProjection) { + const uint64_t numBuildRows = 20; + std::vector probeVectors = + makeBatches(10, [&](int32_t /*unused*/) { + return makeRowVector({makeFlatVector( + 1'000, [](auto row) { return row % 25; })}); + }); + + std::vector buildVectors = + makeBatches(1, [&](int32_t /*unused*/) { + return makeRowVector( + {"u_c0", "u_c1", "u_c2", "u_c3", "u_c4"}, + {makeFlatVector( + numBuildRows, [](auto row) { return row; }), + makeFlatVector( + numBuildRows, + [](auto /* row */) { return std::string(4096, 'a'); }), + makeFlatVector( + numBuildRows, + [](auto /* row */) { return std::string(4096, 'a'); }), + makeFlatVector( + numBuildRows, + [](auto row) { + // Row that has too large of size variation. + if (row == 0) { + return std::string(4096, 'a'); + } else { + return std::string(1, 'a'); + } + }), + makeFlatVector(numBuildRows, [](auto row) { + // Row that has tolerable size variation. + if (row == 0) { + return std::string(4096, 'a'); + } else { + return std::string(256, 'a'); + } + })}); + }); + + createDuckDbTable("t", {probeVectors}); + createDuckDbTable("u", {buildVectors}); + + struct TestParam { + std::vector varSizeColumns; + int32_t numExpectedBatches; + std::string referenceQuery; + std::string debugString() const { + std::stringstream ss; + ss << "varSizeColumns ["; + for (const auto& columnIndex : varSizeColumns) { + ss << columnIndex << ", "; + } + ss << "] "; + ss << "numExpectedBatches " << numExpectedBatches << ", referenceQuery '" + << referenceQuery << "'"; + return ss.str(); + } + }; + + std::vector testParams{ + {{}, 10, "SELECT t.c0 FROM t JOIN u ON t.c0 = u.u_c0"}, + {{1}, 4000, "SELECT t.c0, u.u_c1 FROM t JOIN u ON t.c0 = u.u_c0"}, + {{1, 2}, + 8000, + "SELECT t.c0, u.u_c1, u.u_c2 FROM t JOIN u ON t.c0 = u.u_c0"}, + {{3}, 210, "SELECT t.c0, u.u_c3 FROM t JOIN u ON t.c0 = u.u_c0"}, + {{4}, 2670, "SELECT t.c0, u.u_c4 FROM t JOIN u ON t.c0 = u.u_c0"}}; + + for (const auto& testParam : testParams) { + SCOPED_TRACE(testParam.debugString()); + core::PlanNodeId joinNodeId; + std::vector outputLayout; + outputLayout.push_back("c0"); + for (int32_t i = 0; i < testParam.varSizeColumns.size(); i++) { + outputLayout.push_back(fmt::format("u_c{}", testParam.varSizeColumns[i])); + } + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors) + .hashJoin( + {"c0"}, + {"u_c0"}, + PlanBuilder(planNodeIdGenerator) + .values({buildVectors}) + .planNode(), + "", + outputLayout) + .capturePlanNodeId(joinNodeId) + .planNode(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(plan)) + .config(core::QueryConfig::kPreferredOutputBatchBytes, "8192") + .injectSpill(false) + .referenceQuery(testParam.referenceQuery) + .verifier([&](const std::shared_ptr& task, bool /* unused */) { + auto planStats = toPlanStats(task->taskStats()); + auto outputBatches = planStats.at(joinNodeId).outputVectors; + ASSERT_EQ(outputBatches, testParam.numExpectedBatches); + }) + .run(); + } +} + +DEBUG_ONLY_TEST_F(HashJoinTest, reclaimDuringInputProcessing) { + constexpr int64_t kMaxBytes = 1LL << 30; // 1GB + VectorFuzzer fuzzer({.vectorSize = 1000}, pool()); + const int32_t numBuildVectors = 10; + std::vector buildVectors; + for (int32_t i = 0; i < numBuildVectors; ++i) { + buildVectors.push_back(fuzzer.fuzzRow(buildType_)); + } + const int32_t numProbeVectors = 5; + std::vector probeVectors; + for (int32_t i = 0; i < numProbeVectors; ++i) { + probeVectors.push_back(fuzzer.fuzzRow(probeType_)); + } + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + struct { + // 0: trigger reclaim with some input processed. + // 1: trigger reclaim after all the inputs processed. + int triggerCondition; + bool spillEnabled; + bool expectedReclaimable; + + std::string debugString() const { + return fmt::format( + "triggerCondition {}, spillEnabled {}, expectedReclaimable {}", + triggerCondition, + spillEnabled, + expectedReclaimable); + } + } testSettings[] = { + {0, true, true}, {0, true, true}, {0, false, false}, {0, false, false}}; + for (const auto& testData : testSettings) { + SCOPED_TRACE(testData.debugString()); + + auto tempDirectory = exec::test::TempDirectoryPath::create(); + auto queryPool = memory::memoryManager()->addRootPool( + "", kMaxBytes, memory::MemoryReclaimer::create()); + + core::PlanNodeId probeScanId; + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, false) + .hashJoin( + {"t_k1"}, + {"u_k1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, false) + .planNode(), + "", + concat(probeType_->names(), buildType_->names())) + .planNode(); + + folly::EventCount driverWait; + auto driverWaitKey = driverWait.prepareWait(); + folly::EventCount testWait; + auto testWaitKey = testWait.prepareWait(); + + std::atomic numInputs{0}; + Operator* op; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::addInput", + std::function(([&](Operator* testOp) { + if (testOp->operatorType() != "HashBuild") { + return; + } + op = testOp; + ++numInputs; + if (testData.triggerCondition == 0) { + if (numInputs != 2) { + return; + } + } + if (testData.triggerCondition == 1) { + if (numInputs != numBuildVectors) { + return; + } + } + ASSERT_EQ(op->canReclaim(), testData.expectedReclaimable); + uint64_t reclaimableBytes{0}; + const bool reclaimable = op->reclaimableBytes(reclaimableBytes); + ASSERT_EQ(reclaimable, testData.expectedReclaimable); + if (testData.expectedReclaimable) { + ASSERT_GT(reclaimableBytes, 0); + } else { + ASSERT_EQ(reclaimableBytes, 0); + } + testWait.notify(); + driverWait.wait(driverWaitKey); + }))); + + std::thread taskThread([&]() { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .planNode(plan) + .queryPool(std::move(queryPool)) + .injectSpill(false) + .spillDirectory(testData.spillEnabled ? tempDirectory->getPath() : "") + .referenceQuery( + "SELECT t_k1, t_k2, t_v1, u_k1, u_k2, u_v1 FROM t, u WHERE t.t_k1 = u.u_k1") + .config(core::QueryConfig::kSpillStartPartitionBit, "29") + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + const auto statsPair = taskSpilledStats(*task); + if (testData.expectedReclaimable) { + ASSERT_GT(statsPair.first.spilledBytes, 0); + ASSERT_EQ(statsPair.first.spilledPartitions, 8); + ASSERT_GT(statsPair.second.spilledBytes, 0); + ASSERT_EQ(statsPair.second.spilledPartitions, 8); + verifyTaskSpilledRuntimeStats(*task, true); + } else { + ASSERT_EQ(statsPair.first.spilledBytes, 0); + ASSERT_EQ(statsPair.first.spilledPartitions, 0); + ASSERT_EQ(statsPair.second.spilledBytes, 0); + ASSERT_EQ(statsPair.second.spilledPartitions, 0); + verifyTaskSpilledRuntimeStats(*task, false); + } + }) + .run(); + }); + + testWait.wait(testWaitKey); + ASSERT_TRUE(op != nullptr); + auto task = op->operatorCtx()->task(); + auto taskPauseWait = task->requestPause(); + driverWait.notify(); + taskPauseWait.wait(); + + uint64_t reclaimableBytes{0}; + const bool reclaimable = op->reclaimableBytes(reclaimableBytes); + ASSERT_EQ(op->canReclaim(), testData.expectedReclaimable); + ASSERT_EQ(reclaimable, testData.expectedReclaimable); + if (testData.expectedReclaimable) { + ASSERT_GT(reclaimableBytes, 0); + } else { + ASSERT_EQ(reclaimableBytes, 0); + } + + if (testData.expectedReclaimable) { + { + memory::ScopedMemoryArbitrationContext ctx(op->pool()); + op->pool()->reclaim( + folly::Random::oneIn(2) ? 0 : folly::Random::rand32(), + 0, + reclaimerStats_); + } + ASSERT_GT(reclaimerStats_.reclaimExecTimeUs, 0); + ASSERT_GT(reclaimerStats_.reclaimedBytes, 0); + reclaimerStats_.reset(); + ASSERT_EQ(op->pool()->usedBytes(), 0); + } else { + VELOX_ASSERT_THROW( + op->reclaim( + folly::Random::oneIn(2) ? 0 : folly::Random::rand32(), + reclaimerStats_), + ""); + } + + Task::resume(task); + task.reset(); + + taskThread.join(); + } + ASSERT_EQ(reclaimerStats_, memory::MemoryReclaimer::Stats{}); +} + +DEBUG_ONLY_TEST_F(HashJoinTest, reclaimDuringReserve) { + constexpr int64_t kMaxBytes = 1LL << 30; // 1GB + const int32_t numBuildVectors = 3; + std::vector buildVectors; + for (int32_t i = 0; i < numBuildVectors; ++i) { + const size_t size = i == 0 ? 1 : 1'000; + VectorFuzzer fuzzer({.vectorSize = size}, pool()); + buildVectors.push_back(fuzzer.fuzzRow(buildType_)); + } + + const int32_t numProbeVectors = 3; + std::vector probeVectors; + for (int32_t i = 0; i < numProbeVectors; ++i) { + VectorFuzzer fuzzer({.vectorSize = 1'000}, pool()); + probeVectors.push_back(fuzzer.fuzzRow(probeType_)); + } + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + auto tempDirectory = exec::test::TempDirectoryPath::create(); + auto queryPool = memory::memoryManager()->addRootPool( + "", kMaxBytes, memory::MemoryReclaimer::create()); + + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, false) + .hashJoin( + {"t_k1"}, + {"u_k1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, false) + .planNode(), + "", + concat(probeType_->names(), buildType_->names())) + .planNode(); + + folly::EventCount driverWait; + std::atomic_bool driverWaitFlag{true}; + folly::EventCount testWait; + std::atomic_bool testWaitFlag{true}; + + Operator* op{nullptr}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::addInput", + std::function(([&](Operator* testOp) { + if (testOp->operatorType() != "HashBuild") { + return; + } + op = testOp; + }))); + + std::atomic injectOnce{true}; + SCOPED_TESTVALUE_SET( + "facebook::velox::common::memory::MemoryPoolImpl::maybeReserve", + std::function( + ([&](memory::MemoryPoolImpl* pool) { + ASSERT_TRUE(op != nullptr); + if (!isHashBuildMemoryPool(*pool)) { + return; + } + ASSERT_TRUE(op->canReclaim()); + if (op->pool()->usedBytes() == 0) { + // We skip trigger memory reclaim when the hash table is empty on + // memory reservation. + return; + } + if (!injectOnce.exchange(false)) { + return; + } + uint64_t reclaimableBytes{0}; + const bool reclaimable = op->reclaimableBytes(reclaimableBytes); + ASSERT_TRUE(reclaimable); + ASSERT_GT(reclaimableBytes, 0); + auto* driver = op->operatorCtx()->driver(); + TestSuspendedSection suspendedSection(driver); + testWaitFlag = false; + testWait.notifyAll(); + driverWait.await([&]() { return !driverWaitFlag.load(); }); + }))); + + std::thread taskThread([&]() { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .planNode(plan) + .queryPool(std::move(queryPool)) + .injectSpill(false) + .spillDirectory(tempDirectory->getPath()) + .referenceQuery( + "SELECT t_k1, t_k2, t_v1, u_k1, u_k2, u_v1 FROM t, u WHERE t.t_k1 = u.u_k1") + .config(core::QueryConfig::kSpillStartPartitionBit, "29") + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + const auto statsPair = taskSpilledStats(*task); + ASSERT_GT(statsPair.first.spilledBytes, 0); + ASSERT_EQ(statsPair.first.spilledPartitions, 8); + ASSERT_GT(statsPair.second.spilledBytes, 0); + ASSERT_EQ(statsPair.second.spilledPartitions, 8); + verifyTaskSpilledRuntimeStats(*task, true); + }) + .run(); + }); + + testWait.await([&]() { return !testWaitFlag.load(); }); + ASSERT_TRUE(op != nullptr); + auto task = op->operatorCtx()->task(); + task->requestPause().wait(); + + uint64_t reclaimableBytes{0}; + const bool reclaimable = op->reclaimableBytes(reclaimableBytes); + ASSERT_TRUE(op->canReclaim()); + ASSERT_TRUE(reclaimable); + ASSERT_GT(reclaimableBytes, 0); + + { + memory::ScopedMemoryArbitrationContext ctx(op->pool()); + uint64_t reclaimedBytes = task->pool()->reclaim( + folly::Random::oneIn(2) ? 0 : folly::Random::rand32(), + 0, + reclaimerStats_); + ASSERT_GT(reclaimedBytes, 0); + } + ASSERT_GT(reclaimerStats_.reclaimedBytes, 0); + ASSERT_GT(reclaimerStats_.reclaimExecTimeUs, 0); + ASSERT_EQ(op->pool()->usedBytes(), 0); + + driverWaitFlag = false; + driverWait.notifyAll(); + Task::resume(task); + task.reset(); + + taskThread.join(); +} + +DEBUG_ONLY_TEST_F(HashJoinTest, reclaimDuringAllocation) { + constexpr int64_t kMaxBytes = 1LL << 30; // 1GB + VectorFuzzer fuzzer({.vectorSize = 1000}, pool()); + const int32_t numBuildVectors = 10; + std::vector buildVectors; + for (int32_t i = 0; i < numBuildVectors; ++i) { + buildVectors.push_back(fuzzer.fuzzRow(buildType_)); + } + const int32_t numProbeVectors = 5; + std::vector probeVectors; + for (int32_t i = 0; i < numProbeVectors; ++i) { + probeVectors.push_back(fuzzer.fuzzRow(probeType_)); + } + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + const std::vector enableSpillings = {false, true}; + for (const auto enableSpilling : enableSpillings) { + SCOPED_TRACE(fmt::format("enableSpilling {}", enableSpilling)); + + auto tempDirectory = exec::test::TempDirectoryPath::create(); + auto queryPool = memory::memoryManager()->addRootPool("", kMaxBytes); + + core::PlanNodeId probeScanId; + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, false) + .hashJoin( + {"t_k1"}, + {"u_k1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, false) + .planNode(), + "", + concat(probeType_->names(), buildType_->names())) + .planNode(); + + folly::EventCount driverWait; + auto driverWaitKey = driverWait.prepareWait(); + folly::EventCount testWait; + auto testWaitKey = testWait.prepareWait(); + + Operator* op; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::addInput", + std::function(([&](Operator* testOp) { + if (testOp->operatorType() != "HashBuild") { + return; + } + op = testOp; + }))); + + std::atomic injectOnce{true}; + SCOPED_TESTVALUE_SET( + "facebook::velox::common::memory::MemoryPoolImpl::allocateNonContiguous", + std::function( + ([&](memory::MemoryPoolImpl* pool) { + ASSERT_TRUE(op != nullptr); + const std::string re(".*HashBuild"); + if (!RE2::FullMatch(pool->name(), re)) { + return; + } + if (!injectOnce.exchange(false)) { + return; + } + ASSERT_EQ(op->canReclaim(), enableSpilling); + uint64_t reclaimableBytes{0}; + const bool reclaimable = op->reclaimableBytes(reclaimableBytes); + ASSERT_EQ(reclaimable, enableSpilling); + if (enableSpilling) { + ASSERT_GE(reclaimableBytes, 0); + } else { + ASSERT_EQ(reclaimableBytes, 0); + } + auto* driver = op->operatorCtx()->driver(); + TestSuspendedSection suspendedSection(driver); + testWait.notify(); + driverWait.wait(driverWaitKey); + }))); + + std::thread taskThread([&]() { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .planNode(plan) + .queryPool(std::move(queryPool)) + .injectSpill(false) + .spillDirectory(enableSpilling ? tempDirectory->getPath() : "") + .referenceQuery( + "SELECT t_k1, t_k2, t_v1, u_k1, u_k2, u_v1 FROM t, u WHERE t.t_k1 = u.u_k1") + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + const auto statsPair = taskSpilledStats(*task); + ASSERT_EQ(statsPair.first.spilledBytes, 0); + ASSERT_EQ(statsPair.first.spilledPartitions, 0); + ASSERT_EQ(statsPair.second.spilledBytes, 0); + ASSERT_EQ(statsPair.second.spilledPartitions, 0); + verifyTaskSpilledRuntimeStats(*task, false); + }) + .run(); + }); + + testWait.wait(testWaitKey); + ASSERT_TRUE(op != nullptr); + auto task = op->operatorCtx()->task(); + auto taskPauseWait = task->requestPause(); + taskPauseWait.wait(); + + uint64_t reclaimableBytes{0}; + const bool reclaimable = op->reclaimableBytes(reclaimableBytes); + ASSERT_EQ(op->canReclaim(), enableSpilling); + ASSERT_EQ(reclaimable, enableSpilling); + if (enableSpilling) { + ASSERT_GE(reclaimableBytes, 0); + } else { + ASSERT_EQ(reclaimableBytes, 0); + } + VELOX_ASSERT_THROW( + op->reclaim( + folly::Random::oneIn(2) ? 0 : folly::Random::rand32(), + reclaimerStats_), + ""); + + driverWait.notify(); + Task::resume(task); + task.reset(); + + taskThread.join(); + } + ASSERT_EQ(reclaimerStats_, memory::MemoryReclaimer::Stats{0}); +} + +DEBUG_ONLY_TEST_F(HashJoinTest, reclaimDuringOutputProcessing) { + constexpr int64_t kMaxBytes = 1LL << 30; // 1GB + VectorFuzzer fuzzer({.vectorSize = 1000}, pool()); + const int32_t numBuildVectors = 10; + std::vector buildVectors; + for (int32_t i = 0; i < numBuildVectors; ++i) { + buildVectors.push_back(fuzzer.fuzzRow(buildType_)); + } + const int32_t numProbeVectors = 5; + std::vector probeVectors; + for (int32_t i = 0; i < numProbeVectors; ++i) { + probeVectors.push_back(fuzzer.fuzzRow(probeType_)); + } + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + const std::vector enableSpillings = {false, true}; + for (const auto enableSpilling : enableSpillings) { + SCOPED_TRACE(fmt::format("enableSpilling {}", enableSpilling)); + auto tempDirectory = exec::test::TempDirectoryPath::create(); + auto queryPool = memory::memoryManager()->addRootPool( + "", kMaxBytes, memory::MemoryReclaimer::create()); + + core::PlanNodeId probeScanId; + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, false) + .hashJoin( + {"t_k1"}, + {"u_k1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, false) + .planNode(), + "", + concat(probeType_->names(), buildType_->names())) + .planNode(); + + std::atomic_bool driverWaitFlag{true}; + folly::EventCount driverWait; + std::atomic_bool testWaitFlag{true}; + folly::EventCount testWait; + + std::atomic injectOnce{true}; + Operator* op; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::noMoreInput", + std::function(([&](Operator* testOp) { + if (testOp->operatorType() != "HashBuild") { + return; + } + op = testOp; + if (!injectOnce.exchange(false)) { + return; + } + ASSERT_EQ(op->canReclaim(), enableSpilling); + uint64_t reclaimableBytes{0}; + const bool reclaimable = op->reclaimableBytes(reclaimableBytes); + ASSERT_EQ(reclaimable, enableSpilling); + if (enableSpilling) { + ASSERT_GT(reclaimableBytes, 0); + } else { + ASSERT_EQ(reclaimableBytes, 0); + } + testWaitFlag = false; + testWait.notifyAll(); + driverWait.await([&]() { return !driverWaitFlag.load(); }); + }))); + + std::thread taskThread([&]() { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .planNode(plan) + .queryPool(std::move(queryPool)) + .injectSpill(false) + .spillDirectory(enableSpilling ? tempDirectory->getPath() : "") + .referenceQuery( + "SELECT t_k1, t_k2, t_v1, u_k1, u_k2, u_v1 FROM t, u WHERE t.t_k1 = u.u_k1") + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + const auto statsPair = taskSpilledStats(*task); + ASSERT_EQ(statsPair.first.spilledBytes, 0); + ASSERT_EQ(statsPair.first.spilledPartitions, 0); + ASSERT_EQ(statsPair.second.spilledBytes, 0); + ASSERT_EQ(statsPair.second.spilledPartitions, 0); + verifyTaskSpilledRuntimeStats(*task, false); + }) + .run(); + }); + + testWait.await([&]() { return !testWaitFlag.load(); }); + ASSERT_TRUE(op != nullptr); + auto task = op->operatorCtx()->task(); + auto taskPauseWait = task->requestPause(); + driverWaitFlag = false; + driverWait.notifyAll(); + taskPauseWait.wait(); + + uint64_t reclaimableBytes{0}; + const bool reclaimable = op->reclaimableBytes(reclaimableBytes); + ASSERT_EQ(op->canReclaim(), enableSpilling); + ASSERT_EQ(reclaimable, enableSpilling); + + if (enableSpilling) { + ASSERT_GT(reclaimableBytes, 0); + const auto usedMemoryBytes = op->pool()->usedBytes(); + { + memory::ScopedMemoryArbitrationContext ctx(op->pool()); + op->pool()->reclaim( + folly::Random::oneIn(2) ? 0 : folly::Random::rand32(), + 0, + reclaimerStats_); + } + ASSERT_GE(reclaimerStats_.reclaimedBytes, 0); + ASSERT_GT(reclaimerStats_.reclaimExecTimeUs, 0); + // No reclaim as the operator has started output processing. + ASSERT_EQ(usedMemoryBytes, op->pool()->usedBytes()); + } else { + ASSERT_EQ(reclaimableBytes, 0); + VELOX_ASSERT_THROW( + op->reclaim( + folly::Random::oneIn(2) ? 0 : folly::Random::rand32(), + reclaimerStats_), + ""); + } + + Task::resume(task); + task.reset(); + + taskThread.join(); + } + ASSERT_EQ(reclaimerStats_.numNonReclaimableAttempts, 1); +} + +DEBUG_ONLY_TEST_F(HashJoinTest, reclaimDuringWaitForProbe) { + constexpr int64_t kMaxBytes = 1LL << 30; // 1GB + VectorFuzzer fuzzer({.vectorSize = 1000}, pool()); + const int32_t numBuildVectors = 10; + std::vector buildVectors; + for (int32_t i = 0; i < numBuildVectors; ++i) { + buildVectors.push_back(fuzzer.fuzzRow(buildType_)); + } + const int32_t numProbeVectors = 5; + std::vector probeVectors; + for (int32_t i = 0; i < numProbeVectors; ++i) { + probeVectors.push_back(fuzzer.fuzzRow(probeType_)); + } + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + auto tempDirectory = exec::test::TempDirectoryPath::create(); + auto queryPool = memory::memoryManager()->addRootPool( + "", kMaxBytes, memory::MemoryReclaimer::create()); + + core::PlanNodeId probeScanId; + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, false) + .hashJoin( + {"t_k1"}, + {"u_k1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, false) + .planNode(), + "", + concat(probeType_->names(), buildType_->names())) + .planNode(); + + std::atomic_bool driverWaitFlag{true}; + folly::EventCount driverWait; + std::atomic_bool testWaitFlag{true}; + folly::EventCount testWait; + + Operator* op{nullptr}; + std::atomic injectSpillOnce{true}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::HashBuild::finishHashBuild", + std::function(([&](Operator* testOp) { + if (testOp->operatorType() != "HashBuild") { + return; + } + op = testOp; + if (!injectSpillOnce.exchange(false)) { + return; + } + auto* driver = op->operatorCtx()->driver(); + auto task = driver->task(); + memory::ScopedMemoryArbitrationContext ctx(op->pool()); + Operator::ReclaimableSectionGuard guard(testOp); + testingRunArbitration(testOp->pool()); + }))); + + std::atomic injectOnce{true}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::noMoreInput", + std::function(([&](Operator* testOp) { + if (testOp->operatorType() != "HashProbe") { + return; + } + if (!injectOnce.exchange(false)) { + return; + } + ASSERT_TRUE(op != nullptr); + ASSERT_TRUE(op->canReclaim()); + uint64_t reclaimableBytes{0}; + const bool reclaimable = op->reclaimableBytes(reclaimableBytes); + ASSERT_TRUE(reclaimable); + ASSERT_GT(reclaimableBytes, 0); + testWaitFlag = false; + testWait.notifyAll(); + auto* driver = testOp->operatorCtx()->driver(); + auto task = driver->task(); + TestSuspendedSection suspendedSection(driver); + driverWait.await([&]() { return !driverWaitFlag.load(); }); + }))); + + std::thread taskThread([&]() { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .planNode(plan) + .queryPool(std::move(queryPool)) + .injectSpill(false) + .spillDirectory(tempDirectory->getPath()) + .referenceQuery( + "SELECT t_k1, t_k2, t_v1, u_k1, u_k2, u_v1 FROM t, u WHERE t.t_k1 = u.u_k1") + .config(core::QueryConfig::kSpillStartPartitionBit, "29") + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + const auto statsPair = taskSpilledStats(*task); + ASSERT_GT(statsPair.first.spilledBytes, 0); + ASSERT_EQ(statsPair.first.spilledPartitions, 8); + ASSERT_GT(statsPair.second.spilledBytes, 0); + ASSERT_EQ(statsPair.second.spilledPartitions, 8); + }) + .run(); + }); + + testWait.await([&]() { return !testWaitFlag.load(); }); + ASSERT_TRUE(op != nullptr); + auto task = op->operatorCtx()->task(); + auto taskPauseWait = task->requestPause(); + taskPauseWait.wait(); + + uint64_t reclaimableBytes{0}; + const bool reclaimable = op->reclaimableBytes(reclaimableBytes); + ASSERT_TRUE(op->canReclaim()); + ASSERT_TRUE(reclaimable); + ASSERT_GT(reclaimableBytes, 0); + + const auto usedMemoryBytes = op->pool()->usedBytes(); + reclaimerStats_.reset(); + { + memory::ScopedMemoryArbitrationContext ctx(op->pool()); + op->pool()->reclaim( + folly::Random::oneIn(2) ? 0 : folly::Random::rand32(), + 0, + reclaimerStats_); + } + ASSERT_GE(reclaimerStats_.reclaimedBytes, 0); + ASSERT_GT(reclaimerStats_.reclaimExecTimeUs, 0); + // No reclaim as the build operator is not in building table state. + ASSERT_EQ(usedMemoryBytes, op->pool()->usedBytes()); + + driverWaitFlag = false; + driverWait.notifyAll(); + Task::resume(task); + task.reset(); + + taskThread.join(); + ASSERT_EQ(reclaimerStats_.numNonReclaimableAttempts, 1); +} + +DEBUG_ONLY_TEST_F(HashJoinTest, hashBuildAbortDuringOutputProcessing) { + const auto buildVectors = makeVectors(buildType_, 10, 128); + const auto probeVectors = makeVectors(probeType_, 5, 128); + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + struct { + bool abortFromRootMemoryPool; + int numDrivers; + + std::string debugString() const { + return fmt::format( + "abortFromRootMemoryPool {} numDrivers {}", + abortFromRootMemoryPool, + numDrivers); + } + } testSettings[] = {{true, 1}, {false, 1}, {true, 4}, {false, 4}}; + + for (const auto& testData : testSettings) { + SCOPED_TRACE(testData.debugString()); + + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, true) + .hashJoin( + {"t_k1"}, + {"u_k1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, true) + .planNode(), + "", + concat(probeType_->names(), buildType_->names())) + .planNode(); + + std::atomic injectOnce{true}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::noMoreInput", + std::function(([&](Operator* op) { + if (op->operatorType() != "HashBuild") { + return; + } + if (!injectOnce.exchange(false)) { + return; + } + ASSERT_GT(op->pool()->usedBytes(), 0); + auto* driver = op->operatorCtx()->driver(); + ASSERT_EQ( + driver->task()->enterSuspended(driver->state()), + StopReason::kNone); + testData.abortFromRootMemoryPool ? abortPool(op->pool()->root()) + : abortPool(op->pool()); + // We can't directly reclaim memory from this hash build operator as + // its driver thread is running and in suspension state. + ASSERT_GT(op->pool()->root()->usedBytes(), 0); + ASSERT_EQ( + driver->task()->leaveSuspended(driver->state()), + StopReason::kAlreadyTerminated); + ASSERT_TRUE(op->pool()->aborted()); + ASSERT_TRUE(op->pool()->root()->aborted()); + VELOX_MEM_POOL_ABORTED("Memory pool aborted"); + }))); + + VELOX_ASSERT_THROW( + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .planNode(plan) + .injectSpill(false) + .referenceQuery( + "SELECT t_k1, t_k2, t_v1, u_k1, u_k2, u_v1 FROM t, u WHERE t.t_k1 = u.u_k1") + .run(), + "Manual MemoryPool Abortion"); + waitForAllTasksToBeDeleted(); + } +} + +DEBUG_ONLY_TEST_F(HashJoinTest, hashBuildAbortDuringInputProcessing) { + const auto buildVectors = makeVectors(buildType_, 10, 128); + const auto probeVectors = makeVectors(probeType_, 5, 128); + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + struct { + bool abortFromRootMemoryPool; + int numDrivers; + + std::string debugString() const { + return fmt::format( + "abortFromRootMemoryPool {} numDrivers {}", + abortFromRootMemoryPool, + numDrivers); + } + } testSettings[] = {{true, 1}, {false, 1}, {true, 4}, {false, 4}}; + + for (const auto& testData : testSettings) { + SCOPED_TRACE(testData.debugString()); + + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, true) + .hashJoin( + {"t_k1"}, + {"u_k1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, true) + .planNode(), + "", + concat(probeType_->names(), buildType_->names())) + .planNode(); + + std::atomic numInputs{0}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::addInput", + std::function(([&](Operator* op) { + if (op->operatorType() != "HashBuild") { + return; + } + if (++numInputs != 2) { + return; + } + ASSERT_GT(op->pool()->usedBytes(), 0); + auto* driver = op->operatorCtx()->driver(); + ASSERT_EQ( + driver->task()->enterSuspended(driver->state()), + StopReason::kNone); + testData.abortFromRootMemoryPool ? abortPool(op->pool()->root()) + : abortPool(op->pool()); + // We can't directly reclaim memory from this hash build operator as + // its driver thread is running and in suspension state. + ASSERT_GT(op->pool()->root()->usedBytes(), 0); + ASSERT_EQ( + driver->task()->leaveSuspended(driver->state()), + StopReason::kAlreadyTerminated); + ASSERT_TRUE(op->pool()->aborted()); + ASSERT_TRUE(op->pool()->root()->aborted()); + VELOX_MEM_POOL_ABORTED("Memory pool aborted"); + }))); + + VELOX_ASSERT_THROW( + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .planNode(plan) + .injectSpill(false) + .referenceQuery( + "SELECT t_k1, t_k2, t_v1, u_k1, u_k2, u_v1 FROM t, u WHERE t.t_k1 = u.u_k1") + .run(), + "Manual MemoryPool Abortion"); + + waitForAllTasksToBeDeleted(); + } +} + +DEBUG_ONLY_TEST_F(HashJoinTest, hashBuildAbortDuringAllocation) { + const auto buildVectors = makeVectors(buildType_, 10, 128); + const auto probeVectors = makeVectors(probeType_, 5, 128); + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + struct { + bool abortFromRootMemoryPool; + int numDrivers; + + std::string debugString() const { + return fmt::format( + "abortFromRootMemoryPool {} numDrivers {}", + abortFromRootMemoryPool, + numDrivers); + } + } testSettings[] = {{true, 1}, {false, 1}, {true, 4}, {false, 4}}; + + for (const auto& testData : testSettings) { + SCOPED_TRACE(testData.debugString()); + + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, true) + .hashJoin( + {"t_k1"}, + {"u_k1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, true) + .planNode(), + "", + concat(probeType_->names(), buildType_->names())) + .planNode(); + + std::atomic_bool injectOnce{true}; + SCOPED_TESTVALUE_SET( + "facebook::velox::common::memory::MemoryPoolImpl::allocateNonContiguous", + std::function( + ([&](memory::MemoryPoolImpl* pool) { + if (!isHashBuildMemoryPool(*pool)) { + return; + } + if (!injectOnce.exchange(false)) { + return; + } + + const auto* driverCtx = driverThreadContext()->driverCtx(); + ASSERT_EQ( + driverCtx->task->enterSuspended(driverCtx->driver->state()), + StopReason::kNone); + testData.abortFromRootMemoryPool ? abortPool(pool->root()) + : abortPool(pool); + // We can't directly reclaim memory from this hash build operator + // as its driver thread is running and in suspegnsion state. + ASSERT_GE(pool->root()->usedBytes(), 0); + ASSERT_EQ( + driverCtx->task->leaveSuspended(driverCtx->driver->state()), + StopReason::kAlreadyTerminated); + ASSERT_TRUE(pool->aborted()); + ASSERT_TRUE(pool->root()->aborted()); + VELOX_MEM_POOL_ABORTED("Memory pool aborted"); + }))); + + VELOX_ASSERT_THROW( + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .planNode(plan) + .injectSpill(false) + .referenceQuery( + "SELECT t_k1, t_k2, t_v1, u_k1, u_k2, u_v1 FROM t, u WHERE t.t_k1 = u.u_k1") + .run(), + "Manual MemoryPool Abortion"); + + waitForAllTasksToBeDeleted(); + } +} + +DEBUG_ONLY_TEST_F(HashJoinTest, hashProbeAbortDuringInputProcessing) { + const auto buildVectors = makeVectors(buildType_, 10, 128); + const auto probeVectors = makeVectors(probeType_, 5, 128); + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + struct { + bool abortFromRootMemoryPool; + int numDrivers; + + std::string debugString() const { + return fmt::format( + "abortFromRootMemoryPool {} numDrivers {}", + abortFromRootMemoryPool, + numDrivers); + } + } testSettings[] = {{true, 1}, {false, 1}, {true, 4}, {false, 4}}; + + for (const auto& testData : testSettings) { + SCOPED_TRACE(testData.debugString()); + + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, true) + .hashJoin( + {"t_k1"}, + {"u_k1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, true) + .planNode(), + "", + concat(probeType_->names(), buildType_->names())) + .planNode(); + + std::atomic numInputs{0}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::addInput", + std::function(([&](Operator* op) { + if (op->operatorType() != "HashProbe") { + return; + } + if (++numInputs != 2) { + return; + } + auto* driver = op->operatorCtx()->driver(); + ASSERT_EQ( + driver->task()->enterSuspended(driver->state()), + StopReason::kNone); + testData.abortFromRootMemoryPool ? abortPool(op->pool()->root()) + : abortPool(op->pool()); + ASSERT_EQ( + driver->task()->leaveSuspended(driver->state()), + StopReason::kAlreadyTerminated); + ASSERT_TRUE(op->pool()->aborted()); + ASSERT_TRUE(op->pool()->root()->aborted()); + VELOX_MEM_POOL_ABORTED("Memory pool aborted"); + }))); + + VELOX_ASSERT_THROW( + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .planNode(plan) + .injectSpill(false) + .referenceQuery( + "SELECT t_k1, t_k2, t_v1, u_k1, u_k2, u_v1 FROM t, u WHERE t.t_k1 = u.u_k1") + .run(), + "Manual MemoryPool Abortion"); + waitForAllTasksToBeDeleted(); + } +} + +TEST_F(HashJoinTest, leftJoinWithMissAtEndOfBatch) { + // Tests some cases where the row at the end of an output batch fails the + // filter. + auto probeVectors = std::vector{makeRowVector( + {"t_k1", "t_k2"}, + {makeFlatVector(20, [](auto row) { return 1 + row % 2; }), + makeFlatVector(20, [](auto row) { return row; })})}; + auto buildVectors = std::vector{ + makeRowVector({"u_k1"}, {makeFlatVector({1, 2})})}; + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", {buildVectors}); + auto planNodeIdGenerator = std::make_shared(); + + auto test = [&](const std::string& filter) { + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, true) + .hashJoin( + {"t_k1"}, + {"u_k1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, true) + .planNode(), + filter, + {"t_k1", "u_k1"}, + core::JoinType::kLeft) + .planNode(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(plan) + .injectSpill(false) + .checkSpillStats(false) + .maxSpillLevel(0) + .numDrivers(1) + .config( + core::QueryConfig::kPreferredOutputBatchRows, std::to_string(10)) + .referenceQuery(fmt::format( + "SELECT t_k1, u_k1 from t left join u on t_k1 = u_k1 and {}", + filter)) + .run(); + }; + + // Alternate rows pass this filter and last row of a batch fails. + test("t_k1=1"); + + // All rows fail this filter. + test("t_k1=5"); + + // All rows in the second batch pass this filter. + test("t_k2 > 9"); +} + +TEST_F(HashJoinTest, leftJoinWithMissAtEndOfBatchMultipleBuildMatches) { + // Tests some cases where the row at the end of an output batch fails the + // filter and there are multiple matches with the build side.. + auto probeVectors = std::vector{makeRowVector( + {"t_k1", "t_k2"}, + {makeFlatVector(10, [](auto row) { return 1 + row % 2; }), + makeFlatVector(10, [](auto row) { return row; })})}; + auto buildVectors = std::vector{ + makeRowVector({"u_k1"}, {makeFlatVector({1, 2, 1, 2})})}; + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", {buildVectors}); + auto planNodeIdGenerator = std::make_shared(); + + auto test = [&](const std::string& filter) { + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, true) + .hashJoin( + {"t_k1"}, + {"u_k1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, true) + .planNode(), + filter, + {"t_k1", "u_k1"}, + core::JoinType::kLeft) + .planNode(); + + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(plan) + .injectSpill(false) + .checkSpillStats(false) + .maxSpillLevel(0) + .numDrivers(1) + .config( + core::QueryConfig::kPreferredOutputBatchRows, std::to_string(10)) + .referenceQuery(fmt::format( + "SELECT t_k1, u_k1 from t left join u on t_k1 = u_k1 and {}", + filter)) + .run(); + }; + + // In this case the rows with t_k2 = 4 appear at the end of the first batch, + // meaning the last rows in that output batch are misses, and don't get added. + // The rows with t_k2 = 8 appear in the second batch so only one row is + // written, meaning there is space in the second output batch for the miss + // with tk_2 = 4 to get written. + test("t_k2 != 4 and t_k2 != 8"); +} + +TEST_F(HashJoinTest, leftJoinPreserveProbeOrder) { + const std::vector probeVectors = { + makeRowVector( + {"k1", "v1"}, + { + makeConstant(0, 2), + makeFlatVector({1, 0}), + }), + }; + const std::vector buildVectors = { + makeRowVector( + {"k2", "v2"}, + { + makeConstant(0, 2), + makeConstant(0, 2), + }), + }; + auto planNodeIdGenerator = std::make_shared(); + auto plan = + PlanBuilder(planNodeIdGenerator) + .values(probeVectors) + .hashJoin( + {"k1"}, + {"k2"}, + PlanBuilder(planNodeIdGenerator).values(buildVectors).planNode(), + "v1 % 2 = v2 % 2", + {"v1"}, + core::JoinType::kLeft) + .planNode(); + auto result = AssertQueryBuilder(plan) + .config(core::QueryConfig::kPreferredOutputBatchRows, "1") + .serialExecution(true) + .copyResults(pool_.get()); + ASSERT_EQ(result->size(), 3); + auto* v1 = + result->childAt(0)->loadedVector()->asUnchecked>(); + ASSERT_FALSE(v1->mayHaveNulls()); + ASSERT_EQ(v1->valueAt(0), 1); + ASSERT_EQ(v1->valueAt(1), 0); + ASSERT_EQ(v1->valueAt(2), 0); +} + +DEBUG_ONLY_TEST_F(HashJoinTest, minSpillableMemoryReservation) { + VectorFuzzer fuzzer({.vectorSize = 1000}, pool()); + const int32_t numBuildVectors = 10; + std::vector buildVectors; + for (int32_t i = 0; i < numBuildVectors; ++i) { + buildVectors.push_back(fuzzer.fuzzInputRow(buildType_)); + } + const int32_t numProbeVectors = 5; + std::vector probeVectors; + for (int32_t i = 0; i < numProbeVectors; ++i) { + probeVectors.push_back(fuzzer.fuzzInputRow(probeType_)); + } + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + core::PlanNodeId probeScanId; + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, false) + .hashJoin( + {"t_k1"}, + {"u_k1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, false) + .planNode(), + "", + concat(probeType_->names(), buildType_->names())) + .planNode(); + + for (int32_t minSpillableReservationPct : {5, 50, 100}) { + SCOPED_TRACE(fmt::format( + "minSpillableReservationPct: {}", minSpillableReservationPct)); + + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::HashBuild::addInput", + std::function(([&](exec::HashBuild* hashBuild) { + memory::MemoryPool* pool = hashBuild->pool(); + const auto availableReservationBytes = pool->availableReservation(); + const auto currentUsedBytes = pool->usedBytes(); + // Verifies we always have min reservation after ensuring the input. + ASSERT_GE( + availableReservationBytes, + currentUsedBytes * minSpillableReservationPct / 100); + }))); + + auto tempDirectory = exec::test::TempDirectoryPath::create(); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .planNode(plan) + .injectSpill(false) + .spillDirectory(tempDirectory->getPath()) + .referenceQuery( + "SELECT t_k1, t_k2, t_v1, u_k1, u_k2, u_v1 FROM t, u WHERE t.t_k1 = u.u_k1") + .run(); + } +} + +DEBUG_ONLY_TEST_F(HashJoinTest, exceededMaxSpillLevel) { + VectorFuzzer fuzzer({.vectorSize = 1000}, pool()); + const int32_t numBuildVectors = 10; + std::vector buildVectors; + for (int32_t i = 0; i < numBuildVectors; ++i) { + buildVectors.push_back(fuzzer.fuzzRow(buildType_)); + } + const int32_t numProbeVectors = 5; + std::vector probeVectors; + for (int32_t i = 0; i < numProbeVectors; ++i) { + probeVectors.push_back(fuzzer.fuzzRow(probeType_)); + } + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + core::PlanNodeId probeScanId; + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, false) + .hashJoin( + {"t_k1"}, + {"u_k1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, false) + .planNode(), + "", + concat(probeType_->names(), buildType_->names())) + .planNode(); + + auto tempDirectory = exec::test::TempDirectoryPath::create(); + const int exceededMaxSpillLevelCount = + common::globalSpillStats().spillMaxLevelExceededCount; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::HashBuild::reclaim", + std::function(([&](exec::Operator* op) { + HashBuild* hashBuild = static_cast(op); + ASSERT_FALSE(hashBuild->testingExceededMaxSpillLevelLimit()); + }))); + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::HashProbe::reclaim", + std::function(([&](exec::Operator* op) { + HashProbe* hashProbe = static_cast(op); + ASSERT_FALSE(hashProbe->testingExceededMaxSpillLevelLimit()); + }))); + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::HashBuild::finishHashBuild", + std::function(([&](exec::HashBuild* hashBuild) { + Operator::ReclaimableSectionGuard guard(hashBuild); + testingRunArbitration(hashBuild->pool()); + }))); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(1) + .planNode(plan) + // Always trigger spilling. + .injectSpill(false) + .maxSpillLevel(0) + .spillDirectory(tempDirectory->getPath()) + .referenceQuery( + "SELECT t_k1, t_k2, t_v1, u_k1, u_k2, u_v1 FROM t, u WHERE t.t_k1 = u.u_k1") + .config(core::QueryConfig::kSpillStartPartitionBit, "29") + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + auto opStats = toOperatorStats(task->taskStats()); + ASSERT_EQ( + opStats.at("HashProbe") + .runtimeStats[Operator::kExceededMaxSpillLevel] + .sum, + 8); + ASSERT_EQ( + opStats.at("HashProbe") + .runtimeStats[Operator::kExceededMaxSpillLevel] + .count, + 1); + ASSERT_EQ( + opStats.at("HashBuild") + .runtimeStats[Operator::kExceededMaxSpillLevel] + .sum, + 8); + ASSERT_EQ( + opStats.at("HashBuild") + .runtimeStats[Operator::kExceededMaxSpillLevel] + .count, + 1); + }) + .run(); + ASSERT_EQ( + common::globalSpillStats().spillMaxLevelExceededCount, + exceededMaxSpillLevelCount + 16); +} + +// disable because spilling is not supported for velox-cudf hash join +TEST_F(HashJoinTest, DISABLED_maxSpillBytes) { + const auto rowType = + ROW({"c0", "c1", "c2"}, {INTEGER(), INTEGER(), VARCHAR()}); + const auto probeVectors = createVectors(rowType, 1024, 10 << 20); + const auto buildVectors = createVectors(rowType, 1024, 10 << 20); + + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, true) + .project({"c0", "c1", "c2"}) + .hashJoin( + {"c0"}, + {"u1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, true) + .project({"c0 AS u0", "c1 AS u1", "c2 AS u2"}) + .planNode(), + "", + {"c0", "c1", "c2"}, + core::JoinType::kInner) + .planNode(); + + auto spillDirectory = exec::test::TempDirectoryPath::create(); + auto queryCtx = core::QueryCtx::create(executor_.get()); + + struct { + int32_t maxSpilledBytes; + bool expectedExceedLimit; + std::string debugString() const { + return fmt::format("maxSpilledBytes {}", maxSpilledBytes); + } + } testSettings[] = {{1 << 30, false}, {16 << 20, true}, {0, false}}; + + for (const auto& testData : testSettings) { + SCOPED_TRACE(testData.debugString()); + try { + TestScopedSpillInjection scopedSpillInjection(100); + AssertQueryBuilder(plan) + .spillDirectory(spillDirectory->getPath()) + .queryCtx(queryCtx) + .config(core::QueryConfig::kSpillEnabled, true) + .config(core::QueryConfig::kJoinSpillEnabled, true) + .config(core::QueryConfig::kMaxSpillBytes, testData.maxSpilledBytes) + .copyResults(pool_.get()); + ASSERT_FALSE(testData.expectedExceedLimit); + } catch (const VeloxRuntimeError& e) { + ASSERT_TRUE(testData.expectedExceedLimit); + ASSERT_NE( + e.message().find( + "Query exceeded per-query local spill limit of 16.00MB"), + std::string::npos); + ASSERT_EQ( + e.errorCode(), facebook::velox::error_code::kSpillLimitExceeded); + } + } +} + +// disable because spilling is not supported for velox-cudf hash join +TEST_F(HashJoinTest, DISABLED_onlyHashBuildMaxSpillBytes) { + const auto rowType = + ROW({"c0", "c1", "c2"}, {INTEGER(), INTEGER(), VARCHAR()}); + const auto probeVectors = createVectors(rowType, 32, 128); + const auto buildVectors = createVectors(rowType, 1024, 10 << 20); + + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, true) + .hashJoin( + {"c0"}, + {"u1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, true) + .project({"c0 AS u0", "c1 AS u1", "c2 AS u2"}) + .planNode(), + "", + {"c0", "c1", "c2"}, + core::JoinType::kInner) + .planNode(); + + auto spillDirectory = exec::test::TempDirectoryPath::create(); + auto queryCtx = core::QueryCtx::create(executor_.get()); + + struct { + int32_t maxSpilledBytes; + bool expectedExceedLimit; + std::string debugString() const { + return fmt::format("maxSpilledBytes {}", maxSpilledBytes); + } + } testSettings[] = {{1 << 30, false}, {16 << 20, true}, {0, false}}; + + for (const auto& testData : testSettings) { + SCOPED_TRACE(testData.debugString()); + try { + TestScopedSpillInjection scopedSpillInjection(100); + AssertQueryBuilder(plan) + .spillDirectory(spillDirectory->getPath()) + .queryCtx(queryCtx) + .config(core::QueryConfig::kSpillEnabled, true) + .config(core::QueryConfig::kJoinSpillEnabled, true) + .config(core::QueryConfig::kMaxSpillBytes, testData.maxSpilledBytes) + .copyResults(pool_.get()); + ASSERT_FALSE(testData.expectedExceedLimit); + } catch (const VeloxRuntimeError& e) { + ASSERT_TRUE(testData.expectedExceedLimit); + ASSERT_NE( + e.message().find( + "Query exceeded per-query local spill limit of 16.00MB"), + std::string::npos); + ASSERT_EQ( + e.errorCode(), facebook::velox::error_code::kSpillLimitExceeded); + } + } +} + +// disable because spilling is not supported for velox-cudf hash join +TEST_F(HashJoinTest, DISABLED_reclaimFromJoinBuilderWithMultiDrivers) { + auto rowType = ROW({ + {"c0", INTEGER()}, + {"c1", INTEGER()}, + {"c2", VARCHAR()}, + }); + const auto vectors = createVectors(rowType, 64 << 20, fuzzerOpts_); + const int numDrivers = 4; + + memory::MemoryManagerOptions options; + options.allocatorCapacity = 8L << 30; + auto memoryManagerWithoutArbitrator = + std::make_unique(options); + const auto expectedResult = + runHashJoinTask( + vectors, + newQueryCtx( + memoryManagerWithoutArbitrator.get(), executor_.get(), 8L << 30), + false, + numDrivers, + pool(), + false) + .data; + + auto memoryManagerWithArbitrator = createMemoryManager(); + const auto& arbitrator = memoryManagerWithArbitrator->arbitrator(); + // Create a query ctx with a small capacity to trigger spilling. + auto result = runHashJoinTask( + vectors, + newQueryCtx( + memoryManagerWithArbitrator.get(), executor_.get(), 128 << 20), + false, + numDrivers, + pool(), + true, + expectedResult); + auto taskStats = exec::toPlanStats(result.task->taskStats()); + auto& planStats = taskStats.at(result.planNodeId); + ASSERT_GT(planStats.spilledBytes, 0); + result.task.reset(); + + // This test uses on-demand created memory manager instead of the global + // one. We need to make sure any used memory got cleaned up before exiting + // the scope + waitForAllTasksToBeDeleted(); + ASSERT_GT(arbitrator->stats().numRequests, 0); + ASSERT_GT(arbitrator->stats().reclaimedUsedBytes, 0); +} + +DEBUG_ONLY_TEST_F( + HashJoinTest, + failedToReclaimFromHashJoinBuildersInNonReclaimableSection) { + auto rowType = ROW({ + {"c0", INTEGER()}, + {"c1", INTEGER()}, + {"c2", VARCHAR()}, + }); + const auto vectors = createVectors(rowType, 64 << 20, fuzzerOpts_); + const int numDrivers = 1; + std::shared_ptr queryCtx = + newQueryCtx(memory::memoryManager(), executor_.get(), 512 << 20); + const auto expectedResult = + runHashJoinTask(vectors, queryCtx, false, numDrivers, pool(), false).data; + + std::atomic_bool nonReclaimableSectionWaitFlag{true}; + std::atomic_bool reclaimerInitializationWaitFlag{true}; + folly::EventCount nonReclaimableSectionWait; + std::atomic_bool memoryArbitrationWaitFlag{true}; + folly::EventCount memoryArbitrationWait; + + std::atomic numInitializedDrivers{0}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal", + std::function([&](exec::Driver* driver) { + numInitializedDrivers++; + // We need to make sure reclaimers on both build and probe side are set + // (in Operator::initialize) to avoid race conditions, producing + // consistent test results. + if (numInitializedDrivers.load() == 2) { + reclaimerInitializationWaitFlag = false; + nonReclaimableSectionWait.notifyAll(); + } + })); + + std::atomic injectNonReclaimableSectionOnce{true}; + SCOPED_TESTVALUE_SET( + "facebook::velox::common::memory::MemoryPoolImpl::allocateNonContiguous", + std::function( + ([&](memory::MemoryPoolImpl* pool) { + if (!isHashBuildMemoryPool(*pool)) { + return; + } + if (!injectNonReclaimableSectionOnce.exchange(false)) { + return; + } + + // Signal the test control that one of the hash build operator has + // entered into non-reclaimable section. + nonReclaimableSectionWaitFlag = false; + nonReclaimableSectionWait.notifyAll(); + + // Suspend the driver to simulate the arbitration. + pool->reclaimer()->enterArbitration(); + // Wait for the memory arbitration to complete. + memoryArbitrationWait.await( + [&]() { return !memoryArbitrationWaitFlag.load(); }); + pool->reclaimer()->leaveArbitration(); + }))); + + std::thread joinThread([&]() { + const auto result = runHashJoinTask( + vectors, queryCtx, false, numDrivers, pool(), true, expectedResult); + auto taskStats = exec::toPlanStats(result.task->taskStats()); + auto& planStats = taskStats.at(result.planNodeId); + ASSERT_EQ(planStats.spilledBytes, 0); + }); + + // Wait for the hash build operators to enter into non-reclaimable section. + nonReclaimableSectionWait.await([&]() { + return ( + !nonReclaimableSectionWaitFlag.load() && + !reclaimerInitializationWaitFlag.load()); + }); + + // We expect capacity grow fails as we can't reclaim from hash join operators. + memory::testingRunArbitration(); + + // Notify the hash build operator that memory arbitration has been done. + memoryArbitrationWaitFlag = false; + memoryArbitrationWait.notifyAll(); + + joinThread.join(); + + // This test uses on-demand created memory manager instead of the global + // one. We need to make sure any used memory got cleaned up before exiting + // the scope + waitForAllTasksToBeDeleted(); + ASSERT_EQ( + memory::memoryManager()->arbitrator()->stats().numNonReclaimableAttempts, + 2); +} + +DEBUG_ONLY_TEST_F(HashJoinTest, reclaimDuringTableBuild) { + VectorFuzzer fuzzer({.vectorSize = 1000}, pool()); + const int32_t numBuildVectors = 5; + std::vector buildVectors; + for (int32_t i = 0; i < numBuildVectors; ++i) { + buildVectors.push_back(fuzzer.fuzzRow(buildType_)); + } + const int32_t numProbeVectors = 5; + std::vector probeVectors; + for (int32_t i = 0; i < numProbeVectors; ++i) { + probeVectors.push_back(fuzzer.fuzzRow(probeType_)); + } + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + core::PlanNodeId probeScanId; + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, false) + .hashJoin( + {"t_k1"}, + {"u_k1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, false) + .planNode(), + "", + concat(probeType_->names(), buildType_->names())) + .planNode(); + + std::atomic_bool injectSpillOnce{true}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::HashBuild::finishHashBuild", + std::function([&](Operator* op) { + if (!injectSpillOnce.exchange(false)) { + return; + } + Operator::ReclaimableSectionGuard guard(op); + testingRunArbitration(op->pool()); + })); + + auto tempDirectory = exec::test::TempDirectoryPath::create(); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(4) + .planNode(plan) + .injectSpill(false) + .maxSpillLevel(0) + .spillDirectory(tempDirectory->getPath()) + .referenceQuery( + "SELECT t_k1, t_k2, t_v1, u_k1, u_k2, u_v1 FROM t, u WHERE t.t_k1 = u.u_k1") + .config(core::QueryConfig::kSpillStartPartitionBit, "29") + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + auto opStats = toOperatorStats(task->taskStats()); + ASSERT_GT( + opStats.at("HashBuild").runtimeStats[Operator::kSpillWrites].sum, + 0); + }) + .run(); +} + +DEBUG_ONLY_TEST_F(HashJoinTest, exceptionDuringFinishJoinBuild) { + // This test is to make sure there is no memory leak when exceptions are + // thrown while parallelly preparing join table. + auto memoryManager = memory::memoryManager(); + const auto& arbitrator = memoryManager->arbitrator(); + const uint64_t numDrivers = 2; + const auto expectedFreeCapacityBytes = arbitrator->stats().freeCapacityBytes; + + const uint64_t numBuildSideRows = 500; + auto buildKeyVector = makeFlatVector( + numBuildSideRows, + [](vector_size_t row) { return folly::Random::rand64(); }); + auto buildSideVector = + makeRowVector({"b0", "b1"}, {buildKeyVector, buildKeyVector}); + std::vector buildSideVectors; + for (int i = 0; i < numDrivers; ++i) { + buildSideVectors.push_back(buildSideVector); + } + createDuckDbTable("build", buildSideVectors); + + const uint64_t numProbeSideRows = 10; + auto probeKeyVector = makeFlatVector( + numProbeSideRows, + [&](vector_size_t row) { return buildKeyVector->valueAt(row); }); + auto probeSideVector = + makeRowVector({"p0", "p1"}, {probeKeyVector, probeKeyVector}); + std::vector probeSideVectors; + for (int i = 0; i < numDrivers; ++i) { + probeSideVectors.push_back(probeSideVector); + } + createDuckDbTable("probe", probeSideVectors); + + ASSERT_EQ(arbitrator->stats().freeCapacityBytes, expectedFreeCapacityBytes); + + // We set the task to fail right before we reserve memory for other operators. + // We rely on the driver suspension before parallel join build to throw + // exceptions (suspension on an already terminated task throws). + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::HashBuild::ensureTableFits", + std::function([&](HashBuild* buildOp) { + try { + VELOX_FAIL("Simulated failure"); + } catch (VeloxException&) { + buildOp->operatorCtx()->task()->setError(std::current_exception()); + } + })); + + std::vector probeInput = {probeSideVector}; + std::vector buildInput = {buildSideVector}; + auto planNodeIdGenerator = std::make_shared(); + const auto spillDirectory = exec::test::TempDirectoryPath::create(); + + ASSERT_EQ(arbitrator->stats().freeCapacityBytes, expectedFreeCapacityBytes); + VELOX_ASSERT_THROW( + AssertQueryBuilder(duckDbQueryRunner_) + .spillDirectory(spillDirectory->getPath()) + .config(core::QueryConfig::kSpillEnabled, true) + .config(core::QueryConfig::kJoinSpillEnabled, true) + .queryCtx( + newQueryCtx(memoryManager, executor_.get(), kMemoryCapacity)) + .maxDrivers(numDrivers) + .plan(PlanBuilder(planNodeIdGenerator) + .values(probeInput, true) + .hashJoin( + {"p0"}, + {"b0"}, + PlanBuilder(planNodeIdGenerator) + .values(buildInput, true) + .planNode(), + "", + {"p0", "p1", "b0", "b1"}, + core::JoinType::kInner) + .planNode()) + .assertResults( + "SELECT probe.p0, probe.p1, build.b0, build.b1 FROM probe " + "INNER JOIN build ON probe.p0 = build.b0"), + "Simulated failure"); + // This test uses on-demand created memory manager instead of the global + // one. We need to make sure any used memory got cleaned up before exiting + // the scope + waitForAllTasksToBeDeleted(); + ASSERT_EQ(arbitrator->stats().freeCapacityBytes, expectedFreeCapacityBytes); +} + +DEBUG_ONLY_TEST_F(HashJoinTest, arbitrationTriggeredDuringParallelJoinBuild) { + std::unique_ptr memoryManager = createMemoryManager(); + const uint64_t numDrivers = 2; + + // Large build side key product to bump hash mode to kHash instead of kArray + // to trigger parallel join build. + const uint64_t numBuildSideRows = 500; + auto buildKeyVector = makeFlatVector( + numBuildSideRows, + [](vector_size_t row) { return folly::Random::rand64(); }); + auto buildSideVector = makeRowVector( + {"b0", "b1", "b2"}, {buildKeyVector, buildKeyVector, buildKeyVector}); + std::vector buildSideVectors; + for (int i = 0; i < numDrivers; ++i) { + buildSideVectors.push_back(buildSideVector); + } + createDuckDbTable("build", buildSideVectors); + + const uint64_t numProbeSideRows = 10; + auto probeKeyVector = makeFlatVector( + numProbeSideRows, + [&](vector_size_t row) { return buildKeyVector->valueAt(row); }); + auto probeSideVector = makeRowVector( + {"p0", "p1", "p2"}, {probeKeyVector, probeKeyVector, probeKeyVector}); + std::vector probeSideVectors; + for (int i = 0; i < numDrivers; ++i) { + probeSideVectors.push_back(probeSideVector); + } + createDuckDbTable("probe", probeSideVectors); + + std::shared_ptr joinQueryCtx = + newQueryCtx(memoryManager.get(), executor_.get(), kMemoryCapacity); + + const int64_t allocSize = 512LL << 20; + std::atomic parallelBuildTriggered{false}; + std::atomic joinBuildPool{nullptr}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::HashTable::parallelJoinBuild", + std::function([&](memory::MemoryPool* pool) { + parallelBuildTriggered = true; + // Pick the last running driver threads' pool for later memory + // allocation. This pick is rather arbitrary, as it is un-important + // which pool is going to be allocated from later in a parallel join's + // off-driver thread. + joinBuildPool = pool; + })); + + std::atomic_bool offThreadAllocationTriggered{false}; + folly::EventCount asyncMoveWait; + std::atomic asyncMoveWaitFlag{true}; + SCOPED_TESTVALUE_SET( + "facebook::velox::AsyncSource::prepare", + std::function([&](void* /* unused */) { + if (!offThreadAllocationTriggered.exchange(true)) { + SCOPE_EXIT { + asyncMoveWaitFlag = false; + asyncMoveWait.notifyAll(); + }; + // Executed by the first thread hitting the test value location. This + // allocation will trigger arbitration and fail. + VELOX_ASSERT_THROW( + joinBuildPool.load()->allocate(allocSize), + "Exceeded memory pool cap"); + } + })); + + // Wait for allocation (hence arbitration) on the prepare thread to finish + // before calling AsyncSource::move(). This is to ensure no other AsyncSource + // (hence arbitration) is running on the driver thread (on-thread) before the + // ongoing arbitration finishes. Without ensuring this, the on-thread + // arbitration (triggered by calling AsyncSource::move() first) has + // thread-local driver context by default, defying the purpose of this test. + SCOPED_TESTVALUE_SET( + "facebook::velox::AsyncSource::move", + std::function([&](void* /* unused */) { + asyncMoveWait.await([&]() { return !asyncMoveWaitFlag.load(); }); + })); + + std::vector probeInput = {probeSideVector}; + std::vector buildInput = {buildSideVector}; + auto planNodeIdGenerator = std::make_shared(); + const auto spillDirectory = exec::test::TempDirectoryPath::create(); + AssertQueryBuilder(duckDbQueryRunner_) + .spillDirectory(spillDirectory->getPath()) + .config(core::QueryConfig::kSpillEnabled, true) + .config(core::QueryConfig::kJoinSpillEnabled, true) + // Set very low table size threshold to trigger parallel build. + .config(core::QueryConfig::kMinTableRowsForParallelJoinBuild, 0) + // Set multiple hash build drivers to trigger parallel build. + .maxDrivers(numDrivers) + .queryCtx(joinQueryCtx) + .plan(PlanBuilder(planNodeIdGenerator) + .values(probeInput, true) + .hashJoin( + {"p0", "p1", "p2"}, + {"b0", "b1", "b2"}, + PlanBuilder(planNodeIdGenerator) + .values(buildInput, true) + .planNode(), + "", + {"p0", "p1", "b0", "b1"}, + core::JoinType::kInner) + .planNode()) + .assertResults( + "SELECT probe.p0, probe.p1, build.b0, build.b1 FROM probe " + "INNER JOIN build ON probe.p0 = build.b0 AND probe.p1 = build.b1 AND " + "probe.p2 = build.b2"); + ASSERT_TRUE(parallelBuildTriggered); + + // This test uses on-demand created memory manager instead of the global + // one. We need to make sure any used memory got cleaned up before exiting + // the scope + waitForAllTasksToBeDeleted(); +} + +DEBUG_ONLY_TEST_F(HashJoinTest, arbitrationTriggeredByEnsureJoinTableFit) { + // Use manual spill injection other than spill injection framework. This is + // because spill injection framework does not allow fine grain spill within a + // single operator (We do not want to spill during addInput() but only during + // finishHashBuild()). + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::HashBuild::ensureTableFits", + std::function(([&](Operator* op) { + Operator::ReclaimableSectionGuard guard(op); + memory::testingRunArbitration(op->pool()); + }))); + auto tempDirectory = exec::test::TempDirectoryPath::create(); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers_) + .injectSpill(false) + .spillDirectory(tempDirectory->getPath()) + .keyTypes({BIGINT()}) + .probeVectors(1600, 5) + .buildVectors(1500, 5) + .referenceQuery( + "SELECT t_k0, t_data, u_k0, u_data FROM t, u WHERE t.t_k0 = u.u_k0") + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + const auto statsPair = taskSpilledStats(*task); + ASSERT_GT(statsPair.first.spilledBytes, 0); + }) + .run(); +} + +DEBUG_ONLY_TEST_F(HashJoinTest, joinBuildSpillError) { + const int kMemoryCapacity = 32 << 20; + // Set a small memory capacity to trigger spill. + std::unique_ptr memoryManager = + createMemoryManager(kMemoryCapacity, 0); + const auto& arbitrator = memoryManager->arbitrator(); + auto rowType = ROW( + {{"c0", INTEGER()}, + {"c1", INTEGER()}, + {"c2", VARCHAR()}, + {"c3", VARCHAR()}}); + + std::vector vectors = createVectors(16, rowType, fuzzerOpts_); + createDuckDbTable(vectors); + + std::shared_ptr joinQueryCtx = + newQueryCtx(memoryManager.get(), executor_.get(), kMemoryCapacity); + + const int numDrivers = 4; + std::atomic numAppends{0}; + const std::string injectedErrorMsg("injected spillError"); + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::SpillState::appendToPartition", + std::function([&](exec::SpillState* state) { + if (++numAppends != numDrivers) { + return; + } + VELOX_FAIL(injectedErrorMsg); + })); + + auto planNodeIdGenerator = std::make_shared(); + const auto spillDirectory = exec::test::TempDirectoryPath::create(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(vectors) + .project({"c0 AS t0", "c1 AS t1", "c2 AS t2"}) + .hashJoin( + {"t0"}, + {"u0"}, + PlanBuilder(planNodeIdGenerator) + .values(vectors) + .project({"c0 AS u0", "c1 AS u1", "c2 AS u2"}) + .planNode(), + "", + {"t1"}, + core::JoinType::kAnti) + .planNode(); + VELOX_ASSERT_THROW( + AssertQueryBuilder(plan) + .queryCtx(joinQueryCtx) + .spillDirectory(spillDirectory->getPath()) + .config(core::QueryConfig::kSpillEnabled, true) + .copyResults(pool()), + injectedErrorMsg); + + waitForAllTasksToBeDeleted(); + ASSERT_EQ(arbitrator->stats().numFailures, 1); + + // Wait again here as this test uses on-demand created memory manager instead + // of the global one. We need to make sure any used memory got cleaned up + // before exiting the scope + waitForAllTasksToBeDeleted(); +} + +DEBUG_ONLY_TEST_F(HashJoinTest, probeSpillOnWaitForPeers) { + // This test creates a scenario when tester probe thread finishes processing + // input, entering kWaitForPeers state, and the other thread is still + // processing, spill is triggered properly performed. + + folly::EventCount startWait; + folly::Synchronized testerOpName; + std::atomic_bool injectedSpillOnce{false}; + + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::getOutput", + std::function([&](Operator* op) { + if (!isHashProbeMemoryPool(*op->pool())) { + return; + } + testerOpName.withWLock([&](std::string& opName) { + if (opName.empty()) { + opName = op->pool()->name(); + } + }); + if (op->pool()->name() == *testerOpName.rlock()) { + // Do not block tester thread. + return; + } + startWait.await([&]() { return injectedSpillOnce.load(); }); + })); + + // tester probe operator is guaranteed to be in kWaitForPeers state the next + // isBlocked() is called after noMoreInput() is called. + std::atomic_bool noMoreInputCalled{false}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::noMoreInput", + std::function([&](Operator* op) { + if (!isHashProbeMemoryPool(*op->pool())) { + return; + } + noMoreInputCalled = true; + })); + + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::isBlocked", + std::function([&](Operator* op) { + if (!isHashProbeMemoryPool(*op->pool())) { + return; + } + if (injectedSpillOnce || !noMoreInputCalled) { + return; + } + injectedSpillOnce = true; + EXPECT_EQ( + dynamic_cast(op)->testingState(), + ProbeOperatorState::kWaitForPeers); + testingRunArbitration(op->pool()); + })); + + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Task::requestPauseLocked", + std::function([&](Task* task) { startWait.notifyAll(); })); + + const uint64_t numDrivers{2}; + std::shared_ptr joinQueryCtx = + newQueryCtx(memory::memoryManager(), executor_.get(), kMemoryCapacity); + auto rowType = ROW({{"c0", INTEGER()}, {"c1", INTEGER()}}); + fuzzerOpts_.vectorSize = 20; + std::vector vectors = createVectors(6, rowType, fuzzerOpts_); + std::vector totalVectors; + for (auto i = 0; i < numDrivers; ++i) { + totalVectors.insert(totalVectors.end(), vectors.begin(), vectors.end()); + } + createDuckDbTable(totalVectors); + auto spillDirectory = exec::test::TempDirectoryPath::create(); + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(vectors, true) + .project({"c0 AS t0", "c1 AS t1"}) + .hashJoin( + {"t0"}, + {"u0"}, + PlanBuilder(planNodeIdGenerator) + .values(vectors, true) + .project({"c0 AS u0", "c1 AS u1"}) + .planNode(), + "", + {"t1"}, + core::JoinType::kInner) + .planNode(); + + { + auto task = + AssertQueryBuilder(duckDbQueryRunner_) + .plan(plan) + .queryCtx(joinQueryCtx) + .spillDirectory(spillDirectory->getPath()) + .config(core::QueryConfig::kSpillEnabled, true) + .maxDrivers(numDrivers) + .assertResults("SELECT a.c1 from tmp a join tmp b on a.c0 = b.c0"); + + auto opStats = toOperatorStats(task->taskStats()); + ASSERT_GT(opStats.at("HashProbe").spilledBytes, 0); + ASSERT_EQ(opStats.at("HashBuild").spilledBytes, 0); + + const auto* arbitrator = memory::memoryManager()->arbitrator(); + ASSERT_GT(arbitrator->stats().reclaimedUsedBytes, 0); + } + waitForAllTasksToBeDeleted(); +} + +DEBUG_ONLY_TEST_F(HashJoinTest, taskWaitTimeout) { + const int queryMemoryCapacity = 128 << 20; + // Creates a large number of vectors based on the query capacity to trigger + // memory arbitration. + fuzzerOpts_.vectorSize = 10'000; + auto rowType = ROW( + {{"c0", INTEGER()}, + {"c1", INTEGER()}, + {"c2", VARCHAR()}, + {"c3", VARCHAR()}}); + const auto vectors = + createVectors(rowType, queryMemoryCapacity / 2, fuzzerOpts_); + const int numDrivers = 4; + const auto expectedResult = + runHashJoinTask(vectors, nullptr, false, numDrivers, pool(), false).data; + + for (uint64_t timeoutMs : {1'000, 30'000}) { + SCOPED_TRACE(fmt::format("timeout {}", succinctMillis(timeoutMs))); + auto memoryManager = createMemoryManager(512 << 20, 0, timeoutMs); + auto queryCtx = + newQueryCtx(memoryManager.get(), executor_.get(), queryMemoryCapacity); + + // Set test injection to block one hash build operator to inject delay when + // memory reclaim waits for task to pause. + folly::EventCount buildBlockWait; + std::atomic buildBlockWaitFlag{true}; + std::atomic blockOneBuild{true}; + SCOPED_TESTVALUE_SET( + "facebook::velox::common::memory::MemoryPoolImpl::maybeReserve", + std::function([&](memory::MemoryPool* pool) { + const std::string re(".*HashBuild"); + if (!RE2::FullMatch(pool->name(), re)) { + return; + } + if (!blockOneBuild.exchange(false)) { + return; + } + buildBlockWait.await([&]() { return !buildBlockWaitFlag.load(); }); + })); + + folly::EventCount taskPauseWait; + std::atomic taskPauseWaitFlag{false}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Task::requestPauseLocked", + std::function(([&](Task* /*unused*/) { + taskPauseWaitFlag = true; + taskPauseWait.notifyAll(); + }))); + + std::thread queryThread([&]() { + // We expect failure on short time out. + if (timeoutMs == 1'000) { + VELOX_ASSERT_THROW( + runHashJoinTask( + vectors, + queryCtx, + false, + numDrivers, + pool(), + true, + expectedResult), + "Memory reclaim failed to wait"); + } else { + // We expect succeed on large time out or no timeout. + const auto result = runHashJoinTask( + vectors, queryCtx, false, numDrivers, pool(), true, expectedResult); + auto taskStats = exec::toPlanStats(result.task->taskStats()); + auto& planStats = taskStats.at(result.planNodeId); + ASSERT_GT(planStats.spilledBytes, 0); + } + }); + + // Wait for task pause to reach, and then delay for a while before unblock + // the blocked hash build operator. + taskPauseWait.await([&]() { return taskPauseWaitFlag.load(); }); + // Wait for two seconds and expect the short reclaim wait timeout. + std::this_thread::sleep_for(std::chrono::seconds(2)); + // Unblock the blocked build operator to let memory reclaim proceed. + buildBlockWaitFlag = false; + buildBlockWait.notifyAll(); + + queryThread.join(); + + // This test uses on-demand created memory manager instead of the global + // one. We need to make sure any used memory got cleaned up before exiting + // the scope + waitForAllTasksToBeDeleted(); + } +} + +DEBUG_ONLY_TEST_F(HashJoinTest, hashProbeSpill) { + struct { + bool triggerBuildSpill; + // Triggers after no more input or not. + bool afterNoMoreInput; + // The index of get output call to trigger probe side spilling. + int probeOutputIndex; + + std::string debugString() const { + return fmt::format( + "triggerBuildSpill: {}, afterNoMoreInput: {}, probeOutputIndex: {}", + triggerBuildSpill, + afterNoMoreInput, + probeOutputIndex); + } + } testSettings[] = { + {false, false, 0}, + {false, false, 1}, + {false, false, 10}, + {false, true, 0}, + {true, false, 0}, + {true, false, 1}, + {true, false, 10}, + {true, true, 0}}; + + for (const auto& testData : testSettings) { + SCOPED_TRACE(testData.debugString()); + + std::atomic_bool injectBuildSpillOnce{true}; + std::atomic_int buildInputCount{0}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::addInput", + std::function([&](Operator* op) { + if (!testData.triggerBuildSpill) { + return; + } + if (!isHashBuildMemoryPool(*op->pool())) { + return; + } + if (buildInputCount++ != 1) { + return; + } + if (!injectBuildSpillOnce.exchange(false)) { + return; + } + testingRunArbitration(op->pool()); + })); + + std::atomic_bool injectProbeSpillOnce{true}; + std::atomic_int probeOutputCount{0}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::getOutput", + std::function([&](Operator* op) { + if (!isHashProbeMemoryPool(*op->pool())) { + return; + } + if (testData.afterNoMoreInput) { + if (!op->testingNoMoreInput()) { + return; + } + } else { + if (probeOutputCount++ != testData.probeOutputIndex) { + return; + } + } + if (!injectProbeSpillOnce.exchange(false)) { + return; + } + testingRunArbitration(op->pool()); + })); + + fuzzerOpts_.vectorSize = 128; + auto probeVectors = createVectors(10, probeType_, fuzzerOpts_); + auto buildVectors = createVectors(20, buildType_, fuzzerOpts_); + const auto spillDirectory = exec::test::TempDirectoryPath::create(); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(1) + .spillDirectory(spillDirectory->getPath()) + .probeKeys({"t_k1"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u_k1"}) + .buildVectors(std::move(buildVectors)) + .config(core::QueryConfig::kJoinSpillEnabled, "true") + .joinType(core::JoinType::kRight) + .joinOutputLayout({"t_k1", "t_k2", "u_k1", "t_v1"}) + .referenceQuery( + "SELECT t.t_k1, t.t_k2, u.u_k1, t.t_v1 FROM t RIGHT JOIN u ON t.t_k1 = u.u_k1") + .injectSpill(false) + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + auto opStats = toOperatorStats(task->taskStats()); + ASSERT_GT(opStats.at("HashProbe").spilledBytes, 0); + if (testData.triggerBuildSpill) { + ASSERT_GT(opStats.at("HashBuild").spilledBytes, 0); + } else { + ASSERT_EQ(opStats.at("HashBuild").spilledBytes, 0); + } + + const auto* arbitrator = memory::memoryManager()->arbitrator(); + ASSERT_GT(arbitrator->stats().numRequests, 0); + ASSERT_GT(arbitrator->stats().reclaimedUsedBytes, 0); + }) + .run(); + } +} + +DEBUG_ONLY_TEST_F(HashJoinTest, hashProbeSpillInMiddeOfLastOutputProcessing) { + std::atomic_int outputCountAfterNoMoreInout{0}; + std::atomic_bool injectOnce{true}; + ::facebook::velox::common::testutil::ScopedTestValue abc( + "facebook::velox::exec::Driver::runInternal::getOutput", + std::function([&](Operator* op) { + if (!isHashProbeMemoryPool(*op->pool())) { + return; + } + if (!op->testingNoMoreInput()) { + return; + } + if (outputCountAfterNoMoreInout++ != 1) { + return; + } + if (!injectOnce.exchange(false)) { + return; + } + testingRunArbitration(op->pool()); + })); + + fuzzerOpts_.vectorSize = 128; + auto probeVectors = createVectors(10, probeType_, fuzzerOpts_); + auto buildVectors = createVectors(20, buildType_, fuzzerOpts_); + + const auto spillDirectory = exec::test::TempDirectoryPath::create(); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(1) + .spillDirectory(spillDirectory->getPath()) + .probeKeys({"t_k1"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u_k1"}) + .buildVectors(std::move(buildVectors)) + .config(core::QueryConfig::kJoinSpillEnabled, "true") + .config(core::QueryConfig::kPreferredOutputBatchRows, std::to_string(10)) + .joinType(core::JoinType::kRight) + .joinOutputLayout({"t_k1", "t_k2", "u_k1", "t_v1"}) + .referenceQuery( + "SELECT t.t_k1, t.t_k2, u.u_k1, t.t_v1 FROM t RIGHT JOIN u ON t.t_k1 = u.u_k1") + .injectSpill(false) + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + auto opStats = toOperatorStats(task->taskStats()); + ASSERT_GT(opStats.at("HashProbe").spilledBytes, 0); + // Verifies that we only spill the output which is single partitioned + // but not the hash table. + ASSERT_EQ(opStats.at("HashProbe").spilledPartitions, 1); + }) + .run(); +} + +// Inject probe-side spilling in the middle of output processing. If +// 'recursiveSpill' is true, we trigger probe-spilling when probe the hash table +// built from spilled data. +DEBUG_ONLY_TEST_F(HashJoinTest, hashProbeSpillInMiddeOfOutputProcessing) { + for (bool recursiveSpill : {false, true}) { + std::atomic_int buildInputCount{0}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::addInput", + std::function([&](Operator* op) { + if (!isHashBuildMemoryPool(*op->pool())) { + return; + } + if (!recursiveSpill) { + return; + } + // Trigger spill after the build side has processed some rows. + if (buildInputCount++ != 1) { + return; + } + testingRunArbitration(op->pool()); + })); + + std::atomic_bool injectProbeSpillOnce{true}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::getOutput", + std::function([&](Operator* op) { + if (!isHashProbeMemoryPool(*op->pool())) { + return; + } + + if (op->testingHasInput()) { + return; + } + if (recursiveSpill) { + if (static_cast(op)->testingHasInputSpiller()) { + return; + } + } + if (!injectProbeSpillOnce.exchange(false)) { + return; + } + testingRunArbitration(op->pool()); + })); + + fuzzerOpts_.vectorSize = 128; + auto probeVectors = createVectors(10, probeType_, fuzzerOpts_); + auto buildVectors = createVectors(20, buildType_, fuzzerOpts_); + + const auto spillDirectory = exec::test::TempDirectoryPath::create(); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(1) + .spillDirectory(spillDirectory->getPath()) + .probeKeys({"t_k1"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u_k1"}) + .buildVectors(std::move(buildVectors)) + .config(core::QueryConfig::kJoinSpillEnabled, "true") + .config( + core::QueryConfig::kPreferredOutputBatchRows, std::to_string(10)) + .joinType(core::JoinType::kRight) + .joinOutputLayout({"t_k1", "t_k2", "u_k1", "t_v1"}) + .referenceQuery( + "SELECT t.t_k1, t.t_k2, u.u_k1, t.t_v1 FROM t RIGHT JOIN u ON t.t_k1 = u.u_k1") + .injectSpill(false) + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + auto opStats = toOperatorStats(task->taskStats()); + ASSERT_GT(opStats.at("HashProbe").spilledBytes, 0); + ASSERT_GT(opStats.at("HashProbe").spilledPartitions, 1); + }) + .run(); + } +} + +DEBUG_ONLY_TEST_F(HashJoinTest, hashProbeSpillWhenOneOfProbeFinish) { + const int numDrivers{3}; + + std::atomic_bool probeWaitFlag{true}; + folly::EventCount probeWait; + std::atomic_int numBlockedProbeOps{0}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::getOutput", + std::function([&](Operator* op) { + if (!isHashProbeMemoryPool(*op->pool())) { + return; + } + if (++numBlockedProbeOps <= numDrivers - 1) { + probeWait.await([&]() { return !probeWaitFlag.load(); }); + return; + } + })); + + std::atomic_bool notifyOnce{true}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::noMoreInput", + std::function([&](Operator* op) { + if (!isHashProbeMemoryPool(*op->pool())) { + return; + } + if (!notifyOnce.exchange(false)) { + return; + } + probeWaitFlag = false; + probeWait.notifyAll(); + })); + + std::thread queryThread([&]() { + const auto spillDirectory = exec::test::TempDirectoryPath::create(); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers, true, true) + .spillDirectory(spillDirectory->getPath()) + .keyTypes({BIGINT()}) + .probeVectors(32, 5) + .buildVectors(32, 5) + .config(core::QueryConfig::kJoinSpillEnabled, "true") + .referenceQuery( + "SELECT t_k0, t_data, u_k0, u_data FROM t, u WHERE t.t_k0 = u.u_k0") + .injectSpill(false) + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + auto opStats = toOperatorStats(task->taskStats()); + ASSERT_EQ(opStats.at("HashBuild").spilledBytes, 0); + ASSERT_GT(opStats.at("HashProbe").spilledBytes, 0); + }) + .run(); + }); + // Wait until one of the hash probe operator has finished. + probeWait.await([&]() { return !probeWaitFlag.load(); }); + memory::testingRunArbitration(); + queryThread.join(); +} + +DEBUG_ONLY_TEST_F(HashJoinTest, hashProbeSpillExceedLimit) { + // If 'buildTriggerSpill' is true, then spilling is triggered by hash build. + for (const bool buildTriggerSpill : {false, true}) { + SCOPED_TRACE(fmt::format("buildTriggerSpill {}", buildTriggerSpill)); + + SCOPED_TESTVALUE_SET( + "facebook::velox::common::memory::MemoryPoolImpl::maybeReserve", + std::function([&](memory::MemoryPool* pool) { + if (buildTriggerSpill && !isHashBuildMemoryPool(*pool)) { + return; + } + if (!buildTriggerSpill && !isHashProbeMemoryPool(*pool)) { + return; + } + testingRunArbitration(pool); + })); + + fuzzerOpts_.vectorSize = 128; + auto probeVectors = createVectors(32, probeType_, fuzzerOpts_); + auto buildVectors = createVectors(64, buildType_, fuzzerOpts_); + for (int i = 0; i < probeVectors.size(); ++i) { + const auto probeKeyChannel = probeType_->getChildIdx("t_k1"); + const auto buildKeyChannle = buildType_->getChildIdx("u_k1"); + probeVectors[i]->childAt(probeKeyChannel) = + buildVectors[i]->childAt(buildKeyChannle); + } + + const auto spillDirectory = exec::test::TempDirectoryPath::create(); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(1) + .spillDirectory(spillDirectory->getPath()) + .probeKeys({"t_k1"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u_k1"}) + .buildVectors(std::move(buildVectors)) + .config(core::QueryConfig::kMaxSpillLevel, "1") + .config(core::QueryConfig::kSpillNumPartitionBits, "1") + .config(core::QueryConfig::kJoinSpillEnabled, "true") + // Set small write buffer size to have small vectors to read from + // spilled data. + .config(core::QueryConfig::kSpillWriteBufferSize, "1") + .config( + core::QueryConfig::kPreferredOutputBatchRows, std::to_string(10)) + .joinType(core::JoinType::kRight) + .joinOutputLayout({"t_k1", "t_k2", "u_k1", "t_v1"}) + .referenceQuery( + "SELECT t.t_k1, t.t_k2, u.u_k1, t.t_v1 FROM t RIGHT JOIN u ON t.t_k1 = u.u_k1") + .injectSpill(false) + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + auto opStats = toOperatorStats(task->taskStats()); + if (buildTriggerSpill) { + ASSERT_GT(opStats.at("HashProbe").spilledBytes, 0); + ASSERT_GT(opStats.at("HashBuild").spilledBytes, 0); + } else { + ASSERT_GT(opStats.at("HashProbe").spilledBytes, 0); + ASSERT_EQ(opStats.at("HashBuild").spilledBytes, 0); + } + ASSERT_GT( + opStats.at("HashProbe") + .runtimeStats[Operator::kExceededMaxSpillLevel] + .sum, + 0); + ASSERT_GT( + opStats.at("HashBuild") + .runtimeStats[Operator::kExceededMaxSpillLevel] + .sum, + 0); + }) + .run(); + } +} + +DEBUG_ONLY_TEST_F(HashJoinTest, hashProbeSpillUnderNonReclaimableSection) { + std::atomic_bool injectOnce{true}; + SCOPED_TESTVALUE_SET( + "facebook::velox::common::memory::MemoryPoolImpl::allocateNonContiguous", + std::function([&](memory::MemoryPool* pool) { + if (!isHashProbeMemoryPool(*pool)) { + return; + } + if (!injectOnce.exchange(false)) { + return; + } + auto* arbitrator = memory::memoryManager()->arbitrator(); + const auto numNonReclaimableAttempts = + arbitrator->stats().numNonReclaimableAttempts; + testingRunArbitration(pool); + // Verifies that we run into non-reclaimable section when reclaim from + // hash probe. + ASSERT_EQ( + arbitrator->stats().numNonReclaimableAttempts, + numNonReclaimableAttempts + 1); + })); + + const auto spillDirectory = exec::test::TempDirectoryPath::create(); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(1) + .spillDirectory(spillDirectory->getPath()) + .keyTypes({BIGINT()}) + .probeVectors(32, 5) + .buildVectors(32, 5) + .config(core::QueryConfig::kJoinSpillEnabled, "true") + .referenceQuery( + "SELECT t_k0, t_data, u_k0, u_data FROM t, u WHERE t.t_k0 = u.u_k0") + .injectSpill(false) + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + auto opStats = toOperatorStats(task->taskStats()); + ASSERT_EQ(opStats.at("HashProbe").spilledBytes, 0); + ASSERT_EQ(opStats.at("HashBuild").spilledBytes, 0); + }) + .run(); +} + +// This test case is to cover the case that hash probe trigger spill for right +// semi join types and the pending input needs to be processed in multiple +// steps. +DEBUG_ONLY_TEST_F(HashJoinTest, spillOutputWithRightSemiJoins) { + for (const auto joinType : + {core::JoinType::kRightSemiFilter, core::JoinType::kRightSemiProject}) { + std::atomic_bool injectOnce{true}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::getOutput", + std::function([&](Operator* op) { + if (op->operatorCtx()->operatorType() != "HashProbe") { + return; + } + if (!op->testingHasInput()) { + return; + } + if (!injectOnce.exchange(false)) { + return; + } + testingRunArbitration(op->pool()); + })); + + std::string duckDbSqlReference; + std::vector joinOutputLayout; + bool nullAware{false}; + if (joinType == core::JoinType::kRightSemiProject) { + duckDbSqlReference = "SELECT u_k2, u_k1 IN (SELECT t_k1 FROM t) FROM u"; + joinOutputLayout = {"u_k2", "match"}; + // Null aware is only supported for semi projection join type. + nullAware = true; + } else { + duckDbSqlReference = + "SELECT u_k2 FROM u WHERE u_k1 IN (SELECT t_k1 FROM t)"; + joinOutputLayout = {"u_k2"}; + } + + const auto spillDirectory = exec::test::TempDirectoryPath::create(); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(1) + .spillDirectory(spillDirectory->getPath()) + .probeType(probeType_) + .probeVectors(128, 3) + .probeKeys({"t_k1"}) + .buildType(buildType_) + .buildVectors(128, 4) + .buildKeys({"u_k1"}) + .joinType(joinType) + // Set a small number of output rows to process the input in multiple + // steps. + .config( + core::QueryConfig::kPreferredOutputBatchRows, std::to_string(10)) + .injectSpill(false) + .joinOutputLayout(std::move(joinOutputLayout)) + .nullAware(nullAware) + .referenceQuery(duckDbSqlReference) + .run(); + } +} + +DEBUG_ONLY_TEST_F(HashJoinTest, spillCheckOnLeftSemiFilterWithDynamicFilters) { + const int32_t numSplits = 10; + const int32_t numRowsProbe = 333; + const int32_t numRowsBuild = 100; + + std::vector probeVectors; + probeVectors.reserve(numSplits); + + std::vector> tempFiles; + for (int32_t i = 0; i < numSplits; ++i) { + auto rowVector = makeRowVector({ + makeFlatVector( + numRowsProbe, [&](auto row) { return row - i * 10; }), + makeFlatVector(numRowsProbe, [](auto row) { return row; }), + }); + probeVectors.push_back(rowVector); + tempFiles.push_back(TempFilePath::create()); + writeToFile(tempFiles.back()->getPath(), rowVector); + } + auto makeInputSplits = [&](const core::PlanNodeId& nodeId) { + return [&] { + std::vector probeSplits; + for (auto& file : tempFiles) { + probeSplits.push_back( + exec::Split(makeHiveConnectorSplit(file->getPath()))); + } + SplitInput splits; + splits.emplace(nodeId, probeSplits); + return splits; + }; + }; + + // 100 key values in [35, 233] range. + std::vector buildVectors; + for (int i = 0; i < 5; ++i) { + buildVectors.push_back(makeRowVector({ + makeFlatVector( + numRowsBuild / 5, + [i](auto row) { return 35 + 2 * (row + i * numRowsBuild / 5); }), + makeFlatVector(numRowsBuild / 5, [](auto row) { return row; }), + })); + } + std::vector keyOnlyBuildVectors; + for (int i = 0; i < 5; ++i) { + keyOnlyBuildVectors.push_back( + makeRowVector({makeFlatVector(numRowsBuild / 5, [i](auto row) { + return 35 + 2 * (row + i * numRowsBuild / 5); + })})); + } + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + auto probeType = ROW({"c0", "c1"}, {INTEGER(), BIGINT()}); + + auto planNodeIdGenerator = std::make_shared(); + + auto buildSide = PlanBuilder(planNodeIdGenerator, pool_.get()) + .values(buildVectors) + .project({"c0 AS u_c0", "c1 AS u_c1"}) + .planNode(); + + // Left semi join. + core::PlanNodeId probeScanId; + core::PlanNodeId joinNodeId; + const auto op = PlanBuilder(planNodeIdGenerator, pool_.get()) + .tableScan(probeType) + .capturePlanNodeId(probeScanId) + .hashJoin( + {"c0"}, + {"u_c0"}, + buildSide, + "", + {"c0", "c1"}, + core::JoinType::kLeftSemiFilter) + .capturePlanNodeId(joinNodeId) + .project({"c0", "c1 + 1"}) + .planNode(); + + std::atomic_bool injectOnce{true}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::getOutput", + std::function([&](Operator* op) { + if (op->operatorCtx()->operatorType() != "HashProbe") { + return; + } + if (!op->testingHasInput()) { + return; + } + if (!injectOnce.exchange(false)) { + return; + } + testingRunArbitration(op->pool()); + })); + + auto spillDirectory = exec::test::TempDirectoryPath::create(); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .planNode(std::move(op)) + .makeInputSplits(makeInputSplits(probeScanId)) + .spillDirectory(spillDirectory->getPath()) + .injectSpill(false) + .referenceQuery( + "SELECT t.c0, t.c1 + 1 FROM t WHERE t.c0 IN (SELECT c0 FROM u)") + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + // Verify spill hasn't triggered. + auto taskStats = exec::toPlanStats(task->taskStats()); + auto& planStats = taskStats.at(joinNodeId); + ASSERT_GT(planStats.spilledBytes, 0); + }) + .run(); +} + +// This test is to verify there is no memory reservation made before hash probe +// start processing. This can cause unnecessary spill and query OOM under some +// real workload with many stages as each hash probe might reserve non-trivial +// amount of memory. +DEBUG_ONLY_TEST_F( + HashJoinTest, + hashProbeMemoryReservationCheckBeforeProbeStartWithSpillEnabled) { + fuzzerOpts_.vectorSize = 128; + auto probeVectors = createVectors(10, probeType_, fuzzerOpts_); + auto buildVectors = createVectors(20, buildType_, fuzzerOpts_); + + std::atomic_bool checkOnce{true}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::addInput", + std::function(([&](Operator* op) { + if (op->operatorType() != "HashProbe") { + return; + } + if (!checkOnce.exchange(false)) { + return; + } + ASSERT_EQ(op->pool()->usedBytes(), 0); + ASSERT_EQ(op->pool()->reservedBytes(), 0); + }))); + + const auto spillDirectory = exec::test::TempDirectoryPath::create(); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(1) + .spillDirectory(spillDirectory->getPath()) + .probeKeys({"t_k1"}) + .probeVectors(std::move(probeVectors)) + .buildKeys({"u_k1"}) + .buildVectors(std::move(buildVectors)) + .config(core::QueryConfig::kJoinSpillEnabled, "true") + .joinType(core::JoinType::kInner) + .joinOutputLayout({"t_k1", "t_k2", "u_k1", "t_v1"}) + .referenceQuery( + "SELECT t.t_k1, t.t_k2, u.u_k1, t.t_v1 FROM t JOIN u ON t.t_k1 = u.u_k1") + .injectSpill(true) + .verifier([&](const std::shared_ptr& task, bool injectSpill) { + if (!injectSpill) { + return; + } + auto opStats = toOperatorStats(task->taskStats()); + ASSERT_GT(opStats.at("HashProbe").spilledBytes, 0); + ASSERT_GE(opStats.at("HashProbe").spilledPartitions, 1); + }) + .run(); +} + +TEST_F(HashJoinTest, nanKeys) { + // Verify the NaN values with different binary representations are considered + // equal. + static const double kNan = std::numeric_limits::quiet_NaN(); + static const double kSNaN = std::numeric_limits::signaling_NaN(); + auto probeInput = makeRowVector( + {makeFlatVector({kNan, kSNaN}), makeFlatVector({1, 2})}); + auto buildInput = makeRowVector({makeFlatVector({kNan, 1})}); + + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values({probeInput}) + .project({"c0 AS t0", "c1 AS t1"}) + .hashJoin( + {"t0"}, + {"u0"}, + PlanBuilder(planNodeIdGenerator) + .values({buildInput}) + .project({"c0 AS u0"}) + .planNode(), + "", + {"t0", "u0", "t1"}, + core::JoinType::kLeft) + .planNode(); + auto queryCtx = core::QueryCtx::create(executor_.get()); + auto result = + AssertQueryBuilder(plan).queryCtx(queryCtx).copyResults(pool_.get()); + auto expected = makeRowVector( + {makeFlatVector({kNan, kNan}), + makeFlatVector({kNan, kNan}), + makeFlatVector({1, 2})}); + facebook::velox::test::assertEqualVectors(expected, result); +} + +DEBUG_ONLY_TEST_F(HashJoinTest, spillOnBlockedProbe) { + auto blockedOperatorFactoryUniquePtr = + std::make_unique(); + auto blockedOperatorFactory = blockedOperatorFactoryUniquePtr.get(); + Operator::registerOperator(std::move(blockedOperatorFactoryUniquePtr)); + + std::vector unblockPromises; + std::atomic_bool shouldBlock{true}; + blockedOperatorFactory->setBlockedCb([&](ContinueFuture* future) { + if (!shouldBlock) { + return BlockingReason::kNotBlocked; + } + auto [p, f] = makeVeloxContinuePromiseContract("Blocked Operator"); + *future = std::move(f); + unblockPromises.push_back(std::move(p)); + return BlockingReason::kWaitForConsumer; + }); + + folly::EventCount arbitrationWait; + std::atomic arbitrationWaitFlag{true}; + ::facebook::velox::common::testutil::ScopedTestValue _scopedTestValue15( + "facebook::velox::exec::HashBuild::finishHashBuild", + std::function([&](Operator* /* unused */) { + arbitrationWaitFlag = false; + arbitrationWait.notifyAll(); + })); + std::thread arbitrationThread([&]() { + arbitrationWait.await([&]() { return !arbitrationWaitFlag.load(); }); + memory::memoryManager()->shrinkPools(); + shouldBlock = false; + for (auto& unblockPromise : unblockPromises) { + unblockPromise.setValue(); + } + }); + + auto rowType = ROW({{"c0", INTEGER()}, {"c1", INTEGER()}}); + std::vector vectors = createVectors(1, rowType, fuzzerOpts_); + createDuckDbTable(vectors); + auto planNodeIdGenerator = std::make_shared(); + const auto spillDirectory = exec::test::TempDirectoryPath::create(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(vectors) + .project({"c0 AS t0", "c1 AS t1"}) + .hashJoin( + {"t0"}, + {"u0"}, + PlanBuilder(planNodeIdGenerator) + .values(vectors) + .project({"c0 AS u0", "c1 AS u1"}) + .planNode(), + "", + {"t1"}, + core::JoinType::kInner) + .addNode([&](std::string id, core::PlanNodePtr input) { + return std::make_shared(id, input); + }) + .planNode(); + + { + auto task = + AssertQueryBuilder(duckDbQueryRunner_) + .plan(plan) + .queryCtx(newQueryCtx( + memory::memoryManager(), executor_.get(), kMemoryCapacity)) + .spillDirectory(spillDirectory->getPath()) + .config(core::QueryConfig::kSpillEnabled, true) + .maxDrivers(1) + .assertResults("SELECT a.c1 from tmp a join tmp b on a.c0 = b.c0"); + auto joinSpillStats = taskSpilledStats(*task); + auto buildSpillStats = joinSpillStats.first; + ASSERT_GT(buildSpillStats.spilledBytes, 0); + } + arbitrationThread.join(); + waitForAllTasksToBeDeleted(30'000'000); +} + +DEBUG_ONLY_TEST_F(HashJoinTest, buildReclaimedMemoryReport) { + constexpr int64_t kMaxBytes = 1LL << 30; // 1GB + const int32_t numBuildVectors = 3; + std::vector buildVectors; + for (int32_t i = 0; i < numBuildVectors; ++i) { + VectorFuzzer fuzzer({.vectorSize = 200}, pool()); + buildVectors.push_back(fuzzer.fuzzRow(buildType_)); + } + + const int32_t numProbeVectors = 3; + std::vector probeVectors; + for (int32_t i = 0; i < numProbeVectors; ++i) { + VectorFuzzer fuzzer({.vectorSize = 200}, pool()); + probeVectors.push_back(fuzzer.fuzzRow(probeType_)); + } + + const int numDrivers{2}; + // duckdb need double probe and build inputs as we run two drivers for hash + // join. + std::vector totalProbeVectors = probeVectors; + totalProbeVectors.insert( + totalProbeVectors.end(), probeVectors.begin(), probeVectors.end()); + std::vector totalBuildVectors = buildVectors; + totalBuildVectors.insert( + totalBuildVectors.end(), buildVectors.begin(), buildVectors.end()); + + createDuckDbTable("t", totalProbeVectors); + createDuckDbTable("u", totalBuildVectors); + + auto tempDirectory = exec::test::TempDirectoryPath::create(); + auto queryPool = memory::memoryManager()->addRootPool( + "", kMaxBytes, memory::MemoryReclaimer::create()); + + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, true) + .hashJoin( + {"t_k1"}, + {"u_k1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, true) + .planNode(), + "", + concat(probeType_->names(), buildType_->names())) + .planNode(); + + folly::EventCount driverWait; + std::atomic_bool driverWaitFlag{true}; + folly::EventCount taskWait; + std::atomic_bool taskWaitFlag{true}; + + Operator* op{nullptr}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::HashBuild::finishHashBuild", + std::function(([&](Operator* testOp) { op = testOp; }))); + + std::atomic_bool injectOnce{true}; + SCOPED_TESTVALUE_SET( + "facebook::velox::common::memory::MemoryPoolImpl::maybeReserve", + std::function( + ([&](memory::MemoryPoolImpl* pool) { + if (op == nullptr || op->pool() != pool) { + return; + } + ASSERT_TRUE(isHashBuildMemoryPool(*pool)); + ASSERT_TRUE(op->canReclaim()); + ASSERT_GT(op->pool()->usedBytes(), 0); + ASSERT_GT( + op->pool()->parent()->reservedBytes(), + op->pool()->reservedBytes()); + if (!injectOnce.exchange(false)) { + return; + } + uint64_t reclaimableBytes{0}; + const bool reclaimable = op->reclaimableBytes(reclaimableBytes); + ASSERT_TRUE(reclaimable); + ASSERT_GT(reclaimableBytes, 0); + auto* driver = op->operatorCtx()->driver(); + TestSuspendedSection suspendedSection(driver); + taskWaitFlag = false; + taskWait.notifyAll(); + driverWait.await([&]() { return !driverWaitFlag.load(); }); + }))); + + std::thread taskThread([&]() { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(numDrivers) + .planNode(plan) + .queryPool(std::move(queryPool)) + .injectSpill(false) + .spillDirectory(tempDirectory->getPath()) + .referenceQuery( + "SELECT t_k1, t_k2, t_v1, u_k1, u_k2, u_v1 FROM t, u WHERE t.t_k1 = u.u_k1") + .config(core::QueryConfig::kSpillStartPartitionBit, "29") + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + const auto statsPair = taskSpilledStats(*task); + ASSERT_GT(statsPair.first.spilledBytes, 0); + ASSERT_EQ(statsPair.first.spilledPartitions, 16); + ASSERT_GT(statsPair.second.spilledBytes, 0); + ASSERT_EQ(statsPair.second.spilledPartitions, 16); + verifyTaskSpilledRuntimeStats(*task, true); + }) + .run(); + }); + + taskWait.await([&]() { return !taskWaitFlag.load(); }); + ASSERT_TRUE(op != nullptr); + auto task = op->operatorCtx()->task(); + auto* nodePool = op->pool()->parent(); + const auto nodeMemoryUsage = nodePool->reservedBytes(); + { + memory::ScopedMemoryArbitrationContext ctx(op->pool()); + const uint64_t reclaimedBytes = task->pool()->reclaim( + task->pool()->capacity(), 1'000'000, reclaimerStats_); + ASSERT_GT(reclaimedBytes, 0); + ASSERT_EQ(nodeMemoryUsage - nodePool->reservedBytes(), reclaimedBytes); + } + // Verify all the memory has been freed. + ASSERT_EQ(nodePool->reservedBytes(), 0); + + driverWaitFlag = false; + driverWait.notifyAll(); + task.reset(); + + taskThread.join(); +} + +DEBUG_ONLY_TEST_F(HashJoinTest, probeReclaimedMemoryReport) { + constexpr int64_t kMaxBytes = 1LL << 30; // 1GB + const int32_t numBuildVectors = 3; + std::vector buildVectors; + for (int32_t i = 0; i < numBuildVectors; ++i) { + VectorFuzzer fuzzer({.vectorSize = 200}, pool()); + buildVectors.push_back(fuzzer.fuzzRow(buildType_)); + } + + const int32_t numProbeVectors = 3; + std::vector probeVectors; + for (int32_t i = 0; i < numProbeVectors; ++i) { + VectorFuzzer fuzzer({.vectorSize = 200}, pool()); + probeVectors.push_back(fuzzer.fuzzRow(probeType_)); + } + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + auto tempDirectory = exec::test::TempDirectoryPath::create(); + auto queryPool = memory::memoryManager()->addRootPool( + "", kMaxBytes, memory::MemoryReclaimer::create()); + + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, true) + .hashJoin( + {"t_k1"}, + {"u_k1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, true) + .planNode(), + "", + concat(probeType_->names(), buildType_->names())) + .planNode(); + + folly::EventCount driverWait; + std::atomic_bool driverWaitFlag{true}; + folly::EventCount taskWait; + std::atomic_bool taskWaitFlag{true}; + + Operator* op{nullptr}; + std::atomic_int probeInputCount{0}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::addInput", + std::function(([&](Operator* testOp) { + if (testOp->operatorType() != "HashProbe") { + return; + } + op = testOp; + + ASSERT_TRUE(op->canReclaim()); + if (probeInputCount++ != 1) { + return; + } + auto* driver = op->operatorCtx()->driver(); + TestSuspendedSection suspendedSection(driver); + taskWaitFlag = false; + taskWait.notifyAll(); + driverWait.await([&]() { return !driverWaitFlag.load(); }); + }))); + + std::thread taskThread([&]() { + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(1) + .planNode(plan) + .queryPool(std::move(queryPool)) + .injectSpill(false) + .spillDirectory(tempDirectory->getPath()) + .referenceQuery( + "SELECT t_k1, t_k2, t_v1, u_k1, u_k2, u_v1 FROM t, u WHERE t.t_k1 = u.u_k1") + .config(core::QueryConfig::kSpillStartPartitionBit, "29") + .verifier([&](const std::shared_ptr& task, bool /*unused*/) { + const auto statsPair = taskSpilledStats(*task); + // The spill triggered at the probe side. + ASSERT_EQ(statsPair.first.spilledBytes, 0); + ASSERT_EQ(statsPair.first.spilledPartitions, 0); + ASSERT_GT(statsPair.second.spilledBytes, 0); + ASSERT_EQ(statsPair.second.spilledPartitions, 16); + }) + .run(); + }); + + taskWait.await([&]() { return !taskWaitFlag.load(); }); + ASSERT_TRUE(op != nullptr); + auto task = op->operatorCtx()->task(); + auto* nodePool = op->pool()->parent(); + const auto nodeMemoryUsage = nodePool->reservedBytes(); + { + memory::ScopedMemoryArbitrationContext ctx(op->pool()); + const uint64_t reclaimedBytes = task->pool()->reclaim( + task->pool()->capacity(), 1'000'000, reclaimerStats_); + ASSERT_GT(reclaimedBytes, 0); + ASSERT_EQ(nodeMemoryUsage - nodePool->reservedBytes(), reclaimedBytes); + } + // Verify all the memory has been freed, except for the ones for hash lookup. + ASSERT_EQ(nodePool->reservedBytes(), 1048576); + + driverWaitFlag = false; + driverWait.notifyAll(); + task.reset(); + + taskThread.join(); +} + +DEBUG_ONLY_TEST_F(HashJoinTest, hashTableCleanupAfterProbeFinish) { + auto buildVectors = makeVectors(buildType_, 5, 100); + auto probeVectors = makeVectors(probeType_, 5, 100); + + createDuckDbTable("t", probeVectors); + createDuckDbTable("u", buildVectors); + + HashProbe* probeOp{nullptr}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::getOutput", + std::function([&](Operator* op) { + if (probeOp == nullptr && op->operatorType() == "HashProbe") { + probeOp = dynamic_cast(op); + } + })); + + bool tableEmpty{false}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::noMoreInput", + std::function([&](Operator* op) { + if (op->operatorType() == "FilterProject") { + tableEmpty = (probeOp->testingTable()->numDistinct() == 0); + } + })); + + auto planNodeIdGenerator = std::make_shared(); + auto plan = PlanBuilder(planNodeIdGenerator) + .values(probeVectors, true) + .hashJoin( + {"t_k1"}, + {"u_k1"}, + PlanBuilder(planNodeIdGenerator) + .values(buildVectors, true) + .planNode(), + "", + concat(probeType_->names(), buildType_->names())) + .project({"t_k1", "t_k2", "t_v1", "u_k1", "u_k2", "u_v1"}) + .planNode(); + + auto tempDirectory = exec::test::TempDirectoryPath::create(); + HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get()) + .numDrivers(1) + .planNode(plan) + .injectSpill(false) + .spillDirectory(tempDirectory->getPath()) + .referenceQuery( + "SELECT t_k1, t_k2, t_v1, u_k1, u_k2, u_v1 FROM t, u WHERE t.t_k1 = u.u_k1") + .config(core::QueryConfig::kSpillStartPartitionBit, "29") + .run(); + ASSERT_TRUE(tableEmpty); +} +} // namespace diff --git a/velox/experimental/cudf/tests/LimitTest.cpp b/velox/experimental/cudf/tests/LimitTest.cpp new file mode 100644 index 00000000000..7fb01f49cbe --- /dev/null +++ b/velox/experimental/cudf/tests/LimitTest.cpp @@ -0,0 +1,89 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#include "velox/experimental/cudf/exec/ToCudf.h" + +#include "velox/exec/OutputBufferManager.h" +#include "velox/exec/tests/utils/HiveConnectorTestBase.h" +#include "velox/exec/tests/utils/PlanBuilder.h" + +using namespace facebook::velox; +using namespace facebook::velox::exec; +using namespace facebook::velox::exec::test; + +class LimitTest : public HiveConnectorTestBase { + void SetUp() override { + HiveConnectorTestBase::SetUp(); + cudf_velox::registerCudf(); + } + + void TearDown() override { + cudf_velox::unregisterCudf(); + HiveConnectorTestBase::TearDown(); + } +}; + +TEST_F(LimitTest, basic) { + vector_size_t batchSize = 1'000; + std::vector vectors; + for (int32_t i = 0; i < 3; ++i) { + auto c0 = makeFlatVector( + batchSize, [&](auto row) { return batchSize * i + row; }, nullEvery(5)); + auto c1 = makeFlatVector( + batchSize, [&](auto row) { return row; }, nullEvery(7)); + auto c2 = makeFlatVector( + batchSize, [](auto row) { return row * 0.1; }, nullEvery(11)); + vectors.push_back(makeRowVector({c0, c1, c2})); + } + createDuckDbTable(vectors); + + auto makePlan = [&](int64_t offset, int64_t limit) { + return PlanBuilder().values(vectors).limit(offset, limit, true).planNode(); + }; + + assertQuery(makePlan(0, 10), "SELECT * FROM tmp LIMIT 10"); + int64_t limit = (int64_t)(std::numeric_limits::max()) + 1000000; + int64_t offset = (int64_t)(std::numeric_limits::max()) + 1000; + assertQuery( + makePlan(0, limit), fmt::format("SELECT * FROM tmp LIMIT {}", limit)); + assertQuery(makePlan(0, 1'234), "SELECT * FROM tmp LIMIT 1234"); + + assertQuery(makePlan(17, 10), "SELECT * FROM tmp OFFSET 17 LIMIT 10"); + assertQuery(makePlan(17, 983), "SELECT * FROM tmp OFFSET 17 LIMIT 983"); + assertQuery(makePlan(17, 2'000), "SELECT * FROM tmp OFFSET 17 LIMIT 2000"); + assertQuery( + makePlan(offset, limit), + fmt::format("SELECT * FROM tmp OFFSET {} LIMIT {}", offset, limit)); + + assertQuery( + makePlan(offset, 2000), + fmt::format("SELECT * FROM tmp OFFSET {} LIMIT 2000", offset)); + + assertQuery(makePlan(1'000, 145), "SELECT * FROM tmp OFFSET 1000 LIMIT 145"); + assertQuery( + makePlan(1'000, 1'000), "SELECT * FROM tmp OFFSET 1000 LIMIT 1000"); + assertQuery( + makePlan(1'000, 1'234), "SELECT * FROM tmp OFFSET 1000 LIMIT 1234"); + + assertQuery(makePlan(1'234, 10), "SELECT * FROM tmp OFFSET 1234 LIMIT 10"); + assertQuery(makePlan(1'234, 983), "SELECT * FROM tmp OFFSET 1234 LIMIT 983"); + assertQuery( + makePlan(1'234, 1'000), "SELECT * FROM tmp OFFSET 1234 LIMIT 1000"); + assertQuery( + makePlan(1'234, 2'000), "SELECT * FROM tmp OFFSET 1234 LIMIT 2000"); + + assertQueryReturnsEmptyResult(makePlan(12'345, 10)); +} diff --git a/velox/experimental/cudf/tests/LocalPartitionTest.cpp b/velox/experimental/cudf/tests/LocalPartitionTest.cpp new file mode 100644 index 00000000000..53b33df2b4f --- /dev/null +++ b/velox/experimental/cudf/tests/LocalPartitionTest.cpp @@ -0,0 +1,162 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ +#include "velox/experimental/cudf/exec/ToCudf.h" + +#include "velox/common/base/tests/GTestUtils.h" +#include "velox/exec/PlanNodeStats.h" +#include "velox/exec/tests/utils/AssertQueryBuilder.h" +#include "velox/exec/tests/utils/HiveConnectorTestBase.h" +#include "velox/exec/tests/utils/PlanBuilder.h" + +namespace facebook::velox::exec::test { +namespace { + +class LocalPartitionTest : public HiveConnectorTestBase { + protected: + void SetUp() override { + HiveConnectorTestBase::SetUp(); + cudf_velox::registerCudf(); + } + + template + FlatVectorPtr makeFlatSequence(T start, vector_size_t size) { + return makeFlatVector(size, [start](auto row) { return start + row; }); + } + + template + FlatVectorPtr makeFlatSequence(T start, T max, vector_size_t size) { + return makeFlatVector( + size, [start, max](auto row) { return (start + row) % max; }); + } + + std::vector> writeToFiles( + const std::vector& vectors) { + auto filePaths = makeFilePaths(vectors.size()); + for (auto i = 0; i < vectors.size(); i++) { + writeToFile(filePaths[i]->getPath(), vectors[i]); + } + return filePaths; + } +}; + +TEST_F(LocalPartitionTest, gather) { + std::vector vectors = { + makeRowVector({makeFlatSequence(0, 100)}), + makeRowVector({makeFlatSequence(53, 100)}), + makeRowVector({makeFlatSequence(-71, 100)}), + }; + + auto planNodeIdGenerator = std::make_shared(); + + auto valuesNode = [&](int index) { + return PlanBuilder(planNodeIdGenerator).values({vectors[index]}).planNode(); + }; + + auto op = PlanBuilder(planNodeIdGenerator) + .localPartition( + {}, + { + valuesNode(0), + valuesNode(1), + valuesNode(2), + }) + .singleAggregation({}, {"min(c0)", "max(c0)"}) + .planNode(); + + auto task = assertQuery(op, "SELECT -71, 152"); + + auto filePaths = writeToFiles(vectors); + + auto rowType = asRowType(vectors[0]->type()); + + std::vector scanNodeIds; + + auto tableScanNode = [&]() { + auto node = PlanBuilder(planNodeIdGenerator).tableScan(rowType).planNode(); + scanNodeIds.push_back(node->id()); + return node; + }; + + op = PlanBuilder(planNodeIdGenerator) + .localPartition( + {}, + { + tableScanNode(), + tableScanNode(), + tableScanNode(), + }) + .singleAggregation({}, {"min(c0)", "max(c0)"}) + .planNode(); + + AssertQueryBuilder queryBuilder(op, duckDbQueryRunner_); + for (auto i = 0; i < filePaths.size(); ++i) { + queryBuilder.split( + scanNodeIds[i], makeHiveConnectorSplit(filePaths[i]->getPath())); + } + + task = queryBuilder.assertResults("SELECT -71, 152"); +} + +TEST_F(LocalPartitionTest, partition) { + std::vector vectors = { + makeRowVector({makeFlatSequence(0, 100)}), + makeRowVector({makeFlatSequence(53, 100)}), + makeRowVector({makeFlatSequence(-71, 100)}), + }; + + auto filePaths = writeToFiles(vectors); + + auto rowType = asRowType(vectors[0]->type()); + + auto planNodeIdGenerator = std::make_shared(); + + std::vector scanNodeIds; + + auto scanAggNode = [&]() { + auto builder = PlanBuilder(planNodeIdGenerator); + auto scanNode = builder.tableScan(rowType).planNode(); + scanNodeIds.push_back(scanNode->id()); + return builder.partialAggregation({"c0"}, {"max(c0)"}).planNode(); + }; + + auto op = PlanBuilder(planNodeIdGenerator) + .localPartition( + {"c0"}, + { + scanAggNode(), + scanAggNode(), + scanAggNode(), + }) + .finalAggregation() + .planNode(); + + createDuckDbTable(vectors); + + AssertQueryBuilder queryBuilder(op, duckDbQueryRunner_); + queryBuilder.maxDrivers(2); + queryBuilder.config(core::QueryConfig::kMaxLocalExchangePartitionCount, "2"); + + for (auto i = 0; i < filePaths.size(); ++i) { + queryBuilder.split( + scanNodeIds[i], makeHiveConnectorSplit(filePaths[i]->getPath())); + } + + auto task = + queryBuilder.assertResults("SELECT c0, max(c0) FROM tmp GROUP BY 1"); +} + +} // namespace +} // namespace facebook::velox::exec::test diff --git a/velox/experimental/cudf/tests/OrderByTest.cpp b/velox/experimental/cudf/tests/OrderByTest.cpp index d3eb5a75867..1ee44e0fdf7 100644 --- a/velox/experimental/cudf/tests/OrderByTest.cpp +++ b/velox/experimental/cudf/tests/OrderByTest.cpp @@ -13,8 +13,8 @@ * See the License for the specific language governing permissions and * limitations under the License. */ +#include "velox/experimental/cudf/exec/CudfConversion.h" #include "velox/experimental/cudf/exec/ToCudf.h" -#include "velox/experimental/cudf/exec/Utilities.h" #include "velox/common/base/tests/GTestUtils.h" #include "velox/core/QueryConfig.h" @@ -25,7 +25,6 @@ #include "velox/exec/tests/utils/PlanBuilder.h" #include -#include using namespace facebook::velox; using namespace facebook::velox::exec; @@ -319,4 +318,74 @@ TEST_F(OrderByTest, varfields) { testSingleKey(vectors, "c2"); } +/// Verifies output batch rows of OrderBy +TEST_F(OrderByTest, outputBatchRows) { + struct { + int numRowsPerBatch; + int preferredOutBatchBytes; + int maxOutBatchRows; + int expectedOutputVectors; + + // TODO: add output size check with spilling enabled + std::string debugString() const { + return fmt::format( + "numRowsPerBatch:{}, preferredOutBatchBytes:{}, maxOutBatchRows:{}, expectedOutputVectors:{}", + numRowsPerBatch, + preferredOutBatchBytes, + maxOutBatchRows, + expectedOutputVectors); + } + } testSettings[] = { + {1024, 1, 100, 1024}, + // estimated size per row is ~2092, set preferredOutBatchBytes to 20920, + // so each batch has 10 rows, so it would return 100 batches + {1000, 20920, 100, 100}, + // same as above, but maxOutBatchRows is 1, so it would return 1000 + // batches + {1000, 20920, 1, 1000}}; + + for (const auto& testData : testSettings) { + SCOPED_TRACE(testData.debugString()); + const vector_size_t batchSize = testData.numRowsPerBatch; + std::vector rowVectors; + auto c0 = makeFlatVector( + batchSize, [&](vector_size_t row) { return row; }, nullEvery(5)); + auto c1 = makeFlatVector( + batchSize, [&](vector_size_t row) { return row; }, nullEvery(11)); + std::vector vectors; + vectors.push_back(c0); + for (int i = 0; i < 256; ++i) { + vectors.push_back(c1); + } + rowVectors.push_back(makeRowVector(vectors)); + createDuckDbTable(rowVectors); + + core::PlanNodeId orderById; + auto plan = PlanBuilder() + .values(rowVectors) + .orderBy({fmt::format("{} ASC NULLS LAST", "c0")}, false) + .capturePlanNodeId(orderById) + .planNode(); + auto queryCtx = core::QueryCtx::create(executor_.get()); + queryCtx->testingOverrideConfigUnsafe( + {{core::QueryConfig::kPreferredOutputBatchBytes, + std::to_string(testData.preferredOutBatchBytes)}, + {core::QueryConfig::kMaxOutputBatchRows, + std::to_string(testData.maxOutBatchRows)}, + {facebook::velox::cudf_velox::CudfToVelox::kPassthroughMode, + "false"}}); + CursorParameters params; + params.planNode = plan; + params.queryCtx = queryCtx; + auto task = assertQueryOrdered( + params, "SELECT * FROM tmp ORDER BY c0 ASC NULLS LAST", {0}); + + EXPECT_EQ( + testData.expectedOutputVectors, + toPlanStats(task->taskStats()) + .at(orderById + "-to-velox") + .outputVectors); + } +} + } // namespace diff --git a/velox/experimental/cudf/tests/TableScanTest.cpp b/velox/experimental/cudf/tests/TableScanTest.cpp new file mode 100644 index 00000000000..84ddd2b43f0 --- /dev/null +++ b/velox/experimental/cudf/tests/TableScanTest.cpp @@ -0,0 +1,391 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#include "velox/experimental/cudf/connectors/parquet/ParquetConfig.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetConnector.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetConnectorSplit.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetDataSource.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetTableHandle.h" +#include "velox/experimental/cudf/tests/utils/ParquetConnectorTestBase.h" + +#include "velox/common/base/tests/GTestUtils.h" +#include "velox/common/file/tests/FaultyFile.h" +#include "velox/common/file/tests/FaultyFileSystem.h" +#include "velox/common/memory/MemoryArbitrator.h" +#include "velox/common/testutil/TestValue.h" +#include "velox/exec/Exchange.h" +#include "velox/exec/PlanNodeStats.h" +#include "velox/exec/TableScan.h" +#include "velox/exec/tests/utils/AssertQueryBuilder.h" +#include "velox/exec/tests/utils/HiveConnectorTestBase.h" +#include "velox/exec/tests/utils/LocalExchangeSource.h" +#include "velox/exec/tests/utils/PlanBuilder.h" +#include "velox/exec/tests/utils/TempDirectoryPath.h" +#include "velox/expression/ExprToSubfieldFilter.h" +#include "velox/type/Type.h" + +#include + +using namespace facebook::velox; +using namespace facebook::velox::core; +using namespace facebook::velox::exec; +using namespace facebook::velox::exec::test; +using namespace facebook::velox::common::test; +using namespace facebook::velox::tests::utils; +using namespace facebook::velox::cudf_velox; +using namespace facebook::velox::cudf_velox::exec; +using namespace facebook::velox::cudf_velox::exec::test; + +class TableScanTest : public virtual ParquetConnectorTestBase { + protected: + void SetUp() override { + ParquetConnectorTestBase::SetUp(); + ExchangeSource::factories().clear(); + ExchangeSource::registerFactory(createLocalExchangeSource); + } + + static void SetUpTestCase() { + ParquetConnectorTestBase::SetUpTestCase(); + } + + std::vector makeVectors( + int32_t count, + int32_t rowsPerVector, + const RowTypePtr& rowType = nullptr) { + auto inputs = rowType ? rowType : rowType_; + return ParquetConnectorTestBase::makeVectors(inputs, count, rowsPerVector); + } + + Split makeParquetSplit(std::string path, int64_t splitWeight = 0) { + return Split(makeParquetConnectorSplit(std::move(path), splitWeight)); + } + + std::shared_ptr assertQuery( + const PlanNodePtr& plan, + const std::shared_ptr& + parquetSplit, + const std::string& duckDbSql) { + return OperatorTestBase::assertQuery(plan, {parquetSplit}, duckDbSql); + } + + std::shared_ptr assertQuery( + const PlanNodePtr& plan, + const Split&& split, + const std::string& duckDbSql) { + return OperatorTestBase::assertQuery(plan, {split}, duckDbSql); + } + + std::shared_ptr assertQuery( + const PlanNodePtr& plan, + const std::vector>& filePaths, + const std::string& duckDbSql) { + return ParquetConnectorTestBase::assertQuery(plan, filePaths, duckDbSql); + } + + // Run query with spill enabled. + std::shared_ptr assertQuery( + const PlanNodePtr& plan, + const std::vector>& filePaths, + const std::string& spillDirectory, + const std::string& duckDbSql) { + return AssertQueryBuilder(plan, duckDbQueryRunner_) + .spillDirectory(spillDirectory) + .config(core::QueryConfig::kSpillEnabled, false) + .config(core::QueryConfig::kAggregationSpillEnabled, false) + .splits(makeParquetConnectorSplits(filePaths)) + .assertResults(duckDbSql); + } + + core::PlanNodePtr tableScanNode() { + return tableScanNode(rowType_); + } + + core::PlanNodePtr tableScanNode(const RowTypePtr& outputType) { + auto tableHandle = makeTableHandle(); + return PlanBuilder(pool_.get()) + .startTableScan() + .outputType(outputType) + .tableHandle(tableHandle) + .endTableScan() + .planNode(); + } + + static PlanNodeStats getTableScanStats(const std::shared_ptr& task) { + auto planStats = toPlanStats(task->taskStats()); + return std::move(planStats.at("0")); + } + + static std::unordered_map + getTableScanRuntimeStats(const std::shared_ptr& task) { + VELOX_NYI("RuntimeStats not yet implemented for the cudf ParquetConnector"); + // return task->taskStats().pipelineStats[0].operatorStats[0].runtimeStats; + } + + static int64_t getSkippedStridesStat(const std::shared_ptr& task) { + VELOX_NYI("RuntimeStats not yet implemented for the cudf ParquetConnector"); + // return getTableScanRuntimeStats(task)["skippedStrides"].sum; + } + + static int64_t getSkippedSplitsStat(const std::shared_ptr& task) { + VELOX_NYI("RuntimeStats not yet implemented for the cudf ParquetConnector"); + // return getTableScanRuntimeStats(task)["skippedSplits"].sum; + } + + static void waitForFinishedDrivers( + const std::shared_ptr& task, + uint32_t n) { + // Limit wait to 10 seconds. + size_t iteration{0}; + while (task->numFinishedDrivers() < n and iteration < 100) { + /* sleep override */ + usleep(100'000); // 0.1 second. + ++iteration; + } + ASSERT_EQ(n, task->numFinishedDrivers()); + } + + RowTypePtr rowType_{ + ROW({"c0", "c1", "c2", "c3", "c4", "c5", "c6"}, + {INTEGER(), + VARCHAR(), + TINYINT(), + DOUBLE(), + BIGINT(), + VARCHAR(), + REAL()})}; +}; + +TEST_F(TableScanTest, allColumns) { + auto vectors = makeVectors(10, 1'000); + auto filePath = TempFilePath::create(); + writeToFile(filePath->getPath(), vectors, "c"); + + createDuckDbTable(vectors); + auto plan = tableScanNode(); + auto task = assertQuery(plan, {filePath}, "SELECT * FROM tmp"); + + // A quick sanity check for memory usage reporting. Check that peak total + // memory usage for the project node is > 0. + auto planStats = toPlanStats(task->taskStats()); + auto scanNodeId = plan->id(); + auto it = planStats.find(scanNodeId); + ASSERT_TRUE(it != planStats.end()); + ASSERT_TRUE(it->second.peakMemoryBytes > 0); + + // Verifies there is no dynamic filter stats. + ASSERT_TRUE(it->second.dynamicFilterStats.empty()); + + // TODO: We are not writing any customStats yet so disable this check + // ASSERT_LT(0, it->second.customStats.at("ioWaitWallNanos").sum); +} + +TEST_F(TableScanTest, directBufferInputRawInputBytes) { + constexpr int kSize = 10; + auto vector = makeRowVector({ + makeFlatVector(kSize, folly::identity), + makeFlatVector(kSize, folly::identity), + makeFlatVector(kSize, folly::identity), + }); + auto filePath = TempFilePath::create(); + createDuckDbTable({vector}); + writeToFile(filePath->getPath(), {vector}, "c"); + + auto tableHandle = makeTableHandle(); + auto plan = PlanBuilder(pool_.get()) + .startTableScan() + .tableHandle(tableHandle) + .outputType(ROW({"c0", "c2"}, {BIGINT(), BIGINT()})) + .endTableScan() + .planNode(); + + std::unordered_map config; + std::unordered_map> + connectorConfigs = {}; + auto queryCtx = core::QueryCtx::create( + executor_.get(), + core::QueryConfig(std::move(config)), + connectorConfigs, + nullptr); + + auto task = AssertQueryBuilder(duckDbQueryRunner_) + .plan(plan) + .splits(makeParquetConnectorSplits({filePath})) + .queryCtx(queryCtx) + .assertResults("SELECT c0, c2 FROM tmp"); + + // A quick sanity check for memory usage reporting. Check that peak total + // memory usage for the project node is > 0. + auto planStats = toPlanStats(task->taskStats()); + auto scanNodeId = plan->id(); + auto it = planStats.find(scanNodeId); + ASSERT_TRUE(it != planStats.end()); + auto rawInputBytes = it->second.rawInputBytes; + // Reduced from 500 to 400 as cudf Parquet writer seems to be writing smaller + // files. + ASSERT_GE(rawInputBytes, 400); + + // TableScan runtime stats not available with Parquet connector yet +#if 0 + auto overreadBytes = + getTableScanRuntimeStats(task).at("overreadBytes").sum; + ASSERT_EQ(overreadBytes, 13); + ASSERT_EQ( + getTableScanRuntimeStats(task).at("storageReadBytes").sum, + rawInputBytes + overreadBytes); + ASSERT_GT(getTableScanRuntimeStats(task)["totalScanTime"].sum, 0); + ASSERT_GT(getTableScanRuntimeStats(task)["ioWaitWallNanos"].sum, 0); +#endif +} + +TEST_F(TableScanTest, columnAliases) { + auto vectors = makeVectors(1, 1'000); + auto filePath = TempFilePath::create(); + writeToFile(filePath->getPath(), vectors, "c"); + createDuckDbTable(vectors); + + std::string tableName = "t"; + std::unordered_map aliases = {{"a", "c0"}}; + auto outputType = ROW({"a"}, {INTEGER()}); + auto tableHandle = makeTableHandle(); + auto op = PlanBuilder(pool_.get()) + .startTableScan() + .tableHandle(tableHandle) + .tableName(tableName) + .outputType(outputType) + .columnAliases(aliases) + .endTableScan() + .planNode(); + assertQuery(op, {filePath}, "SELECT c0 FROM tmp"); +} + +TEST_F(TableScanTest, filterPushdown) { + auto rowType = + ROW({"c0", "c1", "c2", "c3"}, {TINYINT(), BIGINT(), DOUBLE(), BOOLEAN()}); + auto filePaths = makeFilePaths(10); + auto vectors = makeVectors(10, 1'000, rowType); + for (int32_t i = 0; i < vectors.size(); i++) { + writeToFile(filePaths[i]->getPath(), vectors[i]); + } + createDuckDbTable(vectors); + + // c1 >= 0 or null and c3 is true + // common::SubfieldFilters subfieldFilters = + // SubfieldFiltersBuilder() + // .add("c1", greaterThanOrEqual(0, true)) + // .add("c3", std::make_unique(true, false)) + // .build(); + // convert subfieldFilters to a typed expression + // c1 >= 0 or null and c3 is true + auto c1Expr = std::make_shared( + BOOLEAN(), + std::vector{ + std::make_shared(BIGINT(), "c1"), + std::make_shared(BIGINT(), int64_t(0)), + }, + "gte"); + + auto c3Expr = std::make_shared( + BOOLEAN(), + std::vector{ + std::make_shared(BOOLEAN(), "c3"), + std::make_shared(BOOLEAN(), true), + }, + "eq"); + + auto subfieldFilterExpr = std::make_shared( + BOOLEAN(), + std::vector{ + c1Expr, + c3Expr, + }, + "and"); + auto tableHandle = makeTableHandle( + "parquet_table", rowType, true, std::move(subfieldFilterExpr), nullptr); + + auto assignments = + facebook::velox::exec::test::HiveConnectorTestBase::allRegularColumns( + rowType); + + auto task = assertQuery( + PlanBuilder() + .startTableScan() + .outputType(ROW({"c1", "c3", "c0"}, {BIGINT(), BOOLEAN(), TINYINT()})) + .tableHandle(tableHandle) + .assignments(assignments) + .endTableScan() + .planNode(), + filePaths, + "SELECT c1, c3, c0 FROM tmp WHERE (c1 >= 0 ) AND c3"); + + auto tableScanStats = getTableScanStats(task); + // EXPECT_EQ(tableScanStats.rawInputRows, 10'000); + // EXPECT_LT(tableScanStats.inputRows, tableScanStats.rawInputRows); + EXPECT_EQ(tableScanStats.inputRows, tableScanStats.outputRows); + +#if 0 + // Repeat the same but do not project out the filtered columns. + assignments.clear(); + assignments["c0"] = + facebook::velox::exec::test::HiveConnectorTestBase::regularColumn( + "c0", TINYINT()); + assertQuery( + PlanBuilder() + .startTableScan() + .outputType(ROW({"c0"}, {TINYINT()})) + .tableHandle(tableHandle) + .assignments(assignments) + .endTableScan() + .planNode(), + filePaths, + "SELECT c0 FROM tmp WHERE (c1 >= 0 ) AND c3"); + + // TODO: zero column non-empty table is not possible in cudf, need to implement. + // Do the same for count, no columns projected out. + assignments.clear(); + assertQuery( + PlanBuilder() + .startTableScan() + .outputType(ROW({}, {})) + .tableHandle(tableHandle) + .assignments(assignments) + .endTableScan() + .singleAggregation({}, {"sum(1)"}) + .planNode(), + filePaths, + "SELECT count(*) FROM tmp WHERE (c1 >= 0 ) AND c3"); + + // Do the same for count, no filter, no projections. + assignments.clear(); + // subfieldFilters.clear(); // Explicitly clear this. + tableHandle = makeTableHandle( + "parquet_table", + rowType, + false, + nullptr, + nullptr); + assertQuery( + PlanBuilder() + .startTableScan() + .outputType(ROW({}, {})) + .tableHandle(tableHandle) + .assignments(assignments) + .endTableScan() + .singleAggregation({}, {"sum(1)"}) + .planNode(), + filePaths, + "SELECT count(*) FROM tmp"); +#endif +} diff --git a/velox/experimental/cudf/tests/TableWriteTest.cpp b/velox/experimental/cudf/tests/TableWriteTest.cpp new file mode 100644 index 00000000000..eec324453c5 --- /dev/null +++ b/velox/experimental/cudf/tests/TableWriteTest.cpp @@ -0,0 +1,811 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#include "velox/experimental/cudf/connectors/parquet/ParquetConfig.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetConnector.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetConnectorSplit.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetDataSource.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetTableHandle.h" +#include "velox/experimental/cudf/exec/ToCudf.h" +#include "velox/experimental/cudf/tests/utils/CudfPlanBuilder.h" +#include "velox/experimental/cudf/tests/utils/ParquetConnectorTestBase.h" + +#include "folly/dynamic.h" +#include "velox/common/base/Fs.h" +#include "velox/common/testutil/TestValue.h" +#include "velox/dwio/common/Options.h" +#include "velox/dwio/common/WriterFactory.h" +#include "velox/exec/PlanNodeStats.h" +#include "velox/exec/TableWriter.h" +#include "velox/exec/tests/utils/ArbitratorTestUtil.h" +#include "velox/exec/tests/utils/AssertQueryBuilder.h" +#include "velox/exec/tests/utils/PlanBuilder.h" +#include "velox/exec/tests/utils/TempDirectoryPath.h" + +#include + +#include + +using namespace facebook::velox; +using namespace facebook::velox::core; +using namespace facebook::velox::exec; +using namespace facebook::velox::common; +using namespace facebook::velox::connector; +using namespace facebook::velox::exec::test; +using namespace facebook::velox::common::test; +using namespace facebook::velox::common::testutil; +using namespace facebook::velox::dwio::common; + +using namespace facebook::velox::cudf_velox; +using namespace facebook::velox::cudf_velox::exec; +using namespace facebook::velox::cudf_velox::exec::test; + +constexpr uint64_t kQueryMemoryCapacity = 512 * MB; + +namespace { + +static std::shared_ptr generateAggregationNode( + const std::string& name, + const std::vector& groupingKeys, + AggregationNode::Step step, + const PlanNodePtr& source) { + core::TypedExprPtr inputField = + std::make_shared(BIGINT(), name); + auto callExpr = std::make_shared( + BIGINT(), std::vector{inputField}, "min"); + std::vector aggregateNames = {"min"}; + std::vector aggregates = { + core::AggregationNode::Aggregate{ + callExpr, {{BIGINT()}}, nullptr, {}, {}}}; + return std::make_shared( + core::PlanNodeId(), + step, + groupingKeys, + std::vector{}, + aggregateNames, + aggregates, + false, // ignoreNullKeys + source); +} + +} // namespace + +enum class TestMode { + kUnpartitioned, +}; + +std::string testModeString(TestMode mode) { + switch (mode) { + case TestMode::kUnpartitioned: + return "UNPARTITIONED"; + } + VELOX_UNREACHABLE(); +} + +FOLLY_ALWAYS_INLINE std::ostream& operator<<(std::ostream& os, TestMode mode) { + os << testModeString(mode); + return os; +} + +// NOTE: google parameterized test framework can't handle complex test +// parameters properly. So we encode the different test parameters into one +// integer value. +struct TestParam { + uint64_t value; + + explicit TestParam(uint64_t value) : value(value) {} + + TestParam( + FileFormat fileFormat, + TestMode testMode, + CommitStrategy commitStrategy, + bool multiDrivers, + CompressionKind compressionKind) { + value = static_cast(compressionKind) << 32 | + static_cast(static_cast(multiDrivers)) << 24 | + static_cast(fileFormat) << 16 | + static_cast(testMode) << 8 | + static_cast(commitStrategy); + } + + CompressionKind compressionKind() const { + return static_cast((value & ((1L << 40) - 1)) >> 32); + } + + bool multiDrivers() const { + return (value >> 24) != 0; + } + + FileFormat fileFormat() const { + return static_cast((value & ((1L << 24) - 1)) >> 16); + } + + TestMode testMode() const { + return static_cast((value & ((1L << 16) - 1)) >> 8); + } + + CommitStrategy commitStrategy() const { + return static_cast((value & ((1L << 8) - 1))); + } + + std::string toString() const { + return fmt::format( + "FileFormat[{}] TestMode[{}] commitStrategy[{}] multiDrivers[{}] compression[{}]", + dwio::common::toString((fileFormat())), + testModeString(testMode()), + commitStrategyToString(commitStrategy()), + multiDrivers(), + compressionKindToString(compressionKind())); + } +}; + +class TableWriteTest : public ParquetConnectorTestBase { + protected: + explicit TableWriteTest(uint64_t testValue) + : testParam_(static_cast(testValue)), + fileFormat_(dwio::common::FileFormat::PARQUET), + testMode_(testParam_.testMode()), + numTableWriterCount_( + testParam_.multiDrivers() ? kNumTableWriterCount : 1), + commitStrategy_(testParam_.commitStrategy()), + compressionKind_(testParam_.compressionKind()) { + LOG(INFO) << testParam_.toString(); + if (cudfDebugEnabled()) { + std::cout << testParam_.toString() << std::endl; + } + + auto rowType = + ROW({"c0", "c1", "c2", "c3", "c4", "c5"}, + {BIGINT(), INTEGER(), SMALLINT(), REAL(), DOUBLE(), VARCHAR()}); + setDataTypes(rowType); + } + + void SetUp() override { + ParquetConnectorTestBase::SetUp(); + } + + std::shared_ptr assertQueryWithWriterConfigs( + const core::PlanNodePtr& plan, + std::vector> filePaths, + const std::string& duckDbSql, + bool spillEnabled = false) { + std::vector splits; + for (const auto& filePath : filePaths) { + splits.push_back(facebook::velox::exec::Split( + makeParquetConnectorSplit(filePath->getPath()))); + } + if (!spillEnabled) { + return AssertQueryBuilder(plan, duckDbQueryRunner_) + .maxDrivers(2 * kNumTableWriterCount) + .config( + QueryConfig::kTaskWriterCount, + std::to_string(numTableWriterCount_)) + .splits(splits) + .assertResults(duckDbSql); + } + } + + std::shared_ptr assertQueryWithWriterConfigs( + const core::PlanNodePtr& plan, + const std::string& duckDbSql, + bool enableSpill = false) { + if (!enableSpill) { + TestScopedSpillInjection scopedSpillInjection(100); + return AssertQueryBuilder(plan, duckDbQueryRunner_) + .maxDrivers(2 * kNumTableWriterCount) + .config( + QueryConfig::kTaskWriterCount, + std::to_string(numTableWriterCount_)) + .config(core::QueryConfig::kSpillEnabled, "true") + .config(QueryConfig::kWriterSpillEnabled, "true") + // Scale writer settings to trigger partition rebalancing. + .config(QueryConfig::kScaleWriterRebalanceMaxMemoryUsageRatio, "1.0") + .config( + QueryConfig::kScaleWriterMinProcessedBytesRebalanceThreshold, "0") + .config( + QueryConfig:: + kScaleWriterMinPartitionProcessedBytesRebalanceThreshold, + "0") + .assertResults(duckDbSql); + } + } + + RowVectorPtr runQueryWithWriterConfigs( + const core::PlanNodePtr& plan, + bool spillEnabled = false) { + if (!spillEnabled) { + return AssertQueryBuilder(plan, duckDbQueryRunner_) + .maxDrivers(2 * kNumTableWriterCount) + .config( + QueryConfig::kTaskWriterCount, + std::to_string(numTableWriterCount_)) + // Scale writer settings to trigger partition rebalancing. + .config(QueryConfig::kScaleWriterRebalanceMaxMemoryUsageRatio, "1.0") + .config( + QueryConfig::kScaleWriterMinProcessedBytesRebalanceThreshold, "0") + .config( + QueryConfig:: + kScaleWriterMinPartitionProcessedBytesRebalanceThreshold, + "0") + .copyResults(pool()); + } + } + + void setCommitStrategy(CommitStrategy commitStrategy) { + commitStrategy_ = commitStrategy; + } + + void setDataTypes( + const RowTypePtr& inputType, + const RowTypePtr& tableSchema = nullptr) { + rowType_ = inputType; + if (tableSchema != nullptr) { + setTableSchema(tableSchema); + } else { + setTableSchema(rowType_); + } + } + + void setTableSchema(const RowTypePtr& tableSchema) { + tableSchema_ = tableSchema; + } + + std::vector> + makeParquetConnectorSplits( + const std::shared_ptr& directoryPath) { + return makeParquetConnectorSplits(directoryPath->getPath()); + } + + std::vector> + makeParquetConnectorSplits(const std::string& directoryPath) { + std::vector> + splits; + + for (auto& path : fs::recursive_directory_iterator(directoryPath)) { + if (path.is_regular_file()) { + splits.push_back(ParquetConnectorTestBase::makeParquetConnectorSplits( + path.path().string(), 1)[0]); + } + } + + return splits; + } + + // Lists and returns all the regular files from a given directory recursively. + std::vector listAllFiles(const std::string& directoryPath) { + std::vector files; + for (auto& path : fs::recursive_directory_iterator(directoryPath)) { + if (path.is_regular_file()) { + files.push_back(path.path().filename()); + } + } + return files; + } + + // Builds and returns the parquet splits from the list of files with one split + // per each file. + std::vector> + makeParquetConnectorSplits( + const std::vector& filePaths) { + std::vector> + splits; + for (const auto& filePath : filePaths) { + splits.push_back(ParquetConnectorTestBase::makeParquetConnectorSplits( + filePath.string(), 1)[0]); + } + return splits; + } + + std::vector makeVectors( + int32_t numVectors, + int32_t rowsPerVector) { + return ParquetConnectorTestBase::makeVectors( + rowType_, numVectors, rowsPerVector); + } + + RowVectorPtr makeConstantVector(size_t size) { + return makeRowVector( + rowType_->names(), + {makeConstant((int64_t)123'456, size), + makeConstant((int32_t)321, size), + makeConstant((int16_t)12'345, size), + makeConstant(variant(TypeKind::REAL), size), + makeConstant((double)1'234.01, size), + makeConstant(variant(TypeKind::VARCHAR), size)}); + } + + std::vector makeBatches( + vector_size_t numBatches, + std::function makeVector) { + std::vector batches; + batches.reserve(numBatches); + for (int32_t i = 0; i < numBatches; ++i) { + batches.push_back(makeVector(i)); + } + return batches; + } + + std::set getLeafSubdirectories( + const std::string& directoryPath) { + std::set subdirectories; + for (auto& path : fs::recursive_directory_iterator(directoryPath)) { + if (path.is_regular_file()) { + subdirectories.emplace(path.path().parent_path().string()); + } + } + return subdirectories; + } + + std::vector getRecursiveFiles(const std::string& directoryPath) { + std::vector files; + for (auto& path : fs::recursive_directory_iterator(directoryPath)) { + if (path.is_regular_file()) { + files.push_back(path.path().string()); + } + } + return files; + } + + uint32_t countRecursiveFiles(const std::string& directoryPath) { + return getRecursiveFiles(directoryPath).size(); + } + + // Helper method to return InsertTableHandle. + std::shared_ptr createInsertTableHandle( + const RowTypePtr& outputRowType, + const cudf_velox::connector::parquet::LocationHandle::TableType& + outputTableType, + const std::string& outputDirectoryPath, + const std::optional compressionKind = {}) { + return std::make_shared( + kParquetConnectorId, + makeParquetInsertTableHandle( + outputRowType->names(), + outputRowType->children(), + makeLocationHandle(outputDirectoryPath, outputTableType), + compressionKind)); + } + + // Returns a table insert plan node. + PlanNodePtr createInsertPlan( + PlanBuilder& inputPlan, + const RowTypePtr& outputRowType, + const std::string& outputDirectoryPath, + const std::optional compressionKind = {}, + int numTableWriters = 1, + const cudf_velox::connector::parquet::LocationHandle::TableType& + outputTableType = + cudf_velox::connector::parquet::LocationHandle::TableType::kNew, + const CommitStrategy& outputCommitStrategy = CommitStrategy::kNoCommit, + bool aggregateResult = true, + std::shared_ptr aggregationNode = nullptr) { + return createInsertPlan( + inputPlan, + inputPlan.planNode()->outputType(), + outputRowType, + outputDirectoryPath, + compressionKind, + numTableWriters, + outputTableType, + outputCommitStrategy, + aggregateResult, + aggregationNode); + } + + PlanNodePtr createInsertPlan( + PlanBuilder& inputPlan, + const RowTypePtr& inputRowType, + const RowTypePtr& tableRowType, + const std::string& outputDirectoryPath, + const std::optional compressionKind = {}, + int numTableWriters = 1, + const cudf_velox::connector::parquet::LocationHandle::TableType& + outputTableType = + cudf_velox::connector::parquet::LocationHandle::TableType::kNew, + const CommitStrategy& outputCommitStrategy = CommitStrategy::kNoCommit, + bool aggregateResult = true, + std::shared_ptr aggregationNode = nullptr) { + VELOX_CHECK( + numTableWriters == 1, "Multiple CudfTableWriters not yet supported"); + return createInsertPlanWithSingleWriter( + inputPlan, + inputRowType, + tableRowType, + outputDirectoryPath, + compressionKind, + outputTableType, + outputCommitStrategy, + aggregateResult, + aggregationNode); + } + + PlanNodePtr createInsertPlanWithSingleWriter( + PlanBuilder& inputPlan, + const RowTypePtr& inputRowType, + const RowTypePtr& tableRowType, + const std::string& outputDirectoryPath, + const std::optional compressionKind, + const cudf_velox::connector::parquet::LocationHandle::TableType& + outputTableType, + const CommitStrategy& outputCommitStrategy, + bool aggregateResult, + std::shared_ptr aggregationNode) { + const bool addScaleWriterExchange = false; + auto insertPlan = inputPlan; + insertPlan + .addNode(addCudfTableWriter( + inputRowType, + tableRowType->names(), + aggregationNode, + createInsertTableHandle( + tableRowType, + outputTableType, + outputDirectoryPath, + compressionKind), + outputCommitStrategy)) + .capturePlanNodeId(tableWriteNodeId_); + if (aggregateResult) { + insertPlan.project({TableWriteTraits::rowCountColumnName()}) + .singleAggregation( + {}, + {fmt::format("sum({})", TableWriteTraits::rowCountColumnName())}); + } + return insertPlan.planNode(); + } + + // Return the corresponding column names in 'inputRowType' of + // 'tableColumnNames' from 'tableRowType'. + static std::vector inputColumnNames( + const std::vector& tableColumnNames, + const RowTypePtr& tableRowType, + const RowTypePtr& inputRowType) { + std::vector inputNames; + inputNames.reserve(tableColumnNames.size()); + for (const auto& tableColumnName : tableColumnNames) { + const auto columnIdx = tableRowType->getChildIdx(tableColumnName); + inputNames.push_back(inputRowType->nameOf(columnIdx)); + } + return inputNames; + } + + // Parameter partitionName is string formatted in the Parquet style + // key1=value1/key2=value2/... Parameter partitionTypes are types of partition + // keys in the same order as in partitionName.The return value is a SQL + // predicate with values single quoted for string and date and not quoted for + // other supported types, ex., key1='value1' AND key2=value2 AND ... + std::string partitionNameToPredicate( + const std::string& partitionName, + const std::vector& partitionTypes) { + std::vector conjuncts; + + std::vector partitionKeyValues; + folly::split('/', partitionName, partitionKeyValues); + VELOX_CHECK_EQ(partitionKeyValues.size(), partitionTypes.size()); + + for (auto i = 0; i < partitionKeyValues.size(); ++i) { + if (partitionTypes[i]->isVarchar() || partitionTypes[i]->isVarbinary() || + partitionTypes[i]->isDate()) { + conjuncts.push_back( + partitionKeyValues[i] + .replace(partitionKeyValues[i].find("="), 1, "='") + .append("'")); + } else { + conjuncts.push_back(partitionKeyValues[i]); + } + } + + return folly::join(" AND ", conjuncts); + } + + // Verifies if a unbucketed file name is encoded properly based on the + // used commit strategy. + void verifyUnbucketedFilePath( + const std::filesystem::path& filePath, + const std::string& targetDir) { + ASSERT_EQ(filePath.parent_path().string(), targetDir); + if (commitStrategy_ == CommitStrategy::kNoCommit) { + ASSERT_TRUE(RE2::FullMatch( + filePath.filename().string(), + fmt::format( + "test_cursor.+_[0-{}]_{}_.+", + numTableWriterCount_ - 1, + tableWriteNodeId_))) + << filePath.filename().string(); + } else { + ASSERT_TRUE(RE2::FullMatch( + filePath.filename().string(), + fmt::format( + ".tmp.velox.test_cursor.+_[0-{}]_{}_.+", + numTableWriterCount_ - 1, + tableWriteNodeId_))) + << filePath.filename().string(); + } + } + + // Verifies the file layout and data produced by a table writer. + void verifyTableWriterOutput( + const std::string& targetDir, + const RowTypePtr& bucketCheckFileType, + bool verifyPartitionedData = true, + bool verifyBucketedData = true) { + SCOPED_TRACE(testParam_.toString()); + std::vector filePaths; + std::vector dirPaths; + for (auto& path : fs::recursive_directory_iterator(targetDir)) { + if (path.is_regular_file()) { + filePaths.push_back(path.path()); + } else { + dirPaths.push_back(path.path()); + } + } + if (testMode_ == TestMode::kUnpartitioned) { + ASSERT_EQ(dirPaths.size(), 0); + ASSERT_LE(filePaths.size(), numTableWriterCount_); + verifyUnbucketedFilePath(filePaths[0], targetDir); + return; + } + } + + int getNumWriters() { + return numTableWriterCount_; + } + + static inline int kNumTableWriterCount = 1; + + const TestParam testParam_; + const FileFormat fileFormat_ = FileFormat::PARQUET; + const TestMode testMode_; + const int numTableWriterCount_; + + RowTypePtr rowType_; + RowTypePtr tableSchema_; + CommitStrategy commitStrategy_; + std::optional compressionKind_; + std::vector sortColumnIndices_; + std::vector sortedFlags_; + core::PlanNodeId tableWriteNodeId_; +}; + +class BasicTableWriteTest : public ParquetConnectorTestBase {}; + +TEST_F(BasicTableWriteTest, roundTrip) { + vector_size_t size = 1'000; + auto data = makeRowVector({ + makeFlatVector(size, [](auto row) { return row; }), + makeFlatVector( + size, [](auto row) { return row * 2; }, nullEvery(7)), + }); + + auto sourceFilePath = TempFilePath::create(); + writeToFile(sourceFilePath->getPath(), data); + + auto targetDirectoryPath = TempDirectoryPath::create(); + + auto rowType = asRowType(data->type()); + auto plan = PlanBuilder() + .startTableScan() + .outputType(rowType) + .tableHandle(ParquetConnectorTestBase::makeTableHandle()) + .endTableScan() + .addNode(cudfTableWrite(targetDirectoryPath->getPath())) + .planNode(); + + auto results = + AssertQueryBuilder(plan) + .split(makeParquetConnectorSplit(sourceFilePath->getPath())) + .copyResults(pool()); + ASSERT_EQ(2, results->size()); + + // First column has number of rows written in the first row and nulls in other + // rows. + auto rowCount = results->childAt(TableWriteTraits::kRowCountChannel) + ->as>(); + ASSERT_FALSE(rowCount->isNullAt(0)); + ASSERT_EQ(size, rowCount->valueAt(0)); + ASSERT_TRUE(rowCount->isNullAt(1)); + + // Second column contains details about written files. + auto details = results->childAt(TableWriteTraits::kFragmentChannel) + ->as>(); + ASSERT_TRUE(details->isNullAt(0)); + ASSERT_FALSE(details->isNullAt(1)); + folly::dynamic obj = folly::parseJson(details->valueAt(1)); + + ASSERT_EQ(size, obj["rowCount"].asInt()); + auto fileWriteInfos = obj["fileWriteInfos"]; + ASSERT_EQ(1, fileWriteInfos.size()); + auto writeFileName = fileWriteInfos[0]["writeFileName"].asString(); + + // Read from 'writeFileName' and verify the data matches the original. + plan = PlanBuilder() + .startTableScan() + .outputType(rowType) + .tableHandle(ParquetConnectorTestBase::makeTableHandle()) + .endTableScan() + .planNode(); + + auto copy = AssertQueryBuilder(plan) + .split(makeParquetConnectorSplit(fmt::format( + "{}/{}", targetDirectoryPath->getPath(), writeFileName))) + .copyResults(pool()); + assertEqualResults({data}, {copy}); +} + +TEST_F(BasicTableWriteTest, targetFileName) { + constexpr const char* kFileName = "test.parquet"; + auto data = makeRowVector({makeFlatVector(10, folly::identity)}); + auto directory = TempDirectoryPath::create(); + auto plan = PlanBuilder() + .values({data}) + .addNode(cudfTableWrite( + directory->getPath(), + dwio::common::FileFormat::PARQUET, + {}, + nullptr, + kFileName)) + .planNode(); + + auto results = AssertQueryBuilder(plan).copyResults(pool()); + auto* details = results->childAt(TableWriteTraits::kFragmentChannel) + ->asUnchecked>(); + auto detail = folly::parseJson(details->valueAt(1)); + auto fileWriteInfos = detail["fileWriteInfos"]; + ASSERT_EQ(1, fileWriteInfos.size()); + ASSERT_EQ(fileWriteInfos[0]["writeFileName"].asString(), kFileName); + plan = PlanBuilder() + .startTableScan() + .outputType(asRowType(data->type())) + .tableHandle(ParquetConnectorTestBase::makeTableHandle()) + .endTableScan() + .planNode(); + AssertQueryBuilder(plan) + .split(makeParquetConnectorSplit( + fmt::format("{}/{}", directory->getPath(), kFileName))) + .assertResults(data); +} + +class UnpartitionedTableWriterTest + : public TableWriteTest, + public testing::WithParamInterface { + public: + UnpartitionedTableWriterTest() : TableWriteTest(GetParam()) {} + + static std::vector getTestParams() { + std::vector testParams; + const auto multiDriverOptions = std::vector{false, true}; + for (bool multiDrivers : multiDriverOptions) { + testParams.push_back(TestParam{ + FileFormat::PARQUET, + TestMode::kUnpartitioned, + CommitStrategy::kNoCommit, + multiDrivers, + CompressionKind_NONE} + .value); + testParams.push_back(TestParam{ + FileFormat::PARQUET, + TestMode::kUnpartitioned, + CommitStrategy::kTaskCommit, + multiDrivers, + CompressionKind_NONE} + .value); + } + return testParams; + } +}; + +TEST_P(UnpartitionedTableWriterTest, differentCompression) { + std::vector compressions{ + CompressionKind_NONE, + CompressionKind_SNAPPY, + CompressionKind_ZSTD, + CompressionKind_LZ4, + CompressionKind_MAX}; + + for (auto compressionKind : compressions) { + auto input = makeVectors(10, 10); + auto outputDirectory = TempDirectoryPath::create(); + if (compressionKind == CompressionKind_MAX) { + VELOX_ASSERT_THROW( + createInsertPlan( + PlanBuilder().values(input), + rowType_, + outputDirectory->getPath(), + compressionKind, + numTableWriterCount_, + cudf_velox::connector::parquet::LocationHandle::TableType::kNew), + "Unsupported compression type: CompressionKind_MAX"); + return; + } + auto plan = createInsertPlan( + PlanBuilder().values(input), + rowType_, + outputDirectory->getPath(), + compressionKind, + numTableWriterCount_, + cudf_velox::connector::parquet::LocationHandle::TableType::kNew); + + auto result = AssertQueryBuilder(plan) + .config( + QueryConfig::kTaskWriterCount, + std::to_string(numTableWriterCount_)) + .copyResults(pool()); + assertEqualResults( + {makeRowVector({makeConstant(100, 1)})}, {result}); + } +} + +// Test not really needed as we always write a TableType::kNew table in Parquet +// DataSink +TEST_P(UnpartitionedTableWriterTest, immutableSettings) { + struct { + cudf_velox::connector::parquet::LocationHandle::TableType dataType; + bool immutableFilesEnabled; + bool expectedInsertSuccees; + + std::string debugString() const { + return fmt::format( + "dataType:{}, immutableFilesEnabled:{}, operationSuccess:{}", + dataType, + immutableFilesEnabled, + expectedInsertSuccees); + } + } testSettings[] = { + {cudf_velox::connector::parquet::LocationHandle::TableType::kNew, + true, + true}, + {cudf_velox::connector::parquet::LocationHandle::TableType::kNew, + false, + true}}; + + for (auto testData : testSettings) { + SCOPED_TRACE(testData.debugString()); + std::unordered_map propFromFile{ + {"parquet.immutable-files", + testData.immutableFilesEnabled ? "true" : "false"}}; + std::shared_ptr config{ + std::make_shared(std::move(propFromFile))}; + resetParquetConnector(config); + + auto input = makeVectors(10, 10); + auto outputDirectory = TempDirectoryPath::create(); + auto plan = createInsertPlan( + PlanBuilder().values(input), + rowType_, + outputDirectory->getPath(), + CompressionKind_NONE, + numTableWriterCount_, + testData.dataType); + + if (!testData.expectedInsertSuccees) { + VELOX_ASSERT_THROW( + AssertQueryBuilder(plan).copyResults(pool()), + "Parquet tables are immutable."); + } else { + auto result = AssertQueryBuilder(plan) + .config( + QueryConfig::kTaskWriterCount, + std::to_string(numTableWriterCount_)) + .copyResults(pool()); + assertEqualResults( + {makeRowVector({makeConstant(100, 1)})}, {result}); + } + } +} + +VELOX_INSTANTIATE_TEST_SUITE_P( + TableWriterTest, + UnpartitionedTableWriterTest, + testing::ValuesIn(UnpartitionedTableWriterTest::getTestParams())); diff --git a/velox/experimental/cudf/tests/utils/CMakeLists.txt b/velox/experimental/cudf/tests/utils/CMakeLists.txt new file mode 100644 index 00000000000..f94bab870f2 --- /dev/null +++ b/velox/experimental/cudf/tests/utils/CMakeLists.txt @@ -0,0 +1,37 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# Licensed 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. + +add_library(velox_cudf_exec_test_lib ParquetConnectorTestBase.cpp + CudfPlanBuilder.cpp) + +set_target_properties( + velox_cudf_exec_test_lib + PROPERTIES CUDA_ARCHITECTURES native) + +target_link_libraries( + velox_cudf_exec_test_lib + velox_vector_test_lib + velox_temp_path + velox_cursor + cudf::cudf + velox_cudf_exec + velox_cudf_vector + velox_core + velox_exception + velox_expression + velox_parse_parser + velox_duckdb_conversion + velox_file_test_utils + velox_cudf_parquet_connector + velox_aggregates) diff --git a/velox/experimental/cudf/tests/utils/CudfPlanBuilder.cpp b/velox/experimental/cudf/tests/utils/CudfPlanBuilder.cpp new file mode 100644 index 00000000000..7f444f34357 --- /dev/null +++ b/velox/experimental/cudf/tests/utils/CudfPlanBuilder.cpp @@ -0,0 +1,98 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#include "velox/experimental/cudf/tests/utils/CudfPlanBuilder.h" + +#include "velox/dwio/common/Options.h" +#include "velox/exec/TableWriter.h" +#include "velox/exec/tests/utils/PlanBuilder.h" + +namespace facebook::velox::cudf_velox::exec::test { + +std::function addCudfTableWriter( + const RowTypePtr& inputColumns, + const std::vector& tableColumnNames, + const std::shared_ptr& aggregationNode, + const std::shared_ptr& insertHandle, + facebook::velox::connector::CommitStrategy commitStrategy) { + return [=](core::PlanNodeId nodeId, + core::PlanNodePtr source) -> core::PlanNodePtr { + return std::make_shared( + nodeId, + inputColumns, + tableColumnNames, + aggregationNode, + insertHandle, + false, + TableWriteTraits::outputType(aggregationNode), + commitStrategy, + std::move(source)); + }; +} + +std::function cudfTableWrite( + const std::string& outputDirectoryPath, + const dwio::common::FileFormat fileFormat, + const std::shared_ptr& aggregationNode, + const std::shared_ptr& options, + const std::string& outputFileName) { + return cudfTableWrite( + outputDirectoryPath, + fileFormat, + aggregationNode, + kParquetConnectorId, + {}, + options, + outputFileName); +} + +std::function cudfTableWrite( + const std::string& outputDirectoryPath, + const dwio::common::FileFormat fileFormat, + const std::shared_ptr& aggregationNode, + const std::string_view& connectorId, + const std::unordered_map& serdeParameters, + const std::shared_ptr& options, + const std::string& outputFileName, + const common::CompressionKind compression, + const RowTypePtr& schema) { + return [=](core::PlanNodeId nodeId, + core::PlanNodePtr source) -> core::PlanNodePtr { + auto rowType = schema ? schema : source->outputType(); + + auto locationHandle = ParquetConnectorTestBase::makeLocationHandle( + outputDirectoryPath, + cudf_velox::connector::parquet::LocationHandle::TableType::kNew, + outputFileName); + auto parquetHandle = ParquetConnectorTestBase::makeParquetInsertTableHandle( + rowType->names(), rowType->children(), locationHandle, compression); + auto insertHandle = std::make_shared( + std::string(connectorId), parquetHandle); + + return std::make_shared( + nodeId, + rowType, + rowType->names(), + aggregationNode, + insertHandle, + false, + TableWriteTraits::outputType(aggregationNode), + facebook::velox::connector::CommitStrategy::kNoCommit, + std::move(source)); + }; +} + +} // namespace facebook::velox::cudf_velox::exec::test diff --git a/velox/experimental/cudf/tests/utils/CudfPlanBuilder.h b/velox/experimental/cudf/tests/utils/CudfPlanBuilder.h new file mode 100644 index 00000000000..9626f60b57a --- /dev/null +++ b/velox/experimental/cudf/tests/utils/CudfPlanBuilder.h @@ -0,0 +1,99 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#include "velox/experimental/cudf/connectors/parquet/ParquetDataSink.h" +#include "velox/experimental/cudf/tests/utils/ParquetConnectorTestBase.h" + +#include "velox/dwio/common/Options.h" +#include "velox/exec/tests/utils/PlanBuilder.h" + +#include + +namespace facebook::velox::cudf_velox::exec::test { + +using namespace facebook::velox; +using namespace facebook::velox::core; +using namespace facebook::velox::exec; +using namespace facebook::velox::common; +using namespace facebook::velox::exec::test; +using namespace facebook::velox::common::test; +using namespace facebook::velox::common::testutil; +using namespace facebook::velox::dwio::common; + +// Adds a TableWriter node to write all input columns into a Parquet table. +std::function addCudfTableWriter( + const RowTypePtr& inputColumns, + const std::vector& tableColumnNames, + const std::shared_ptr& aggregationNode, + const std::shared_ptr& insertHandle, + facebook::velox::connector::CommitStrategy commitStrategy = + facebook::velox::connector::CommitStrategy::kNoCommit); + +/// Adds a TableWriteNode to write all input columns into an un-partitioned +/// un-bucketed Parquet table without compression. +/// +/// @param outputDirectoryPath Path to a directory to write data to. +/// @param fileFormat File format to use for the written data. +/// @param aggregationNode AggregationNode for column statistics collection +/// during write. +/// @param polymorphic options object to be passed to the writer. +/// write, supported aggregation types vary for different column types. +/// @param outputFileName Optional file name of the output. If specified +/// (non-empty), use it instead of generating the file name in Velox. Should +/// only be specified in non-bucketing write. +/// For example: +/// Boolean: count, countIf. +/// NumericType/Date/Timestamp: min, max, approx_distinct, count. +/// Varchar: count, approx_distinct, sum_data_size_for_stats, +/// max_data_size_for_stats. +std::function cudfTableWrite( + const std::string& outputDirectoryPath, + const dwio::common::FileFormat fileFormat = + dwio::common::FileFormat::PARQUET, + const std::shared_ptr& aggregationNode = nullptr, + const std::shared_ptr& options = nullptr, + const std::string& outputFileName = ""); + +/// Adds a TableWriteNode to write all input columns into Parquet +/// table with compression. +/// +/// @param outputDirectoryPath Path to a directory to write data to. +/// @param fileFormat File format to use for the written data. +/// @param aggregationNode AggregationNode for column statistics collection +/// during write. +/// @param connectorId Name used to register the connector. +/// @param serdeParameters Additional parameters passed to the writer. +/// @param Option objects passed to the writer. +/// @param outputFileName Optional file name of the output. If specified +/// (non-empty), use it instead of generating the file name in Velox. Should +/// only be specified in non-bucketing write. +/// @param compressionKind Compression scheme to use for writing the +/// output data files. +/// @param schema Output schema to be passed to the writer. By default use the +/// output of the previous operator. +std::function cudfTableWrite( + const std::string& outputDirectoryPath, + const dwio::common::FileFormat fileFormat, + const std::shared_ptr& aggregationNode, + const std::string_view& connectorId = kParquetConnectorId, + const std::unordered_map& serdeParameters = {}, + const std::shared_ptr& options = nullptr, + const std::string& outputFileName = "", + const common::CompressionKind compression = + common::CompressionKind::CompressionKind_NONE, + const RowTypePtr& schema = nullptr); + +} // namespace facebook::velox::cudf_velox::exec::test diff --git a/velox/experimental/cudf/tests/utils/ParquetConnectorTestBase.cpp b/velox/experimental/cudf/tests/utils/ParquetConnectorTestBase.cpp new file mode 100644 index 00000000000..cbb2318bf1c --- /dev/null +++ b/velox/experimental/cudf/tests/utils/ParquetConnectorTestBase.cpp @@ -0,0 +1,303 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#include "velox/experimental/cudf/exec/VeloxCudfInterop.h" +#include "velox/experimental/cudf/tests/utils/ParquetConnectorTestBase.h" +#include "velox/experimental/cudf/vector/CudfVector.h" + +#include "velox/common/base/Exceptions.h" +#include "velox/common/file/FileSystems.h" +#include "velox/common/file/tests/FaultyFileSystem.h" +#include "velox/dwio/common/FileSink.h" +#include "velox/dwio/common/tests/utils/BatchMaker.h" +#include "velox/dwio/dwrf/writer/FlushPolicy.h" +#include "velox/exec/tests/utils/AssertQueryBuilder.h" + +#include +#include +#include +#include + +#include +#include + +namespace facebook::velox::cudf_velox::exec::test { + +namespace { + +void fillColumnNames( + cudf::io::table_input_metadata& tableMeta, + const std::string& prefix) { + // Fill unnamed columns' names in cudf table_meta + std::function + addDefaultName = + [&](cudf::io::column_in_metadata& colMeta, std::string defaultName) { + if (colMeta.get_name().empty()) { + colMeta.set_name(defaultName); + } + for (int32_t i = 0; i < colMeta.num_children(); ++i) { + addDefaultName( + colMeta.child(i), fmt::format("{}_{}", defaultName, i)); + } + }; + for (int32_t i = 0; i < tableMeta.column_metadata.size(); ++i) { + addDefaultName(tableMeta.column_metadata[i], prefix + std::to_string(i)); + } +} + +} // namespace + +ParquetConnectorTestBase::ParquetConnectorTestBase() { + filesystems::registerLocalFileSystem(); + tests::utils::registerFaultyFileSystem(); +} + +void ParquetConnectorTestBase::SetUp() { + OperatorTestBase::SetUp(); + facebook::velox::connector::registerConnectorFactory( + std::make_shared()); + auto parquetConnector = + facebook::velox::connector::getConnectorFactory( + connector::parquet::ParquetConnectorFactory::kParquetConnectorName) + ->newConnector( + kParquetConnectorId, + std::make_shared( + std::unordered_map()), + ioExecutor_.get()); + facebook::velox::connector::registerConnector(parquetConnector); + dwio::common::registerFileSinks(); +} + +void ParquetConnectorTestBase::TearDown() { + // Make sure all pending loads are finished or cancelled before unregister + // connector. + ioExecutor_.reset(); + facebook::velox::connector::unregisterConnector(kParquetConnectorId); + facebook::velox::connector::unregisterConnectorFactory( + facebook::velox::cudf_velox::connector::parquet::ParquetConnectorFactory:: + kParquetConnectorName); + OperatorTestBase::TearDown(); +} + +void ParquetConnectorTestBase::resetParquetConnector( + const std::shared_ptr& config) { + facebook::velox::connector::unregisterConnector(kParquetConnectorId); + auto parquetConnector = + facebook::velox::connector::getConnectorFactory( + facebook::velox::cudf_velox::connector::parquet:: + ParquetConnectorFactory::kParquetConnectorName) + ->newConnector(kParquetConnectorId, config, ioExecutor_.get()); + facebook::velox::connector::registerConnector(parquetConnector); +} + +std::vector ParquetConnectorTestBase::makeVectors( + const RowTypePtr& rowType, + int32_t numVectors, + int32_t rowsPerVector) { + std::vector vectors; + for (int32_t i = 0; i < numVectors; ++i) { + auto vector = std::dynamic_pointer_cast( + velox::test::BatchMaker::createBatch(rowType, rowsPerVector, *pool_)); + vectors.push_back(vector); + } + return vectors; +} + +std::shared_ptr +ParquetConnectorTestBase::assertQuery( + const core::PlanNodePtr& plan, + const std::vector< + std::shared_ptr>& filePaths, + const std::string& duckDbSql) { + return OperatorTestBase::assertQuery( + plan, makeParquetConnectorSplits(filePaths), duckDbSql); +} + +std::shared_ptr +ParquetConnectorTestBase::assertQuery( + const facebook::velox::core::PlanNodePtr& plan, + const std::vector< + std::shared_ptr>& splits, + const std::string& duckDbSql, + const int32_t numPrefetchSplit) { + return facebook::velox::exec::test::AssertQueryBuilder( + plan, duckDbQueryRunner_) + .config( + facebook::velox::core::QueryConfig::kMaxSplitPreloadPerDriver, + std::to_string(numPrefetchSplit)) + .splits(splits) + .assertResults(duckDbSql); +} + +std::vector> +ParquetConnectorTestBase::makeFilePaths(int count) { + std::vector> + filePaths; + filePaths.reserve(count); + for (auto i = 0; i < count; ++i) { + filePaths.emplace_back(facebook::velox::exec::test::TempFilePath::create()); + } + return filePaths; +} + +void ParquetConnectorTestBase::writeToFile( + const std::string& filePath, + const std::vector& vectors, + std::string prefix) { + // Convert all RowVectorPtrs to cudf tables + std::vector> cudfTables; + cudfTables.reserve(vectors.size()); + for (const auto& vector : vectors) { + VELOX_CHECK_NOT_NULL(vector); + if (vector->size()) { + auto stream = cudf::get_default_stream(); + auto cudfTable = with_arrow::toCudfTable(vector, vector->pool(), stream); + stream.synchronize(); + cudfTables.emplace_back(std::move(cudfTable)); + } + } + // Make sure cudfTables has at least one table + if (cudfTables.empty()) { + VELOX_CHECK(not cudfTables.empty()); + return; + } + + // Create a sink and writer + auto const sinkInfo = cudf::io::sink_info(filePath); + auto tableInputMetadata = + cudf::io::table_input_metadata(cudfTables[0]->view()); + fillColumnNames(tableInputMetadata, prefix); + auto options = cudf::io::chunked_parquet_writer_options::builder(sinkInfo) + .metadata(tableInputMetadata) + .build(); + cudf::io::parquet_chunked_writer writer(options); + + // Write all table chunks + for (const auto& table : cudfTables) { + writer.write(table->view()); + } + + // Close the writer + writer.close(); +} + +void ParquetConnectorTestBase::writeToFile( + const std::string& filePath, + RowVectorPtr vector, + std::string prefix) { + auto const sinkInfo = cudf::io::sink_info(filePath); + VELOX_CHECK_NOT_NULL(vector); + auto stream = cudf::get_default_stream(); + auto cudfTable = with_arrow::toCudfTable(vector, vector->pool(), stream); + stream.synchronize(); + auto tableInputMetadata = cudf::io::table_input_metadata(cudfTable->view()); + fillColumnNames(tableInputMetadata, prefix); + auto options = + cudf::io::parquet_writer_options::builder(sinkInfo, cudfTable->view()) + .metadata(tableInputMetadata) + .build(); + cudf::io::write_parquet(options); +} + +std::unique_ptr +ParquetConnectorTestBase::makeColumnHandle( + const std::string& name, + const TypePtr& type, + const std::vector& children) { + return std::make_unique( + name, type, cudf::data_type(cudf::type_id::EMPTY), children); +} + +std::unique_ptr +ParquetConnectorTestBase::makeColumnHandle( + const std::string& name, + const TypePtr& type, + const cudf::data_type data_type, + const std::vector& children) { + return std::make_unique( + name, type, data_type, children); +} + +std::vector> +ParquetConnectorTestBase::makeParquetConnectorSplits( + const std::vector< + std::shared_ptr>& + filePaths) { + std::vector> + splits; + for (const auto& filePath : filePaths) { + splits.push_back(makeParquetConnectorSplit(filePath->getPath())); + } + return splits; +} + +std::vector> +ParquetConnectorTestBase::makeParquetConnectorSplits( + const std::string& filePath, + uint32_t splitCount) { + auto file = + filesystems::getFileSystem(filePath, nullptr)->openFileForRead(filePath); + const int64_t fileSize = file->size(); + // Take the upper bound. + const int64_t splitSize = std::ceil((fileSize) / splitCount); + std::vector> + splits; + // Add all the splits. + for (int i = 0; i < splitCount; i++) { + auto split = ParquetConnectorSplitBuilder(filePath).build(); + splits.push_back(std::move(split)); + } + return splits; +} + +std::shared_ptr +ParquetConnectorTestBase::makeParquetConnectorSplit( + const std::string& filePath, + int64_t splitWeight) { + return ParquetConnectorSplitBuilder(filePath) + .splitWeight(splitWeight) + .build(); +} + +// static +std::shared_ptr +ParquetConnectorTestBase::makeParquetInsertTableHandle( + const std::vector& tableColumnNames, + const std::vector& tableColumnTypes, + std::shared_ptr locationHandle, + const std::optional compressionKind, + const std::unordered_map& serdeParameters, + const std::shared_ptr& writerOptions) { + std::vector> + columnHandles; + + for (int i = 0; i < tableColumnNames.size(); ++i) { + columnHandles.push_back( + std::make_shared( + tableColumnNames.at(i), + tableColumnTypes.at(i), + cudf::data_type{veloxToCudfTypeId(tableColumnTypes.at(i))})); + } + + return std::make_shared( + columnHandles, + locationHandle, + compressionKind, + serdeParameters, + writerOptions); +} + +} // namespace facebook::velox::cudf_velox::exec::test diff --git a/velox/experimental/cudf/tests/utils/ParquetConnectorTestBase.h b/velox/experimental/cudf/tests/utils/ParquetConnectorTestBase.h new file mode 100644 index 00000000000..092af2a19b7 --- /dev/null +++ b/velox/experimental/cudf/tests/utils/ParquetConnectorTestBase.h @@ -0,0 +1,189 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#pragma once + +#include "velox/experimental/cudf/connectors/parquet/ParquetConfig.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetConnector.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetDataSink.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetDataSource.h" +#include "velox/experimental/cudf/connectors/parquet/ParquetTableHandle.h" + +#include "velox/exec/Operator.h" +#include "velox/exec/tests/utils/OperatorTestBase.h" +#include "velox/exec/tests/utils/TempFilePath.h" +#include "velox/type/tests/SubfieldFiltersBuilder.h" + +namespace facebook::velox::cudf_velox::exec::test { + +static const std::string kParquetConnectorId = "test-parquet"; + +using ColumnHandleMap = std::unordered_map< + std::string, + std::shared_ptr>; + +class ParquetConnectorTestBase + : public facebook::velox::exec::test::OperatorTestBase { + public: + ParquetConnectorTestBase(); + + void SetUp() override; + void TearDown() override; + + void resetParquetConnector( + const std::shared_ptr& config); + + void writeToFile( + const std::string& filePath, + RowVectorPtr vector, + std::string prefix = "c"); + + void writeToFile( + const std::string& filePath, + const std::vector& vectors, + std::string prefix = "c"); + + std::vector makeVectors( + const RowTypePtr& rowType, + int32_t numVectors, + int32_t rowsPerVector); + + using facebook::velox::exec::test::OperatorTestBase::assertQuery; + + /// Assumes plan has a single TableScan node. + std::shared_ptr assertQuery( + const facebook::velox::core::PlanNodePtr& plan, + const std::vector< + std::shared_ptr>& + filePaths, + const std::string& duckDbSql); + + std::shared_ptr assertQuery( + const facebook::velox::core::PlanNodePtr& plan, + const std::vector< + std::shared_ptr>& splits, + const std::string& duckDbSql, + const int32_t numPrefetchSplit); + + static std::vector> + makeFilePaths(int count); + + static std::shared_ptr< + facebook::velox::cudf_velox::connector::parquet::ParquetConnectorSplit> + makeParquetConnectorSplit( + const std::string& filePath, + int64_t splitWeight = 0); + + static std::vector< + std::shared_ptr> + makeParquetConnectorSplits( + const std::vector< + std::shared_ptr>& + filePaths); + + static std::vector> + makeParquetConnectorSplits(const std::string& filePath, uint32_t splitCount); + + static std::shared_ptr + makeTableHandle( + const std::string& tableName = "parquet_table", + const RowTypePtr& dataColumns = nullptr, + bool filterPushdownEnabled = false, + const core::TypedExprPtr& subfieldFilterExpr = nullptr, + const core::TypedExprPtr& remainingFilterExpr = nullptr) { + return std::make_shared( + kParquetConnectorId, + tableName, + filterPushdownEnabled, + subfieldFilterExpr, + remainingFilterExpr, + dataColumns); + } + + /// @param name Column name. + /// @param type Column type. + /// @param Required subfields of this column. + static std::unique_ptr + makeColumnHandle( + const std::string& name, + const TypePtr& type, + const std::vector& children); + + /// @param name Column name. + /// @param type Column type. + /// @param type cudf column type. + /// @param Required subfields of this column. + static std::unique_ptr + makeColumnHandle( + const std::string& name, + const TypePtr& type, + const cudf::data_type data_type, + const std::vector& children); + + /// @param targetDirectory Final directory of the target table. + /// @param tableType Whether to create a new table. + static std::shared_ptr makeLocationHandle( + std::string targetDirectory) { + return std::make_shared( + targetDirectory, + connector::parquet::LocationHandle::TableType::kNew, + ""); + } + + /// @param targetDirectory Final directory of the target table. + /// @param tableType Whether to create a new table, insert into an existing + /// table, or write a temporary table. + /// @param targetDirectory Final file name of the target table . + static std::shared_ptr makeLocationHandle( + std::string targetDirectory, + connector::parquet::LocationHandle::TableType tableType = + connector::parquet::LocationHandle::TableType::kNew, + std::string targetFileName = "") { + return std::make_shared( + targetDirectory, tableType, targetFileName); + } + + /// Build a ParquetInsertTableHandle. + /// @param tableColumnNames Column names of the target table. Corresponding + /// type of tableColumnNames[i] is tableColumnTypes[i]. + /// @param tableColumnTypes Column types of the target table. Corresponding + /// name of tableColumnTypes[i] is tableColumnNames[i]. + /// @param locationHandle Location handle for the table write. + /// @param compressionKind compression algorithm to use for table write. + /// @param serdeParameters Table writer configuration parameters. + static std::shared_ptr + makeParquetInsertTableHandle( + const std::vector& tableColumnNames, + const std::vector& tableColumnTypes, + std::shared_ptr locationHandle, + const std::optional compressionKind = {}, + const std::unordered_map& serdeParameters = {}, + const std::shared_ptr& writerOptions = + nullptr); +}; + +/// Same as connector::parquet::ParquetConnectorBuilder, except that this +/// defaults connectorId to kParquetConnectorId. +class ParquetConnectorSplitBuilder + : public connector::parquet::ParquetConnectorSplitBuilder { + public: + explicit ParquetConnectorSplitBuilder(std::string filePath) + : connector::parquet::ParquetConnectorSplitBuilder(filePath) { + connectorId(kParquetConnectorId); + } +}; + +} // namespace facebook::velox::cudf_velox::exec::test diff --git a/velox/experimental/cudf/vector/CMakeLists.txt b/velox/experimental/cudf/vector/CMakeLists.txt new file mode 100644 index 00000000000..d26f0b4c7dc --- /dev/null +++ b/velox/experimental/cudf/vector/CMakeLists.txt @@ -0,0 +1,26 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# Licensed 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. + +add_library(velox_cudf_vector CudfVector.cpp) + +set_target_properties( + velox_cudf_vector + PROPERTIES CUDA_ARCHITECTURES native) + +target_link_libraries( + velox_cudf_vector + cudf::cudf + velox_exception + velox_common_base + velox_vector) diff --git a/velox/experimental/cudf/vector/CudfVector.cpp b/velox/experimental/cudf/vector/CudfVector.cpp new file mode 100644 index 00000000000..93e1c22bd67 --- /dev/null +++ b/velox/experimental/cudf/vector/CudfVector.cpp @@ -0,0 +1,123 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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. + */ + +#include "velox/experimental/cudf/vector/CudfVector.h" + +#include "velox/buffer/Buffer.h" +#include "velox/common/memory/MemoryPool.h" +#include "velox/vector/TypeAliases.h" + +#include +#include + +namespace facebook::velox::cudf_velox { +namespace { + +/// Calculates the total memory size in bytes of a cudf column and reconstructs +/// it. +/// +/// This function disassembles a cudf column to access its underlying memory +/// buffers, calculates the total size including children columns (for nested +/// types), and then reassembles the column. +/// +/// @return A pair containing the total size in bytes and the reconstructed +/// column +std::pair> getColumnSize( + std::unique_ptr column) { + // Store column metadata (type, null count, and size) before releasing it, + // as the release() operation transfers ownership of the underlying buffers + // and invalidates access to these properties. + auto type = column->type(); + auto nullCount = column->null_count(); + auto size = column->size(); + + auto contents = column->release(); + auto bytes = contents.data->size() + contents.null_mask->size(); + + // Recursively get the size of the children columns. + std::vector> children; + for (auto& child : contents.children) { + auto [childBytes, childColumn] = getColumnSize(std::move(child)); + bytes += childBytes; + children.push_back(std::move(childColumn)); + } + + // Reassemble the column with the original metadata. + auto reconstitutedColumn = std::make_unique( + type, + size, + std::move(*contents.data.release()), + std::move(*contents.null_mask.release()), + nullCount, + std::move(children)); + + return std::make_pair(bytes, std::move(reconstitutedColumn)); +} + +/// Calculates the total memory size in bytes of a cudf table and reconstructs +/// it. +/// +/// This function disassembles a cudf table to access its underlying columns, +/// calculates the total size, and then reassembles the table. +/// +/// @note This is a workaround because cudf::table doesn't have an API to get +/// this information without involving estimation and d->h copies. +/// @see https://github.com/rapidsai/cudf/issues/18462 +/// +/// @return A pair containing the total size in bytes and the reconstructed +/// table +std::pair> getTableSize( + std::unique_ptr&& table) { + auto columns = table->release(); + std::vector> columnsOut; + uint64_t totalBytes = 0; + + for (auto& column : columns) { + auto [bytes, columnOut] = getColumnSize(std::move(column)); + totalBytes += bytes; + columnsOut.push_back(std::move(columnOut)); + } + return std::make_pair( + totalBytes, std::make_unique(std::move(columnsOut))); +} + +} // namespace + +CudfVector::CudfVector( + velox::memory::MemoryPool* pool, + TypePtr type, + vector_size_t size, + std::unique_ptr&& table, + rmm::cuda_stream_view stream) + : RowVector( + pool, + std::move(type), + BufferPtr(nullptr), + size, + std::vector(), + std::nullopt), + table_{std::move(table)}, + stream_{stream} { + auto [bytes, tableOut] = getTableSize(std::move(table_)); + flatSize_ = bytes; + table_ = std::move(tableOut); +} + +uint64_t CudfVector::estimateFlatSize() const { + return flatSize_; +} + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/vector/CudfVector.h b/velox/experimental/cudf/vector/CudfVector.h index db1590b3c08..bb27bd4a889 100644 --- a/velox/experimental/cudf/vector/CudfVector.h +++ b/velox/experimental/cudf/vector/CudfVector.h @@ -15,13 +15,13 @@ */ #pragma once -#include "velox/buffer/Buffer.h" #include "velox/common/memory/MemoryPool.h" #include "velox/vector/ComplexVector.h" #include "velox/vector/TypeAliases.h" #include -#include + +#include #include #include @@ -36,16 +36,7 @@ class CudfVector : public RowVector { TypePtr type, vector_size_t size, std::unique_ptr&& table, - rmm::cuda_stream_view stream) - : RowVector( - pool, - std::move(type), - BufferPtr(nullptr), - size, - std::vector(), - std::nullopt), - table_{std::move(table)}, - stream_{stream} {} + rmm::cuda_stream_view stream); rmm::cuda_stream_view stream() const { return stream_; @@ -56,12 +47,16 @@ class CudfVector : public RowVector { } std::unique_ptr&& release() { + flatSize_ = 0; return std::move(table_); } + uint64_t estimateFlatSize() const override; + private: std::unique_ptr table_; rmm::cuda_stream_view stream_; + uint64_t flatSize_; }; using CudfVectorPtr = std::shared_ptr; diff --git a/velox/experimental/gpu/tests/CMakeLists.txt b/velox/experimental/gpu/tests/CMakeLists.txt index 6202a78d0bd..27fcff16c4d 100644 --- a/velox/experimental/gpu/tests/CMakeLists.txt +++ b/velox/experimental/gpu/tests/CMakeLists.txt @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. list(APPEND CMAKE_PREFIX_PATH "${CUDAToolkit_LIBRARY_DIR}/cmake") +find_package(Thrust REQUIRED) find_package(CUB REQUIRED) add_executable(velox_gpu_hash_table_test HashTableTest.cu) @@ -21,4 +22,5 @@ target_link_libraries( gflags::gflags glog::glog CUB::CUB + Thrust::Thrust CUDA::cudart)